diff --git a/packages/leann-backend-diskann/third_party/DiskANN b/packages/leann-backend-diskann/third_party/DiskANN new file mode 160000 index 0000000..015c201 --- /dev/null +++ b/packages/leann-backend-diskann/third_party/DiskANN @@ -0,0 +1 @@ +Subproject commit 015c201141cfd35e2054772358ac5ae7d3dd25a6 diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.clang-format b/packages/leann-backend-diskann/third_party/DiskANN/.clang-format deleted file mode 100644 index ad3192f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.clang-format +++ /dev/null @@ -1,6 +0,0 @@ ---- -BasedOnStyle: Microsoft ---- -Language: Cpp -SortIncludes: false -... diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.gitattributes b/packages/leann-backend-diskann/third_party/DiskANN/.gitattributes deleted file mode 100644 index fbf9358..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.gitattributes +++ /dev/null @@ -1,14 +0,0 @@ -# Set the default behavior, in case people don't have core.autocrlf set. -* text=auto - -# Explicitly declare text files you want to always be normalized and converted -# to native line endings on checkout. -*.c text -*.h text - -# Declare files that will always have CRLF line endings on checkout. -*.sln text eol=crlf - -# Denote all files that are truly binary and should not be modified. -*.png binary -*.jpg binary diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/bug_report.md b/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/bug_report.md deleted file mode 100644 index 829d38d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/bug_report.md +++ /dev/null @@ -1,40 +0,0 @@ ---- -name: Bug report -about: Bug reports help us improve! Thanks for submitting yours! -title: "[BUG] " -labels: bug -assignees: '' - ---- - -## Expected Behavior -Tell us what should happen - -## Actual Behavior -Tell us what happens instead - -## Example Code -Please see [How to create a Minimal, Reproducible example](https://stackoverflow.com/help/minimal-reproducible-example) for some guidance on creating the best possible example of the problem -```bash - -``` - -## Dataset Description -Please tell us about the shape and datatype of your data, (e.g. 128 dimensions, 12.3 billion points, floats) -- Dimensions: -- Number of Points: -- Data type: - -## Error -``` -Paste the full error, with any sensitive information minimally redacted and marked $$REDACTED$$ - -``` - -## Your Environment -* Operating system (e.g. Windows 11 Pro, Ubuntu 22.04.1 LTS) -* DiskANN version (or commit built from) - -## Additional Details -Any other contextual information you might feel is important. - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/config.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/config.yml deleted file mode 100644 index 99d680b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/config.yml +++ /dev/null @@ -1,2 +0,0 @@ -blank_issues_enabled: false - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/feature_request.md b/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/feature_request.md deleted file mode 100644 index 9c3c58c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/feature_request.md +++ /dev/null @@ -1,25 +0,0 @@ ---- -name: Feature request -about: Suggest an idea for this project -title: '' -labels: enhancement -assignees: '' - ---- - -## Is your feature request related to a problem? Please describe. -A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] - -## Describe the solution you'd like -A clear and concise description of what you want to happen. - -## Describe alternatives you've considered -A clear and concise description of any alternative solutions or features you've considered. - -## Provide references (if applicable) -If your feature request is related to a published algorithm/idea, please provide links to -any relevant articles or webpages. - -## Additional context -Add any other context or screenshots about the feature request here. - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/usage-question.md b/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/usage-question.md deleted file mode 100644 index 7532f76..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/ISSUE_TEMPLATE/usage-question.md +++ /dev/null @@ -1,11 +0,0 @@ ---- -name: Usage Question -about: Ask us a question about DiskANN! -title: "[Question]" -labels: question -assignees: '' - ---- - -This is our forum for asking whatever DiskANN question you'd like! No need to feel shy - we're happy to talk about use cases and optimal tuning strategies! - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/PULL_REQUEST_TEMPLATE.md b/packages/leann-backend-diskann/third_party/DiskANN/.github/PULL_REQUEST_TEMPLATE.md deleted file mode 100644 index 0b97019..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/PULL_REQUEST_TEMPLATE.md +++ /dev/null @@ -1,22 +0,0 @@ - -- [ ] Does this PR have a descriptive title that could go in our release notes? -- [ ] Does this PR add any new dependencies? -- [ ] Does this PR modify any existing APIs? - - [ ] Is the change to the API backwards compatible? -- [ ] Should this result in any changes to our documentation, either updating existing docs or adding new ones? - -#### Reference Issues/PRs - - -#### What does this implement/fix? Briefly explain your changes. - -#### Any other comments? - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/build/action.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/build/action.yml deleted file mode 100644 index 219d9d6..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/build/action.yml +++ /dev/null @@ -1,39 +0,0 @@ -name: 'DiskANN Build Bootstrap' -description: 'Prepares DiskANN build environment and executes build' -runs: - using: "composite" - steps: - # ------------ Linux Build --------------- - - name: Prepare and Execute Build - if: ${{ runner.os == 'Linux' }} - run: | - sudo scripts/dev/install-dev-deps-ubuntu.bash - cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DUNIT_TEST=True - cmake --build build -- -j - cmake --install build --prefix="dist" - shell: bash - # ------------ End Linux Build --------------- - # ------------ Windows Build --------------- - - name: Add VisualStudio command line tools into path - if: runner.os == 'Windows' - uses: ilammy/msvc-dev-cmd@v1 - - name: Run configure and build for Windows - if: runner.os == 'Windows' - run: | - mkdir build && cd build && cmake .. -DUNIT_TEST=True && msbuild diskann.sln /m /nologo /t:Build /p:Configuration="Release" /property:Platform="x64" -consoleloggerparameters:"ErrorsOnly;Summary" - cd .. - mkdir dist - mklink /j .\dist\bin .\x64\Release\ - shell: cmd - # ------------ End Windows Build --------------- - # ------------ Windows Build With EXEC_ENV_OLS and USE_BING_INFRA --------------- - - name: Add VisualStudio command line tools into path - if: runner.os == 'Windows' - uses: ilammy/msvc-dev-cmd@v1 - - name: Run configure and build for Windows with Bing feature flags - if: runner.os == 'Windows' - run: | - mkdir build_bing && cd build_bing && cmake .. -DEXEC_ENV_OLS=1 -DUSE_BING_INFRA=1 -DUNIT_TEST=True && msbuild diskann.sln /m /nologo /t:Build /p:Configuration="Release" /property:Platform="x64" -consoleloggerparameters:"ErrorsOnly;Summary" - cd .. - shell: cmd - # ------------ End Windows Build --------------- diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/format-check/action.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/format-check/action.yml deleted file mode 100644 index 6ed08c0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/format-check/action.yml +++ /dev/null @@ -1,13 +0,0 @@ -name: 'Checking code formatting...' -description: 'Ensures code complies with code formatting rules' -runs: - using: "composite" - steps: - - name: Checking code formatting... - run: | - sudo apt install clang-format - find include -name '*.h' -type f -print0 | xargs -0 -P 16 /usr/bin/clang-format --Werror --dry-run - find src -name '*.cpp' -type f -print0 | xargs -0 -P 16 /usr/bin/clang-format --Werror --dry-run - find apps -name '*.cpp' -type f -print0 | xargs -0 -P 16 /usr/bin/clang-format --Werror --dry-run - find python -name '*.cpp' -type f -print0 | xargs -0 -P 16 /usr/bin/clang-format --Werror --dry-run - shell: bash diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-high-dim-random/action.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-high-dim-random/action.yml deleted file mode 100644 index 65e9b7e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-high-dim-random/action.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: 'Generating Random Data (Basic)' -description: 'Generates the random data files used in acceptance tests' -runs: - using: "composite" - steps: - - name: Generate Random Data (Basic) - run: | - mkdir data - - echo "Generating random 1020,1024,1536D float and 4096 int8 vectors for index" - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1020D_5K_norm1.0.bin -D 1020 -N 5000 --norm 1.0 - #dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1024D_5K_norm1.0.bin -D 1024 -N 5000 --norm 1.0 - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1536D_5K_norm1.0.bin -D 1536 -N 5000 --norm 1.0 - dist/bin/rand_data_gen --data_type int8 --output_file data/rand_int8_4096D_5K_norm1.0.bin -D 4096 -N 5000 --norm 1.0 - - echo "Generating random 1020,1024,1536D float and 4096D int8 avectors for query" - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1020D_1K_norm1.0.bin -D 1020 -N 1000 --norm 1.0 - #dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1024D_1K_norm1.0.bin -D 1024 -N 1000 --norm 1.0 - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_1536D_1K_norm1.0.bin -D 1536 -N 1000 --norm 1.0 - dist/bin/rand_data_gen --data_type int8 --output_file data/rand_int8_4096D_1K_norm1.0.bin -D 4096 -N 1000 --norm 1.0 - - echo "Computing ground truth for 1020,1024,1536D float and 4096D int8 avectors for query" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/rand_float_1020D_5K_norm1.0.bin --query_file data/rand_float_1020D_1K_norm1.0.bin --gt_file data/l2_rand_float_1020D_5K_norm1.0_1020D_1K_norm1.0_gt100 --K 100 - #dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/rand_float_1024D_5K_norm1.0.bin --query_file data/rand_float_1024D_1K_norm1.0.bin --gt_file data/l2_rand_float_1024D_5K_norm1.0_1024D_1K_norm1.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/rand_float_1536D_5K_norm1.0.bin --query_file data/rand_float_1536D_1K_norm1.0.bin --gt_file data/l2_rand_float_1536D_5K_norm1.0_1536D_1K_norm1.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type int8 --dist_fn l2 --base_file data/rand_int8_4096D_5K_norm1.0.bin --query_file data/rand_int8_4096D_1K_norm1.0.bin --gt_file data/l2_rand_int8_4096D_5K_norm1.0_4096D_1K_norm1.0_gt100 --K 100 - - shell: bash diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-random/action.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-random/action.yml deleted file mode 100644 index 2755067..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/generate-random/action.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: 'Generating Random Data (Basic)' -description: 'Generates the random data files used in acceptance tests' -runs: - using: "composite" - steps: - - name: Generate Random Data (Basic) - run: | - mkdir data - - echo "Generating random vectors for index" - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_10D_10K_norm1.0.bin -D 10 -N 10000 --norm 1.0 - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_10D_10K_unnorm.bin -D 10 -N 10000 --rand_scaling 2.0 - dist/bin/rand_data_gen --data_type int8 --output_file data/rand_int8_10D_10K_norm50.0.bin -D 10 -N 10000 --norm 50.0 - dist/bin/rand_data_gen --data_type uint8 --output_file data/rand_uint8_10D_10K_norm50.0.bin -D 10 -N 10000 --norm 50.0 - - echo "Generating random vectors for query" - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_10D_1K_norm1.0.bin -D 10 -N 1000 --norm 1.0 - dist/bin/rand_data_gen --data_type float --output_file data/rand_float_10D_1K_unnorm.bin -D 10 -N 1000 --rand_scaling 2.0 - dist/bin/rand_data_gen --data_type int8 --output_file data/rand_int8_10D_1K_norm50.0.bin -D 10 -N 1000 --norm 50.0 - dist/bin/rand_data_gen --data_type uint8 --output_file data/rand_uint8_10D_1K_norm50.0.bin -D 10 -N 1000 --norm 50.0 - - echo "Computing ground truth for floats across l2, mips, and cosine distance functions" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/rand_float_10D_10K_norm1.0.bin --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type float --dist_fn mips --base_file data/rand_float_10D_10K_norm1.0.bin --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/mips_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type float --dist_fn cosine --base_file data/rand_float_10D_10K_norm1.0.bin --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/cosine_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type float --dist_fn cosine --base_file data/rand_float_10D_10K_unnorm.bin --query_file data/rand_float_10D_1K_unnorm.bin --gt_file data/cosine_rand_float_10D_10K_unnorm_10D_1K_unnorm_gt100 --K 100 - - echo "Computing ground truth for int8s across l2, mips, and cosine distance functions" - dist/bin/compute_groundtruth --data_type int8 --dist_fn l2 --base_file data/rand_int8_10D_10K_norm50.0.bin --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type int8 --dist_fn mips --base_file data/rand_int8_10D_10K_norm50.0.bin --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/mips_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type int8 --dist_fn cosine --base_file data/rand_int8_10D_10K_norm50.0.bin --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/cosine_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - - echo "Computing ground truth for uint8s across l2, mips, and cosine distance functions" - dist/bin/compute_groundtruth --data_type uint8 --dist_fn l2 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type uint8 --dist_fn mips --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/mips_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type uint8 --dist_fn cosine --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - - shell: bash diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/python-wheel/action.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/python-wheel/action.yml deleted file mode 100644 index 6a2880c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/actions/python-wheel/action.yml +++ /dev/null @@ -1,22 +0,0 @@ -name: Build Python Wheel -description: Builds a python wheel with cibuildwheel -inputs: - cibw-identifier: - description: "CI build wheel identifier to build" - required: true -runs: - using: "composite" - steps: - - uses: actions/setup-python@v3 - - name: Install cibuildwheel - run: python -m pip install cibuildwheel==2.11.3 - shell: bash - - name: Building Python ${{inputs.cibw-identifier}} Wheel - run: python -m cibuildwheel --output-dir dist - env: - CIBW_BUILD: ${{inputs.cibw-identifier}} - shell: bash - - uses: actions/upload-artifact@v3 - with: - name: wheels - path: ./dist/*.whl diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python-pdoc.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python-pdoc.yml deleted file mode 100644 index 444a7ee..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python-pdoc.yml +++ /dev/null @@ -1,81 +0,0 @@ -name: DiskANN Build PDoc Documentation -on: [workflow_call] -jobs: - build-reference-documentation: - permissions: - contents: write - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Set up Python 3.9 - uses: actions/setup-python@v2 - with: - python-version: 3.9 - - name: Install python build - run: python -m pip install build - shell: bash - # Install required dependencies - - name: Prepare Linux environment - run: | - sudo scripts/dev/install-dev-deps-ubuntu.bash - shell: bash - # We need to build the wheel in order to run pdoc. pdoc does not seem to work if you just point it at - # our source directory. - - name: Building Python Wheel for documentation generation - run: python -m build --wheel --outdir documentation_dist - shell: bash - - name: "Run Reference Documentation Generation" - run: | - pip install pdoc pipdeptree - pip install documentation_dist/*.whl - echo "documentation" > dependencies_documentation.txt - pipdeptree >> dependencies_documentation.txt - pdoc -o docs/python/html diskannpy - - name: Create version environment variable - run: | - echo "DISKANN_VERSION=$(python <> $GITHUB_ENV - - name: Archive documentation version artifact - uses: actions/upload-artifact@v4 - with: - name: dependencies - path: | - ${{ github.run_id }}-dependencies_documentation.txt - overwrite: true - - name: Archive documentation artifacts - uses: actions/upload-artifact@v4 - with: - name: documentation-site - path: | - docs/python/html - # Publish to /dev if we are on the "main" branch - - name: Publish reference docs for latest development version (main branch) - uses: peaceiris/actions-gh-pages@v3 - if: github.ref == 'refs/heads/main' - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/python/html - destination_dir: docs/python/dev - # Publish to / if we are releasing - - name: Publish reference docs by version (main branch) - uses: peaceiris/actions-gh-pages@v3 - if: github.event_name == 'release' - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/python/html - destination_dir: docs/python/${{ env.DISKANN_VERSION }} - # Publish to /latest if we are releasing - - name: Publish latest reference docs (main branch) - uses: peaceiris/actions-gh-pages@v3 - if: github.event_name == 'release' - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/python/html - destination_dir: docs/python/latest diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python.yml deleted file mode 100644 index b825398..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/build-python.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: DiskANN Build Python Wheel -on: [workflow_call] -jobs: - linux-build: - name: Python - Ubuntu - ${{matrix.cibw-identifier}} - strategy: - fail-fast: false - matrix: - cibw-identifier: ["cp39-manylinux_x86_64", "cp310-manylinux_x86_64", "cp311-manylinux_x86_64"] - runs-on: ubuntu-latest - defaults: - run: - shell: bash - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Building python wheel ${{matrix.cibw-identifier}} - uses: ./.github/actions/python-wheel - with: - cibw-identifier: ${{matrix.cibw-identifier}} - windows-build: - name: Python - Windows - ${{matrix.cibw-identifier}} - strategy: - fail-fast: false - matrix: - cibw-identifier: ["cp39-win_amd64", "cp310-win_amd64", "cp311-win_amd64"] - runs-on: windows-latest - defaults: - run: - shell: bash - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - submodules: true - fetch-depth: 1 - - name: Building python wheel ${{matrix.cibw-identifier}} - uses: ./.github/actions/python-wheel - with: - cibw-identifier: ${{matrix.cibw-identifier}} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/common.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/common.yml deleted file mode 100644 index 09c020a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/common.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: DiskANN Common Checks -# common means common to both pr-test and push-test -on: [workflow_call] -jobs: - formatting-check: - strategy: - fail-fast: true - name: Code Formatting Test - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checking code formatting... - uses: ./.github/actions/format-check - docker-container-build: - name: Docker Container Build - needs: [formatting-check] - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Docker build - run: | - docker build . \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/disk-pq.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/disk-pq.yml deleted file mode 100644 index 930d213..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/disk-pq.yml +++ /dev/null @@ -1,117 +0,0 @@ -name: Disk With PQ -on: [workflow_call] -jobs: - acceptance-tests-disk-pq: - name: Disk, PQ - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: build and search disk index (one shot graph build, L2, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_oneshot -R 16 -L 32 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, cosine, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn cosine --data_path data/rand_float_10D_10K_unnorm.bin --index_path_prefix data/disk_index_cosine_rand_float_10D_10K_unnorm_diskfull_oneshot -R 16 -L 32 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type float --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/disk_index_cosine_rand_float_10D_10K_unnorm_diskfull_oneshot --result_path /tmp/res --query_file data/rand_float_10D_1K_unnorm.bin --gt_file data/cosine_rand_float_10D_10K_unnorm_10D_1K_unnorm_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, L2, no diskPQ) (int8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_oneshot -R 16 -L 32 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, L2, no diskPQ) (uint8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_oneshot -R 16 -L 32 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name: build and search disk index (one shot graph build, L2, no diskPQ, build with PQ distance comparisons) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_oneshot_buildpq5 -R 16 -L 32 -B 0.00003 -M 1 --build_PQ_bytes 5 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_oneshot_buildpq5 --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, L2, no diskPQ, build with PQ distance comparisons) (int8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_oneshot_buildpq5 -R 16 -L 32 -B 0.00003 -M 1 --build_PQ_bytes 5 - dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_oneshot_buildpq5 --result_path /tmp/res --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16\ - - name: build and search disk index (one shot graph build, L2, no diskPQ, build with PQ distance comparisons) (uint8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_oneshot_buildpq5 -R 16 -L 32 -B 0.00003 -M 1 --build_PQ_bytes 5 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_oneshot_buildpq5 --result_path /tmp/res --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name: build and search disk index (sharded graph build, L2, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_sharded -R 16 -L 32 -B 0.00003 -M 0.00006 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskfull_sharded --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (sharded graph build, cosine, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn cosine --data_path data/rand_float_10D_10K_unnorm.bin --index_path_prefix data/disk_index_cosine_rand_float_10D_10K_unnorm_diskfull_sharded -R 16 -L 32 -B 0.00003 -M 0.00006 - dist/bin/search_disk_index --data_type float --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/disk_index_cosine_rand_float_10D_10K_unnorm_diskfull_sharded --result_path /tmp/res --query_file data/rand_float_10D_1K_unnorm.bin --gt_file data/cosine_rand_float_10D_10K_unnorm_10D_1K_unnorm_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (sharded graph build, L2, no diskPQ) (int8) - run: | - dist/bin/build_disk_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_sharded -R 16 -L 32 -B 0.00003 -M 0.00006 - dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskfull_sharded --result_path /tmp/res --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (sharded graph build, L2, no diskPQ) (uint8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_sharded -R 16 -L 32 -B 0.00003 -M 0.00006 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskfull_sharded --result_path /tmp/res --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name: build and search disk index (one shot graph build, L2, diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskpq_oneshot -R 16 -L 32 -B 0.00003 -M 1 --PQ_disk_bytes 5 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_10D_10K_norm1.0_diskpq_oneshot --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, L2, diskPQ) (int8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskpq_oneshot -R 16 -L 32 -B 0.00003 -M 1 --PQ_disk_bytes 5 - dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_10D_10K_norm50.0_diskpq_oneshot --result_path /tmp/res --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search disk index (one shot graph build, L2, diskPQ) (uint8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskpq_oneshot -R 16 -L 32 -B 0.00003 -M 1 --PQ_disk_bytes 5 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50.0_diskpq_oneshot --result_path /tmp/res --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name: build and search disk index (sharded graph build, MIPS, diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn mips --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/disk_index_mips_rand_float_10D_10K_norm1.0_diskpq_sharded -R 16 -L 32 -B 0.00003 -M 0.00006 --PQ_disk_bytes 5 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_mips_rand_float_10D_10K_norm1.0_diskpq_sharded --result_path /tmp/res --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/mips_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: disk-pq-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic-labels.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic-labels.yml deleted file mode 100644 index d5dc712..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic-labels.yml +++ /dev/null @@ -1,102 +0,0 @@ -name: Dynamic-Labels -on: [workflow_call] -jobs: - acceptance-tests-dynamic: - name: Dynamic-Labels - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: Generate Labels - run: | - echo "Generating synthetic labels and computing ground truth for filtered search with universal label" - dist/bin/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file data/rand_labels_50_10K.txt --distribution_type random - - echo "Generating synthetic labels with a zipf distribution and computing ground truth for filtered search with universal label" - dist/bin/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file data/zipf_labels_50_10K.txt --distribution_type zipf - - - name: Test a streaming index (float) with labels (Zipf distributed) - run: | - dist/bin/test_streaming_scenario --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --universal_label 0 --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/index_zipf_stream -R 64 --FilteredLbuild 200 -L 50 --alpha 1.2 --insert_threads 8 --consolidate_threads 8 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 3.2 --unique_labels_supported 51 - - echo "Computing groundtruth with filter" - dist/bin/compute_groundtruth_for_filters --data_type float --universal_label 0 --filter_label 1 --dist_fn l2 --base_file data/index_zipf_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_zipf_base-act4000-cons2000-max10000_1 --label_file data/index_zipf_stream.after-streaming-act4000-cons2000-max10000_raw_labels.txt --tags_file data/index_zipf_stream.after-streaming-act4000-cons2000-max10000.tags - echo "Searching with filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --filter_label 1 --fail_if_recall_below 40 --index_path_prefix data/index_zipf_stream.after-streaming-act4000-cons2000-max10000 --result_path data/res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_zipf_base-act4000-cons2000-max10000_1 -K 10 -L 20 40 60 80 100 150 -T 64 --dynamic true --tags 1 - - echo "Computing groundtruth w/o filter" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_zipf_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_zipf_base-act4000-cons2000-max10000 - echo "Searching without filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_zipf_stream.after-streaming-act4000-cons2000-max10000 --result_path res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_zipf_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 - - - name: Test a streaming index (float) with labels (random distributed) - run: | - dist/bin/test_streaming_scenario --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --universal_label 0 --label_file data/rand_labels_50_10K.txt --index_path_prefix data/index_rand_stream -R 64 --FilteredLbuild 200 -L 50 --alpha 1.2 --insert_threads 8 --consolidate_threads 8 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 3.2 --unique_labels_supported 51 - - echo "Computing groundtruth with filter" - dist/bin/compute_groundtruth_for_filters --data_type float --universal_label 0 --filter_label 1 --dist_fn l2 --base_file data/index_rand_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_rand_base-act4000-cons2000-max10000_1 --label_file data/index_rand_stream.after-streaming-act4000-cons2000-max10000_raw_labels.txt --tags_file data/index_rand_stream.after-streaming-act4000-cons2000-max10000.tags - echo "Searching with filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --filter_label 1 --fail_if_recall_below 40 --index_path_prefix data/index_rand_stream.after-streaming-act4000-cons2000-max10000 --result_path data/res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_rand_base-act4000-cons2000-max10000_1 -K 10 -L 20 40 60 80 100 150 -T 64 --dynamic true --tags 1 - - echo "Computing groundtruth w/o filter" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_rand_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_rand_base-act4000-cons2000-max10000 - echo "Searching without filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_rand_stream.after-streaming-act4000-cons2000-max10000 --result_path res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_rand_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 - - - name: Test Insert Delete Consolidate (float) with labels (zipf distributed) - run: | - dist/bin/test_insert_deletes_consolidate --data_type float --dist_fn l2 --universal_label 0 --label_file data/zipf_labels_50_10K.txt --FilteredLbuild 70 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_zipf_ins_del -R 64 -L 10 --alpha 1.2 --points_to_skip 0 --max_points_to_insert 7500 --beginning_index_size 0 --points_per_checkpoint 1000 --checkpoints_per_snapshot 0 --points_to_delete_from_beginning 2500 --start_deletes_after 5000 --do_concurrent true --start_point_norm 3.2 --unique_labels_supported 51 - - echo "Computing groundtruth with filter" - dist/bin/compute_groundtruth_for_filters --data_type float --filter_label 5 --universal_label 0 --dist_fn l2 --base_file data/index_zipf_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_zipf_random10D_1K_wlabel_5 --label_file data/index_zipf_ins_del.after-concurrent-delete-del2500-7500_raw_labels.txt --tags_file data/index_zipf_ins_del.after-concurrent-delete-del2500-7500.tags - echo "Searching with filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --filter_label 5 --fail_if_recall_below 10 --index_path_prefix data/index_zipf_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_zipf_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_zipf_random10D_1K_wlabel_5 -K 10 -L 20 40 60 80 100 150 -T 64 --dynamic true --tags 1 - - echo "Computing groundtruth w/o filter" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_zipf_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_zipf_random10D_1K - echo "Searching without filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_zipf_ins_del.after-concurrent-delete-del2500-7500 --result_path res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_zipf_random10D_1K -K 10 -L 20 40 60 80 100 -T 64 - - - name: Test Insert Delete Consolidate (float) with labels (random distributed) - run: | - dist/bin/test_insert_deletes_consolidate --data_type float --dist_fn l2 --universal_label 0 --label_file data/rand_labels_50_10K.txt --FilteredLbuild 70 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_rand_ins_del -R 64 -L 10 --alpha 1.2 --points_to_skip 0 --max_points_to_insert 7500 --beginning_index_size 0 --points_per_checkpoint 1000 --checkpoints_per_snapshot 0 --points_to_delete_from_beginning 2500 --start_deletes_after 5000 --do_concurrent true --start_point_norm 3.2 --unique_labels_supported 51 - - echo "Computing groundtruth with filter" - dist/bin/compute_groundtruth_for_filters --data_type float --filter_label 5 --universal_label 0 --dist_fn l2 --base_file data/index_rand_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_rand_random10D_1K_wlabel_5 --label_file data/index_rand_ins_del.after-concurrent-delete-del2500-7500_raw_labels.txt --tags_file data/index_rand_ins_del.after-concurrent-delete-del2500-7500.tags - echo "Searching with filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --filter_label 5 --fail_if_recall_below 40 --index_path_prefix data/index_rand_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_rand_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_rand_random10D_1K_wlabel_5 -K 10 -L 20 40 60 80 100 150 -T 64 --dynamic true --tags 1 - - echo "Computing groundtruth w/o filter" - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_rand_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_rand_random10D_1K - echo "Searching without filter" - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_rand_ins_del.after-concurrent-delete-del2500-7500 --result_path res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_rand_random10D_1K -K 10 -L 20 40 60 80 100 -T 64 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: dynamic-labels-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic.yml deleted file mode 100644 index edd691e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/dynamic.yml +++ /dev/null @@ -1,75 +0,0 @@ -name: Dynamic -on: [workflow_call] -jobs: - acceptance-tests-dynamic: - name: Dynamic - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: test a streaming index (float) - run: | - dist/bin/test_streaming_scenario --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 3.2 - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_base-act4000-cons2000-max10000 --tags_file data/index_stream.after-streaming-act4000-cons2000-max10000.tags - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_stream.after-streaming-act4000-cons2000-max10000 --result_path data/res_stream --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 - - name: test a streaming index (int8) - if: success() || failure() - run: | - dist/bin/test_streaming_scenario --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200 - dist/bin/compute_groundtruth --data_type int8 --dist_fn l2 --base_file data/index_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_int8_10D_1K_norm50.0.bin --K 100 --gt_file data/gt100_base-act4000-cons2000-max10000 --tags_file data/index_stream.after-streaming-act4000-cons2000-max10000.tags - dist/bin/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_stream.after-streaming-act4000-cons2000-max10000 --result_path res_stream --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/gt100_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 - - name: test a streaming index - if: success() || failure() - run: | - dist/bin/test_streaming_scenario --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200 - dist/bin/compute_groundtruth --data_type uint8 --dist_fn l2 --base_file data/index_stream.after-streaming-act4000-cons2000-max10000.data --query_file data/rand_uint8_10D_1K_norm50.0.bin --K 100 --gt_file data/gt100_base-act4000-cons2000-max10000 --tags_file data/index_stream.after-streaming-act4000-cons2000-max10000.tags - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_stream.after-streaming-act4000-cons2000-max10000 --result_path data/res_stream --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/gt100_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 - - - name: build and search an incremental index (float) - if: success() || failure() - run: | - dist/bin/test_insert_deletes_consolidate --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_ins_del -R 64 -L 300 --alpha 1.2 -T 8 --points_to_skip 0 --max_points_to_insert 7500 --beginning_index_size 0 --points_per_checkpoint 1000 --checkpoints_per_snapshot 0 --points_to_delete_from_beginning 2500 --start_deletes_after 5000 --do_concurrent true --start_point_norm 3.2; - dist/bin/compute_groundtruth --data_type float --dist_fn l2 --base_file data/index_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_float_10D_1K_norm1.0.bin --K 100 --gt_file data/gt100_random10D_1K-conc-2500-7500 --tags_file data/index_ins_del.after-concurrent-delete-del2500-7500.tags - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_ins_del --query_file data/rand_float_10D_1K_norm1.0.bin --gt_file data/gt100_random10D_1K-conc-2500-7500 -K 10 -L 20 40 60 80 100 -T 8 --dynamic true --tags 1 - - name: build and search an incremental index (int8) - if: success() || failure() - run: | - dist/bin/test_insert_deletes_consolidate --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/index_ins_del -R 64 -L 300 --alpha 1.2 -T 8 --points_to_skip 0 --max_points_to_insert 7500 --beginning_index_size 0 --points_per_checkpoint 1000 --checkpoints_per_snapshot 0 --points_to_delete_from_beginning 2500 --start_deletes_after 5000 --do_concurrent true --start_point_norm 200 - dist/bin/compute_groundtruth --data_type int8 --dist_fn l2 --base_file data/index_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_int8_10D_1K_norm50.0.bin --K 100 --gt_file data/gt100_random10D_1K-conc-2500-7500 --tags_file data/index_ins_del.after-concurrent-delete-del2500-7500.tags - dist/bin/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_ins_del --query_file data/rand_int8_10D_1K_norm50.0.bin --gt_file data/gt100_random10D_1K-conc-2500-7500 -K 10 -L 20 40 60 80 100 -T 8 --dynamic true --tags 1 - - name: build and search an incremental index (uint8) - if: success() || failure() - run: | - dist/bin/test_insert_deletes_consolidate --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/index_ins_del -R 64 -L 300 --alpha 1.2 -T 8 --points_to_skip 0 --max_points_to_insert 7500 --beginning_index_size 0 --points_per_checkpoint 1000 --checkpoints_per_snapshot 0 --points_to_delete_from_beginning 2500 --start_deletes_after 5000 --do_concurrent true --start_point_norm 200 - dist/bin/compute_groundtruth --data_type uint8 --dist_fn l2 --base_file data/index_ins_del.after-concurrent-delete-del2500-7500.data --query_file data/rand_uint8_10D_1K_norm50.0.bin --K 100 --gt_file data/gt100_random10D_10K-conc-2500-7500 --tags_file data/index_ins_del.after-concurrent-delete-del2500-7500.tags - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_ins_del.after-concurrent-delete-del2500-7500 --result_path data/res_ins_del --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/gt100_random10D_10K-conc-2500-7500 -K 10 -L 20 40 60 80 100 -T 8 --dynamic true --tags 1 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: dynamic-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-no-pq.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-no-pq.yml deleted file mode 100644 index 07fc4a2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-no-pq.yml +++ /dev/null @@ -1,81 +0,0 @@ -name: In-Memory Without PQ -on: [workflow_call] -jobs: - acceptance-tests-mem-no-pq: - name: In-Mem, Without PQ - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: build and search in-memory index with L2 metrics (float) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0 - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0 --query_file data/rand_float_10D_1K_norm1.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 -L 16 32 - - name: build and search in-memory index with L2 metrics (int8) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/index_l2_rand_int8_10D_10K_norm50.0 - dist/bin/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_int8_10D_10K_norm50.0 --query_file data/rand_int8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - name: build and search in-memory index with L2 metrics (uint8) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0 - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0 --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - - name: Searching with fast_l2 distance function (float) - if: runner.os != 'Windows' && (success() || failure()) - run: | - dist/bin/search_memory_index --data_type float --dist_fn fast_l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0 --query_file data/rand_float_10D_1K_norm1.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 -L 16 32 - - - name: build and search in-memory index with MIPS metric (float) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type float --dist_fn mips --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_mips_rand_float_10D_10K_norm1.0 - dist/bin/search_memory_index --data_type float --dist_fn mips --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0 --query_file data/rand_float_10D_1K_norm1.0.bin --recall_at 10 --result_path temp --gt_file data/mips_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 -L 16 32 - - - name: build and search in-memory index with cosine metric (float) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type float --dist_fn cosine --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_cosine_rand_float_10D_10K_norm1.0 - dist/bin/search_memory_index --data_type float --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0 --query_file data/rand_float_10D_1K_norm1.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 -L 16 32 - - name: build and search in-memory index with cosine metric (int8) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type int8 --dist_fn cosine --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/index_cosine_rand_int8_10D_10K_norm50.0 - dist/bin/search_memory_index --data_type int8 --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_int8_10D_10K_norm50.0 --query_file data/rand_int8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - name: build and search in-memory index with cosine metric - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn cosine --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/index_cosine_rand_uint8_10D_10K_norm50.0 - dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0 --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: in-memory-no-pq-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-pq.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-pq.yml deleted file mode 100644 index be20f10..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/in-mem-pq.yml +++ /dev/null @@ -1,56 +0,0 @@ -name: In-Memory With PQ -on: [workflow_call] -jobs: - acceptance-tests-mem-pq: - name: In-Mem, PQ - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: build and search in-memory index with L2 metric with PQ based distance comparisons (float) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type float --dist_fn l2 --data_path data/rand_float_10D_10K_norm1.0.bin --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0_buildpq5 --build_PQ_bytes 5 - dist/bin/search_memory_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_float_10D_10K_norm1.0_buildpq5 --query_file data/rand_float_10D_1K_norm1.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_float_10D_10K_norm1.0_10D_1K_norm1.0_gt100 -L 16 32 - - - name: build and search in-memory index with L2 metrics with PQ base distance comparisons (int8) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_10D_10K_norm50.0.bin --index_path_prefix data/index_l2_rand_int8_10D_10K_norm50.0_buildpq5 --build_PQ_bytes 5 - dist/bin/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_int8_10D_10K_norm50.0_buildpq5 --query_file data/rand_int8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_int8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - - name: build and search in-memory index with L2 metrics with PQ base distance comparisons (uint8) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --data_path data/rand_uint8_10D_10K_norm50.0.bin --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0_buildpq5 --build_PQ_bytes 5 - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50.0_buildpq5 --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 16 32 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: in-memory-pq-${{matrix.os}} - path: | - ./dist/** - ./data/** \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/labels.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/labels.yml deleted file mode 100644 index 93995f7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/labels.yml +++ /dev/null @@ -1,120 +0,0 @@ -name: Labels -on: [workflow_call] -jobs: - acceptance-tests-labels: - name: Labels - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-random - - - name: Generate Labels - run: | - echo "Generating synthetic labels and computing ground truth for filtered search with universal label" - dist/bin/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file data/rand_labels_50_10K.txt --distribution_type random - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn l2 --universal_label 0 --filter_label 10 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn mips --universal_label 0 --filter_label 10 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --gt_file data/mips_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn cosine --universal_label 0 --filter_label 10 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - - echo "Generating synthetic labels with a zipf distribution and computing ground truth for filtered search with universal label" - dist/bin/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file data/zipf_labels_50_10K.txt --distribution_type zipf - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn l2 --universal_label 0 --filter_label 5 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn mips --universal_label 0 --filter_label 5 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --gt_file data/mips_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn cosine --universal_label 0 --filter_label 5 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --gt_file data/cosine_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - - echo "Generating synthetic labels and computing ground truth for filtered search without a universal label" - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn l2 --filter_label 5 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel_nouniversal --K 100 - dist/bin/generate_synthetic_labels --num_labels 10 --num_points 1000 --output_file data/query_labels_1K.txt --distribution_type one_per_point - dist/bin/compute_groundtruth_for_filters --data_type uint8 --dist_fn l2 --universal_label 0 --filter_label_file data/query_labels_1K.txt --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --gt_file data/combined_l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --K 100 - - - name: build and search in-memory index with labels using L2 and Cosine metrics (random distributed labels) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50_wlabel - dist/bin/build_memory_index --data_type uint8 --dist_fn cosine --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --index_path_prefix data/index_cosine_rand_uint8_10D_10K_norm50_wlabel - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --filter_label 10 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --filter_label 10 --fail_if_recall_below 70 --index_path_prefix data/index_cosine_rand_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - - echo "Searching without filters" - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 32 64 - dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_cosine_rand_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 32 64 - - - name: build and search disk index with labels using L2 and Cosine metrics (random distributed labels) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --universal_label 0 --FilteredLbuild 90 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50_wlabel -R 32 -L 5 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --filter_label 10 --fail_if_recall_below 50 --index_path_prefix data/disk_index_l2_rand_uint8_10D_10K_norm50_wlabel --result_path temp --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: build and search in-memory index with labels using L2 and Cosine metrics (zipf distributed labels) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel - dist/bin/build_memory_index --data_type uint8 --dist_fn cosine --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/index_cosine_zipf_uint8_10D_10K_norm50_wlabel - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --filter_label 5 --fail_if_recall_below 70 --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --filter_label 5 --fail_if_recall_below 70 --index_path_prefix data/index_cosine_zipf_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - - echo "Searching without filters" - dist/bin/compute_groundtruth --data_type uint8 --dist_fn l2 --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/compute_groundtruth --data_type uint8 --dist_fn cosine --base_file data/rand_uint8_10D_10K_norm50.0.bin --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/cosine_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 --K 100 - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 32 64 - dist/bin/search_memory_index --data_type uint8 --dist_fn cosine --fail_if_recall_below 70 --index_path_prefix data/index_cosine_zipf_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/cosine_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100 -L 32 64 - - - name: build and search disk index with labels using L2 and Cosine metrics (zipf distributed labels) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --universal_label 0 --FilteredLbuild 90 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/disk_index_l2_zipf_uint8_10D_10K_norm50_wlabel -R 32 -L 5 -B 0.00003 -M 1 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --filter_label 5 --fail_if_recall_below 50 --index_path_prefix data/disk_index_l2_zipf_uint8_10D_10K_norm50_wlabel --result_path temp --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - - name : build and search in-memory and disk index (without universal label, zipf distributed) - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel_nouniversal - dist/bin/build_disk_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/disk_index_l2_zipf_uint8_10D_10K_norm50_wlabel_nouniversal -R 32 -L 5 -B 0.00003 -M 1 - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --filter_label 5 --fail_if_recall_below 70 --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel_nouniversal --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel_nouniversal -L 16 32 - dist/bin/search_disk_index --data_type uint8 --dist_fn l2 --filter_label 5 --index_path_prefix data/disk_index_l2_zipf_uint8_10D_10K_norm50_wlabel_nouniversal --result_path temp --query_file data/rand_uint8_10D_1K_norm50.0.bin --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel_nouniversal --recall_at 5 -L 5 12 -W 2 --num_nodes_to_cache 10 -T 16 - - name: Generate combined GT for each query with a separate label and search - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --query_filters_file data/query_labels_1K.txt --fail_if_recall_below 70 --index_path_prefix data/index_l2_zipf_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/combined_l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - - name: build and search in-memory index with pq_dist of 5 with 10 dimensions - if: success() || failure() - run: | - dist/bin/build_memory_index --data_type uint8 --dist_fn l2 --FilteredLbuild 90 --universal_label 0 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/rand_labels_50_10K.txt --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50_wlabel --build_PQ_bytes 5 - dist/bin/search_memory_index --data_type uint8 --dist_fn l2 --filter_label 10 --fail_if_recall_below 70 --index_path_prefix data/index_l2_rand_uint8_10D_10K_norm50_wlabel --query_file data/rand_uint8_10D_1K_norm50.0.bin --recall_at 10 --result_path temp --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -L 16 32 - - name: Build and search stitched vamana with random and zipf distributed labels - if: success() || failure() - run: | - dist/bin/build_stitched_index --num_threads 48 --data_type uint8 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/rand_labels_50_10K.txt -R 32 -L 100 --alpha 1.2 --stitched_R 64 --index_path_prefix data/stit_rand_32_100_64_new --universal_label 0 - dist/bin/build_stitched_index --num_threads 48 --data_type uint8 --data_path data/rand_uint8_10D_10K_norm50.0.bin --label_file data/zipf_labels_50_10K.txt -R 32 -L 100 --alpha 1.2 --stitched_R 64 --index_path_prefix data/stit_zipf_32_100_64_new --universal_label 0 - dist/bin/search_memory_index --num_threads 48 --data_type uint8 --dist_fn l2 --filter_label 10 --index_path_prefix data/stit_rand_32_100_64_new --query_file data/rand_uint8_10D_1K_norm50.0.bin --result_path data/rand_stit_96_10_90_new --gt_file data/l2_rand_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -K 10 -L 16 32 150 - dist/bin/search_memory_index --num_threads 48 --data_type uint8 --dist_fn l2 --filter_label 5 --index_path_prefix data/stit_zipf_32_100_64_new --query_file data/rand_uint8_10D_1K_norm50.0.bin --result_path data/zipf_stit_96_10_90_new --gt_file data/l2_zipf_uint8_10D_10K_norm50.0_10D_1K_norm50.0_gt100_wlabel -K 10 -L 16 32 150 - - - name: upload data and bin - if: success() || failure() - uses: actions/upload-artifact@v4 - with: - name: labels-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/multi-sector-disk-pq.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/multi-sector-disk-pq.yml deleted file mode 100644 index 969467a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/multi-sector-disk-pq.yml +++ /dev/null @@ -1,60 +0,0 @@ -name: Disk With PQ -on: [workflow_call] -jobs: - acceptance-tests-disk-pq: - name: Disk, PQ - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Generate Data - uses: ./.github/actions/generate-high-dim-random - - - name: build and search disk index (1020D, one shot graph build, L2, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_1020D_5K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_1020D_5K_norm1.0_diskfull_oneshot -R 32 -L 500 -B 0.003 -M 1 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_1020D_5K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_float_1020D_1K_norm1.0.bin --gt_file data/l2_rand_float_1020D_5K_norm1.0_1020D_1K_norm1.0_gt100 --recall_at 5 -L 250 -W 2 --num_nodes_to_cache 100 -T 16 - #- name: build and search disk index (1024D, one shot graph build, L2, no diskPQ) (float) - # if: success() || failure() - # run: | - # dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_1024D_5K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_1024D_5K_norm1.0_diskfull_oneshot -R 32 -L 500 -B 0.003 -M 1 - # dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_1024D_5K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_float_1024D_1K_norm1.0.bin --gt_file data/l2_rand_float_1024D_5K_norm1.0_1024D_1K_norm1.0_gt100 --recall_at 5 -L 250 -W 2 --num_nodes_to_cache 100 -T 16 - - name: build and search disk index (1536D, one shot graph build, L2, no diskPQ) (float) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type float --dist_fn l2 --data_path data/rand_float_1536D_5K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_float_1536D_5K_norm1.0_diskfull_oneshot -R 32 -L 500 -B 0.003 -M 1 - dist/bin/search_disk_index --data_type float --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_float_1536D_5K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_float_1536D_1K_norm1.0.bin --gt_file data/l2_rand_float_1536D_5K_norm1.0_1536D_1K_norm1.0_gt100 --recall_at 5 -L 250 -W 2 --num_nodes_to_cache 100 -T 16 - - - name: build and search disk index (4096D, one shot graph build, L2, no diskPQ) (int8) - if: success() || failure() - run: | - dist/bin/build_disk_index --data_type int8 --dist_fn l2 --data_path data/rand_int8_4096D_5K_norm1.0.bin --index_path_prefix data/disk_index_l2_rand_int8_4096D_5K_norm1.0_diskfull_oneshot -R 32 -L 500 -B 0.003 -M 1 - dist/bin/search_disk_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix data/disk_index_l2_rand_int8_4096D_5K_norm1.0_diskfull_oneshot --result_path /tmp/res --query_file data/rand_int8_4096D_1K_norm1.0.bin --gt_file data/l2_rand_int8_4096D_5K_norm1.0_4096D_1K_norm1.0_gt100 --recall_at 5 -L 250 -W 2 --num_nodes_to_cache 100 -T 16 - - - name: upload data and bin - uses: actions/upload-artifact@v4 - with: - name: multi-sector-disk-pq-${{matrix.os}} - path: | - ./dist/** - ./data/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/perf.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/perf.yml deleted file mode 100644 index d4eb9e2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/perf.yml +++ /dev/null @@ -1,26 +0,0 @@ -name: DiskANN Nightly Performance Metrics -on: - schedule: - - cron: "41 14 * * *" # 14:41 UTC, 7:41 PDT, 8:41 PST, 08:11 IST -jobs: - perf-test: - name: Run Perf Test from main - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Build Perf Container - run: | - docker build --build-arg GIT_COMMIT_ISH="$GITHUB_SHA" -t perf -f scripts/perf/Dockerfile scripts - - name: Performance Tests - run: | - mkdir metrics - docker run -v ./metrics:/app/logs perf &> ./metrics/combined_stdouterr.log - - name: Upload Metrics Logs - uses: actions/upload-artifact@v4 - with: - name: metrics-${{matrix.os}} - path: | - ./metrics/** diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/pr-test.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/pr-test.yml deleted file mode 100644 index f84953b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/pr-test.yml +++ /dev/null @@ -1,35 +0,0 @@ -name: DiskANN Pull Request Build and Test -on: [pull_request] -jobs: - common: - strategy: - fail-fast: true - name: DiskANN Common Build Checks - uses: ./.github/workflows/common.yml - unit-tests: - name: Unit tests - uses: ./.github/workflows/unit-tests.yml - in-mem-pq: - name: In-Memory with PQ - uses: ./.github/workflows/in-mem-pq.yml - in-mem-no-pq: - name: In-Memory without PQ - uses: ./.github/workflows/in-mem-no-pq.yml - disk-pq: - name: Disk with PQ - uses: ./.github/workflows/disk-pq.yml - multi-sector-disk-pq: - name: Multi-sector Disk with PQ - uses: ./.github/workflows/multi-sector-disk-pq.yml - labels: - name: Labels - uses: ./.github/workflows/labels.yml - dynamic: - name: Dynamic - uses: ./.github/workflows/dynamic.yml - dynamic-labels: - name: Dynamic Labels - uses: ./.github/workflows/dynamic-labels.yml - python: - name: Python - uses: ./.github/workflows/build-python.yml diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/push-test.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/push-test.yml deleted file mode 100644 index d1261d5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/push-test.yml +++ /dev/null @@ -1,50 +0,0 @@ -name: DiskANN Push Build -on: [push] -jobs: - common: - strategy: - fail-fast: true - name: DiskANN Common Build Checks - uses: ./.github/workflows/common.yml - build-documentation: - permissions: - contents: write - strategy: - fail-fast: true - name: DiskANN Build Documentation - uses: ./.github/workflows/build-python-pdoc.yml - build: - strategy: - fail-fast: false - matrix: - os: [ ubuntu-latest, windows-2019, windows-latest ] - name: Build for ${{matrix.os}} - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: Build diskannpy dependency tree - run: | - pip install diskannpy pipdeptree - echo "dependencies" > dependencies_${{ matrix.os }}.txt - pipdeptree >> dependencies_${{ matrix.os }}.txt - - name: Archive diskannpy dependencies artifact - uses: actions/upload-artifact@v4 - with: - name: dependencies_${{ matrix.os }} - path: | - dependencies_${{ matrix.os }}.txt - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/python-release.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/python-release.yml deleted file mode 100644 index a15d4d1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/python-release.yml +++ /dev/null @@ -1,43 +0,0 @@ -name: Build and Release Python Wheels -on: - release: - types: [published] -jobs: - python-release-wheels: - name: Python - uses: ./.github/workflows/build-python.yml - build-documentation: - strategy: - fail-fast: true - name: DiskANN Build Documentation - uses: ./.github/workflows/build-python-pdoc.yml - release: - permissions: - contents: write - runs-on: ubuntu-latest - needs: python-release-wheels - steps: - - uses: actions/download-artifact@v3 - with: - name: wheels - path: dist/ - - name: Generate SHA256 files for each wheel - run: | - sha256sum dist/*.whl > checksums.txt - cat checksums.txt - - uses: actions/setup-python@v3 - - name: Install twine - run: python -m pip install twine - - name: Publish with twine - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} - run: | - twine upload dist/*.whl - - name: Update release with SHA256 and Artifacts - uses: softprops/action-gh-release@v1 - with: - token: ${{ secrets.GITHUB_TOKEN }} - files: | - dist/*.whl - checksums.txt diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/unit-tests.yml b/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/unit-tests.yml deleted file mode 100644 index 6ae6877..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.github/workflows/unit-tests.yml +++ /dev/null @@ -1,32 +0,0 @@ -name: Unit Tests -on: [workflow_call] -jobs: - acceptance-tests-labels: - name: Unit Tests - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, windows-2019, windows-latest] - runs-on: ${{matrix.os}} - defaults: - run: - shell: bash - steps: - - name: Checkout repository - if: ${{ runner.os == 'Linux' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - name: Checkout repository - if: ${{ runner.os == 'Windows' }} - uses: actions/checkout@v3 - with: - fetch-depth: 1 - submodules: true - - name: DiskANN Build CLI Applications - uses: ./.github/actions/build - - - name: Run Unit Tests - run: | - cd build - ctest -C Release \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.gitignore b/packages/leann-backend-diskann/third_party/DiskANN/.gitignore deleted file mode 100644 index c6a88e7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.gitignore +++ /dev/null @@ -1,384 +0,0 @@ -## Ignore Visual Studio temporary files, build results, and -## files generated by popular Visual Studio add-ons. -## -## Get latest from https://github.com/github/gitignore/blob/master/VisualStudio.gitignore - -# User-specific files -*.rsuser -*.suo -*.user -*.userosscache -*.sln.docstates - -# User-specific files (MonoDevelop/Xamarin Studio) -*.userprefs - -# Mono auto generated files -mono_crash.* - -# Build results -[Dd]ebug/ -[Dd]ebugPublic/ -[Rr]elease/ -[Rr]eleases/ -x64/ -x86/ -[Aa][Rr][Mm]/ -[Aa][Rr][Mm]64/ -bld/ -[Bb]in/ -[Oo]bj/ -[Ll]og/ -[Ll]ogs/ - -# Visual Studio 2015/2017 cache/options directory -.vs/ -# Uncomment if you have tasks that create the project's static files in wwwroot -#wwwroot/ - -# Visual Studio 2017 auto generated files -Generated\ Files/ - -# MSTest test Results -[Tt]est[Rr]esult*/ -[Bb]uild[Ll]og.* - -# NUnit -*.VisualState.xml -TestResult.xml -nunit-*.xml - -# Build Results of an ATL Project -[Dd]ebugPS/ -[Rr]eleasePS/ -dlldata.c - -# Benchmark Results -BenchmarkDotNet.Artifacts/ - -# .NET Core -project.lock.json -project.fragment.lock.json -artifacts/ - -# StyleCop -StyleCopReport.xml - -# Files built by Visual Studio -*_i.c -*_p.c -*_h.h -*.ilk -*.meta -*.obj -*.iobj -*.pch -*.pdb -*.ipdb -*.pgc -*.pgd -*.rsp -*.sbr -*.tlb -*.tli -*.tlh -*.tmp -*.tmp_proj -*_wpftmp.csproj -*.log -*.vspscc -*.vssscc -.builds -*.pidb -*.svclog -*.scc - -# Chutzpah Test files -_Chutzpah* - -# Visual C++ cache files -ipch/ -*.aps -*.ncb -*.opendb -*.opensdf -*.sdf -*.cachefile -*.VC.db -*.VC.VC.opendb - -# Visual Studio profiler -*.psess -*.vsp -*.vspx -*.sap - -# Visual Studio Trace Files -*.e2e - -# TFS 2012 Local Workspace -$tf/ - -# Guidance Automation Toolkit -*.gpState - -# ReSharper is a .NET coding add-in -_ReSharper*/ -*.[Rr]e[Ss]harper -*.DotSettings.user - -# TeamCity is a build add-in -_TeamCity* - -# DotCover is a Code Coverage Tool -*.dotCover - -# AxoCover is a Code Coverage Tool -.axoCover/* -!.axoCover/settings.json - -# Visual Studio code coverage results -*.coverage -*.coveragexml - -# NCrunch -_NCrunch_* -.*crunch*.local.xml -nCrunchTemp_* - -# MightyMoose -*.mm.* -AutoTest.Net/ - -# Web workbench (sass) -.sass-cache/ - -# Installshield output folder -[Ee]xpress/ - -# DocProject is a documentation generator add-in -DocProject/buildhelp/ -DocProject/Help/*.HxT -DocProject/Help/*.HxC -DocProject/Help/*.hhc -DocProject/Help/*.hhk -DocProject/Help/*.hhp -DocProject/Help/Html2 -DocProject/Help/html - -# Click-Once directory -publish/ - -# Publish Web Output -*.[Pp]ublish.xml -*.azurePubxml -# Note: Comment the next line if you want to checkin your web deploy settings, -# but database connection strings (with potential passwords) will be unencrypted -*.pubxml -*.publishproj - -# Microsoft Azure Web App publish settings. Comment the next line if you want to -# checkin your Azure Web App publish settings, but sensitive information contained -# in these scripts will be unencrypted -PublishScripts/ - -# NuGet Packages -*.nupkg -# NuGet Symbol Packages -*.snupkg -# The packages folder can be ignored because of Package Restore -**/[Pp]ackages/* -# except build/, which is used as an MSBuild target. -!**/[Pp]ackages/build/ -# Uncomment if necessary however generally it will be regenerated when needed -#!**/[Pp]ackages/repositories.config -# NuGet v3's project.json files produces more ignorable files -*.nuget.props -*.nuget.targets - -# Microsoft Azure Build Output -csx/ -*.build.csdef - -# Microsoft Azure Emulator -ecf/ -rcf/ - -# Windows Store app package directories and files -AppPackages/ -BundleArtifacts/ -Package.StoreAssociation.xml -_pkginfo.txt -*.appx -*.appxbundle -*.appxupload - -# Visual Studio cache files -# files ending in .cache can be ignored -*.[Cc]ache -# but keep track of directories ending in .cache -!?*.[Cc]ache/ - -# Others -ClientBin/ -~$* -*~ -*.dbmdl -*.dbproj.schemaview -*.jfm -*.pfx -*.publishsettings -orleans.codegen.cs - -# Including strong name files can present a security risk -# (https://github.com/github/gitignore/pull/2483#issue-259490424) -#*.snk - -# Since there are multiple workflows, uncomment next line to ignore bower_components -# (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) -#bower_components/ - -# RIA/Silverlight projects -Generated_Code/ - -# Backup & report files from converting an old project file -# to a newer Visual Studio version. Backup files are not needed, -# because we have git ;-) -_UpgradeReport_Files/ -Backup*/ -UpgradeLog*.XML -UpgradeLog*.htm -ServiceFabricBackup/ -*.rptproj.bak - -# SQL Server files -*.mdf -*.ldf -*.ndf - -# Business Intelligence projects -*.rdl.data -*.bim.layout -*.bim_*.settings -*.rptproj.rsuser -*- [Bb]ackup.rdl -*- [Bb]ackup ([0-9]).rdl -*- [Bb]ackup ([0-9][0-9]).rdl - -# Microsoft Fakes -FakesAssemblies/ - -# GhostDoc plugin setting file -*.GhostDoc.xml - -# Node.js Tools for Visual Studio -.ntvs_analysis.dat -node_modules/ - -# Visual Studio 6 build log -*.plg - -# Visual Studio 6 workspace options file -*.opt - -# Visual Studio 6 auto-generated workspace file (contains which files were open etc.) -*.vbw - -# Visual Studio LightSwitch build output -**/*.HTMLClient/GeneratedArtifacts -**/*.DesktopClient/GeneratedArtifacts -**/*.DesktopClient/ModelManifest.xml -**/*.Server/GeneratedArtifacts -**/*.Server/ModelManifest.xml -_Pvt_Extensions - -# Paket dependency manager -.paket/paket.exe -paket-files/ - -# FAKE - F# Make -.fake/ - -# CodeRush personal settings -.cr/personal - -# Python Tools for Visual Studio (PTVS) -__pycache__/ -*.pyc - -# Cake - Uncomment if you are using it -# tools/** -# !tools/packages.config - -# Tabs Studio -*.tss - -# Telerik's JustMock configuration file -*.jmconfig - -# BizTalk build output -*.btp.cs -*.btm.cs -*.odx.cs -*.xsd.cs - -# OpenCover UI analysis results -OpenCover/ - -# Azure Stream Analytics local run output -ASALocalRun/ - -# MSBuild Binary and Structured Log -*.binlog - -# NVidia Nsight GPU debugger configuration file -*.nvuser - -# MFractors (Xamarin productivity tool) working folder -.mfractor/ - -# Local History for Visual Studio -.localhistory/ - -# BeatPulse healthcheck temp database -healthchecksdb - -# Backup folder for Package Reference Convert tool in Visual Studio 2017 -MigrationBackup/ - -# Ionide (cross platform F# VS Code tools) working folder -.ionide/ - -/vcproj/nsg/x64/Debug/nsg.Build.CppClean.log -/vcproj/test_recall/x64/Debug/test_recall.Build.CppClean.log -/vcproj/test_recall/test_recall.vcxproj.user -/.vs -/out/build/x64-Debug -cscope* - -build/ -build_linux/ -!.github/actions/build - -# jetbrains specific stuff -.idea/ -cmake-build-debug/ - -#python extension module ignores -python/diskannpy.egg-info/ -python/dist/ - -**/*.egg-info -wheelhouse/* -dist/* -venv*/** -*.swp - -gperftools - -# Rust -rust/target - -python/src/*.so - -compile_commands.json \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/.gitmodules b/packages/leann-backend-diskann/third_party/DiskANN/.gitmodules deleted file mode 100644 index 125572b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/.gitmodules +++ /dev/null @@ -1,3 +0,0 @@ -[submodule "gperftools"] - path = gperftools - url = https://github.com/gperftools/gperftools.git diff --git a/packages/leann-backend-diskann/third_party/DiskANN/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/CMakeLists.txt deleted file mode 100644 index 4025861..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/CMakeLists.txt +++ /dev/null @@ -1,563 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -# Parameters: -# -# BOOST_ROOT: -# Specify root of the Boost library if Boost cannot be auto-detected. On Windows, a fallback to a -# downloaded nuget version will be used if Boost cannot be found. -# -# DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS: -# This is a work-in-progress feature, not completed yet. The core DiskANN library will be split into -# build-related and search-related functionality. In build-related functionality, when using tcmalloc, -# it's possible to release memory that's free but reserved by tcmalloc. Setting this to true enables -# such behavior. -# Contact for this feature: gopalrs. - - -# Some variables like MSVC are defined only after project(), so put that first. -cmake_minimum_required(VERSION 3.20) -project(diskann) - -#Set option to use tcmalloc -option(USE_TCMALLOC "Use tcmalloc from gperftools" ON) - -# set tcmalloc to false when on macos -if(APPLE) - set(USE_TCMALLOC OFF) -endif() - -option(PYBIND "Build with Python bindings" ON) - -if(PYBIND) - # Find Python - find_package(Python 3.6 COMPONENTS Interpreter Development REQUIRED) - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import pybind11; print(pybind11.get_cmake_dir())" - OUTPUT_VARIABLE pybind11_DIR - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - find_package(pybind11 CONFIG REQUIRED) - - message(STATUS "Python include dirs: ${Python_INCLUDE_DIRS}") - message(STATUS "Pybind11 include dirs: ${pybind11_INCLUDE_DIRS}") - - # Add pybind11 include directories - include_directories(SYSTEM ${pybind11_INCLUDE_DIRS} ${Python_INCLUDE_DIRS}) - - # Add compilation definitions - add_definitions(-DPYBIND11_EMBEDDED) - - # Set visibility flags - if(NOT MSVC) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fvisibility=hidden") - endif() -endif() - -set(CMAKE_STANDARD 17) -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED ON) - -# if(NOT MSVC) -# set(CMAKE_CXX_COMPILER g++) -# endif() - -set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake;${CMAKE_MODULE_PATH}") - -# Install nuget packages for dependencies. -if (MSVC) - find_program(NUGET_EXE NAMES nuget) - - if (NOT NUGET_EXE) - message(FATAL_ERROR "Cannot find nuget command line tool.\nPlease install it from e.g. https://www.nuget.org/downloads") - endif() - - set(DISKANN_MSVC_PACKAGES_CONFIG ${CMAKE_BINARY_DIR}/packages.config) - set(DISKANN_MSVC_PACKAGES ${CMAKE_BINARY_DIR}/packages) - - message(STATUS "Invoking nuget to download Boost, OpenMP and MKL dependencies...") - configure_file(${PROJECT_SOURCE_DIR}/windows/packages.config.in ${DISKANN_MSVC_PACKAGES_CONFIG}) - exec_program(${NUGET_EXE} ARGS install \"${DISKANN_MSVC_PACKAGES_CONFIG}\" -ExcludeVersion -OutputDirectory \"${DISKANN_MSVC_PACKAGES}\") - if (RESTAPI) - set(DISKANN_MSVC_RESTAPI_PACKAGES_CONFIG ${CMAKE_BINARY_DIR}/restapi/packages.config) - configure_file(${PROJECT_SOURCE_DIR}/windows/packages_restapi.config.in ${DISKANN_MSVC_RESTAPI_PACKAGES_CONFIG}) - exec_program(${NUGET_EXE} ARGS install \"${DISKANN_MSVC_RESTAPI_PACKAGES_CONFIG}\" -ExcludeVersion -OutputDirectory \"${DISKANN_MSVC_PACKAGES}\") - endif() - message(STATUS "Finished setting up nuget dependencies") -endif() - -include_directories(${PROJECT_SOURCE_DIR}/include) - -include(FetchContent) - -if(USE_TCMALLOC) - FetchContent_Declare( - tcmalloc - GIT_REPOSITORY https://github.com/google/tcmalloc.git - GIT_TAG origin/master # or specify a particular version or commit - ) - - FetchContent_MakeAvailable(tcmalloc) -endif() - -if(NOT PYBIND) - set(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS ON) -endif() -# It's necessary to include tcmalloc headers only if calling into MallocExtension interface. -# For using tcmalloc in DiskANN tools, it's enough to just link with tcmalloc. -if (DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) - include_directories(${tcmalloc_SOURCE_DIR}/src) - if (MSVC) - include_directories(${tcmalloc_SOURCE_DIR}/src/windows) - endif() -endif() - -#OpenMP -if (MSVC) - # Do not use find_package here since it would use VisualStudio's built-in OpenMP, but MKL libraries - # refer to Intel's OpenMP. - # - # No extra settings are needed for compilation: it only needs /openmp flag which is set further below, - # in the common MSVC compiler options block. - include_directories(BEFORE "${DISKANN_MSVC_PACKAGES}/intelopenmp.devel.win/lib/native/include") - link_libraries("${DISKANN_MSVC_PACKAGES}/intelopenmp.devel.win/lib/native/win-x64/libiomp5md.lib") - - set(OPENMP_WINDOWS_RUNTIME_FILES - "${DISKANN_MSVC_PACKAGES}/intelopenmp.redist.win/runtimes/win-x64/native/libiomp5md.dll" - "${DISKANN_MSVC_PACKAGES}/intelopenmp.redist.win/runtimes/win-x64/native/libiomp5md.pdb") -elseif(APPLE) - # Check if we're building Python bindings - if(PYBIND) - # First look for PyTorch's OpenMP to avoid conflicts - execute_process( - COMMAND ${Python_EXECUTABLE} -c "import os; import torch; print(os.path.join(os.path.dirname(torch.__file__), 'lib', 'libomp.dylib'))" - RESULT_VARIABLE TORCH_PATH_RESULT - OUTPUT_VARIABLE TORCH_LIBOMP_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE - ERROR_QUIET - ) - - execute_process( - COMMAND brew --prefix libomp - OUTPUT_VARIABLE LIBOMP_ROOT - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - - if(EXISTS "${TORCH_LIBOMP_PATH}") - message(STATUS "Found PyTorch's libomp: ${TORCH_LIBOMP_PATH}") - set(OpenMP_CXX_FLAGS "-Xclang -fopenmp") - set(OpenMP_C_FLAGS "-Xclang -fopenmp") - set(OpenMP_CXX_LIBRARIES "${TORCH_LIBOMP_PATH}") - set(OpenMP_C_LIBRARIES "${TORCH_LIBOMP_PATH}") - set(OpenMP_FOUND TRUE) - - include_directories(${LIBOMP_ROOT}/include) - - # Set compiler flags and link libraries - set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") - link_libraries("${TORCH_LIBOMP_PATH}") - else() - message(STATUS "No PyTorch's libomp found, falling back to normal OpenMP detection") - # Fallback to normal OpenMP detection - execute_process( - COMMAND brew --prefix libomp - OUTPUT_VARIABLE LIBOMP_ROOT - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - - set(OpenMP_ROOT "${LIBOMP_ROOT}") - find_package(OpenMP) - - if (OPENMP_FOUND) - set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") - set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") - link_libraries(OpenMP::OpenMP_CXX) - else() - message(FATAL_ERROR "No OpenMP support") - endif() - endif() - else() - # Regular OpenMP setup for non-Python builds - execute_process( - COMMAND brew --prefix libomp - OUTPUT_VARIABLE LIBOMP_ROOT - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - set(OpenMP_ROOT "${LIBOMP_ROOT}") - find_package(OpenMP) - - if (OPENMP_FOUND) - set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") - link_libraries(OpenMP::OpenMP_CXX) - else() - message(FATAL_ERROR "No OpenMP support") - endif() - endif() -else() - find_package(OpenMP) - - if (OPENMP_FOUND) - set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") - set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") - else() - message(FATAL_ERROR "No OpenMP support") - endif() -endif() - -# DiskANN core uses header-only libraries. Only DiskANN tools need program_options which has a linker library, -# but its size is small. Reduce number of dependent DLLs by linking statically. -if (MSVC) - set(Boost_USE_STATIC_LIBS ON) -endif() - -if(NOT MSVC) - find_package(Boost COMPONENTS program_options) -endif() - -# For Windows, fall back to nuget version if find_package didn't find it. -if (MSVC AND NOT Boost_FOUND) - set(DISKANN_BOOST_INCLUDE "${DISKANN_MSVC_PACKAGES}/boost/lib/native/include") - # Multi-threaded static library. - set(PROGRAM_OPTIONS_LIB_PATTERN "${DISKANN_MSVC_PACKAGES}/boost_program_options-vc${MSVC_TOOLSET_VERSION}/lib/native/libboost_program_options-vc${MSVC_TOOLSET_VERSION}-mt-x64-*.lib") - file(GLOB DISKANN_BOOST_PROGRAM_OPTIONS_LIB ${PROGRAM_OPTIONS_LIB_PATTERN}) - - set(PROGRAM_OPTIONS_DLIB_PATTERN "${DISKANN_MSVC_PACKAGES}/boost_program_options-vc${MSVC_TOOLSET_VERSION}/lib/native/libboost_program_options-vc${MSVC_TOOLSET_VERSION}-mt-gd-x64-*.lib") - file(GLOB DISKANN_BOOST_PROGRAM_OPTIONS_DLIB ${PROGRAM_OPTIONS_DLIB_PATTERN}) - - if (EXISTS ${DISKANN_BOOST_INCLUDE} AND EXISTS ${DISKANN_BOOST_PROGRAM_OPTIONS_LIB} AND EXISTS ${DISKANN_BOOST_PROGRAM_OPTIONS_DLIB}) - set(Boost_FOUND ON) - set(Boost_INCLUDE_DIR ${DISKANN_BOOST_INCLUDE}) - add_library(Boost::program_options STATIC IMPORTED) - set_target_properties(Boost::program_options PROPERTIES IMPORTED_LOCATION_RELEASE "${DISKANN_BOOST_PROGRAM_OPTIONS_LIB}") - set_target_properties(Boost::program_options PROPERTIES IMPORTED_LOCATION_DEBUG "${DISKANN_BOOST_PROGRAM_OPTIONS_DLIB}") - message(STATUS "Falling back to using Boost from the nuget package") - else() - message(WARNING "Couldn't find Boost. Was looking for ${DISKANN_BOOST_INCLUDE} and ${PROGRAM_OPTIONS_LIB_PATTERN}") - endif() -endif() - -if (NOT Boost_FOUND) - message(FATAL_ERROR "Couldn't find Boost dependency") -endif() - -include_directories(${Boost_INCLUDE_DIR}) - -#MKL Config -if (MSVC) - # Only the DiskANN DLL and one of the tools need MKL libraries. Additionally, only a small part of MKL is used. - # Given that and given that MKL DLLs are huge, use static linking to end up with no MKL DLL dependencies and with - # significantly smaller disk footprint. - # - # The compile options are not modified as there's already an unconditional -DMKL_ILP64 define below - # for all architectures, which is all that's needed. - set(DISKANN_MKL_INCLUDE_DIRECTORIES "${DISKANN_MSVC_PACKAGES}/intelmkl.static.win-x64/lib/native/include") - set(DISKANN_MKL_LIB_PATH "${DISKANN_MSVC_PACKAGES}/intelmkl.static.win-x64/lib/native/win-x64") - - set(DISKANN_MKL_LINK_LIBRARIES - "${DISKANN_MKL_LIB_PATH}/mkl_intel_ilp64.lib" - "${DISKANN_MKL_LIB_PATH}/mkl_core.lib" - "${DISKANN_MKL_LIB_PATH}/mkl_intel_thread.lib") -elseif(APPLE) - # no mkl on non-intel devices - find_library(ACCELERATE_LIBRARY Accelerate) - message(STATUS "Found Accelerate (${ACCELERATE_LIBRARY})") - set(DISKANN_ACCEL_LINK_OPTIONS ${ACCELERATE_LIBRARY}) - add_compile_definitions(ACCELERATE_NEW_LAPACK) -else() - # expected path for manual intel mkl installs - set(POSSIBLE_OMP_PATHS "/opt/intel/oneapi/compiler/2025.0/lib/libiomp5.so;/opt/intel/oneapi/compiler/latest/linux/compiler/lib/intel64_lin/libiomp5.so;/usr/lib/x86_64-linux-gnu/libiomp5.so;/opt/intel/lib/intel64_lin/libiomp5.so") - foreach(POSSIBLE_OMP_PATH ${POSSIBLE_OMP_PATHS}) - if (EXISTS ${POSSIBLE_OMP_PATH}) - get_filename_component(OMP_PATH ${POSSIBLE_OMP_PATH} DIRECTORY) - endif() - endforeach() - - if(NOT OMP_PATH) - message(FATAL_ERROR "Could not find Intel OMP in standard locations; use -DOMP_PATH to specify the install location for your environment") - endif() - link_directories(${OMP_PATH}) - - set(POSSIBLE_MKL_LIB_PATHS "/opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_core.so;/usr/lib/x86_64-linux-gnu/libmkl_core.so;/opt/intel/mkl/lib/intel64/libmkl_core.so") - foreach(POSSIBLE_MKL_LIB_PATH ${POSSIBLE_MKL_LIB_PATHS}) - if (EXISTS ${POSSIBLE_MKL_LIB_PATH}) - get_filename_component(MKL_PATH ${POSSIBLE_MKL_LIB_PATH} DIRECTORY) - endif() - endforeach() - - set(POSSIBLE_MKL_INCLUDE_PATHS "/opt/intel/oneapi/mkl/latest/include;/usr/include/mkl;/opt/intel/mkl/include/;") - foreach(POSSIBLE_MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATHS}) - if (EXISTS ${POSSIBLE_MKL_INCLUDE_PATH}) - set(MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATH}) - endif() - endforeach() - if(NOT MKL_PATH) - message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_PATH to specify the install location for your environment") - elseif(NOT MKL_INCLUDE_PATH) - message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_INCLUDE_PATH to specify the install location for headers for your environment") - endif() - if (EXISTS ${MKL_PATH}/libmkl_def.so.2) - set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so.2) - elseif(EXISTS ${MKL_PATH}/libmkl_def.so) - set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so) - else() - message(FATAL_ERROR "Despite finding MKL, libmkl_def.so was not found in expected locations.") - endif() - link_directories(${MKL_PATH}) - include_directories(${MKL_INCLUDE_PATH}) - - # compile flags and link libraries - # if gcc/g++ - if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - add_compile_options(-m64 -Wl,--no-as-needed) - endif() - if (NOT PYBIND) - link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core iomp5 pthread m dl) - else() - # static linking for python so as to minimize customer dependency issues - if (CMAKE_BUILD_TYPE STREQUAL "Debug") - # In debug mode, use dynamic linking to ensure all symbols are available - link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core ${MKL_DEF_SO} iomp5 pthread m dl) - else() - # In release mode, use static linking to minimize dependencies - link_libraries( - ${MKL_PATH}/libmkl_intel_ilp64.a - ${MKL_PATH}/libmkl_intel_thread.a - ${MKL_PATH}/libmkl_core.a - ${MKL_DEF_SO} - iomp5 - pthread - m - dl - ) - endif() - endif() - - add_definitions(-DMKL_ILP64) -endif() - - -# Section for tcmalloc. The DiskANN tools are always linked to tcmalloc. For Windows, they also need to -# force-include the _tcmalloc symbol for enabling tcmalloc. -# -# The DLL itself needs to be linked to tcmalloc only if DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS -# is enabled. -if(USE_TCMALLOC) - if (MSVC) - if (NOT EXISTS "${PROJECT_SOURCE_DIR}/gperftools/gperftools.sln") - message(FATAL_ERROR "The gperftools submodule was not found. " - "Please check-out git submodules by doing 'git submodule init' followed by 'git submodule update'") - endif() - - set(TCMALLOC_LINK_LIBRARY "${PROJECT_SOURCE_DIR}/gperftools/x64/Release-Patch/libtcmalloc_minimal.lib") - set(TCMALLOC_WINDOWS_RUNTIME_FILES - "${PROJECT_SOURCE_DIR}/gperftools/x64/Release-Patch/libtcmalloc_minimal.dll" - "${PROJECT_SOURCE_DIR}/gperftools/x64/Release-Patch/libtcmalloc_minimal.pdb") - - # Tell CMake how to build the tcmalloc linker library from the submodule. - add_custom_target(build_libtcmalloc_minimal DEPENDS ${TCMALLOC_LINK_LIBRARY}) - add_custom_command(OUTPUT ${TCMALLOC_LINK_LIBRARY} - COMMAND ${CMAKE_VS_MSBUILD_COMMAND} gperftools.sln /m /nologo - /t:libtcmalloc_minimal /p:Configuration="Release-Patch" - /property:Platform="x64" - /p:PlatformToolset=v${MSVC_TOOLSET_VERSION} - /p:WindowsTargetPlatformVersion=${CMAKE_VS_WINDOWS_TARGET_PLATFORM_VERSION} - WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/gperftools) - - add_library(libtcmalloc_minimal_for_exe STATIC IMPORTED) - add_library(libtcmalloc_minimal_for_dll STATIC IMPORTED) - - set_target_properties(libtcmalloc_minimal_for_dll PROPERTIES - IMPORTED_LOCATION "${TCMALLOC_LINK_LIBRARY}") - - set_target_properties(libtcmalloc_minimal_for_exe PROPERTIES - IMPORTED_LOCATION "${TCMALLOC_LINK_LIBRARY}" - INTERFACE_LINK_OPTIONS /INCLUDE:_tcmalloc) - - # Ensure libtcmalloc_minimal is built before it's being used. - add_dependencies(libtcmalloc_minimal_for_dll build_libtcmalloc_minimal) - add_dependencies(libtcmalloc_minimal_for_exe build_libtcmalloc_minimal) - - set(DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS libtcmalloc_minimal_for_exe) - elseif(APPLE) # ! Inherited from #474, not been adjusted for TCMalloc Removal - execute_process( - COMMAND brew --prefix gperftools - OUTPUT_VARIABLE GPERFTOOLS_PREFIX - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - set(DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS "-L${GPERFTOOLS_PREFIX}/lib -ltcmalloc") - elseif(NOT PYBIND) - set(DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS "-ltcmalloc") - endif() - - if (DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) - add_definitions(-DRELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) - - if (MSVC) - set(DISKANN_DLL_TCMALLOC_LINK_OPTIONS libtcmalloc_minimal_for_dll) - endif() - endif() -endif() - -if (NOT MSVC AND NOT APPLE) - set(DISKANN_ASYNC_LIB aio) -endif() - -#Main compiler/linker settings -if(MSVC) - #language options - add_compile_options(/permissive- /openmp:experimental /Zc:twoPhase- /Zc:inline /WX- /std:c++17 /Gd /W3 /MP /Zi /FC /nologo) - #code generation options - add_compile_options(/arch:AVX2 /fp:fast /fp:except- /EHsc /GS- /Gy) - #optimization options - add_compile_options(/Ot /Oy /Oi) - #path options - add_definitions(-DUSE_AVX2 -DUSE_ACCELERATED_PQ -D_WINDOWS -DNOMINMAX -DUNICODE) - # Linker options. Exclude VCOMP/VCOMPD.LIB which contain VisualStudio's version of OpenMP. - # MKL was linked against Intel's OpenMP and depends on the corresponding DLL. - add_link_options(/NODEFAULTLIB:VCOMP.LIB /NODEFAULTLIB:VCOMPD.LIB /DEBUG:FULL /OPT:REF /OPT:ICF) - - set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug) - set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug) - set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug) - - set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release) - set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release) - set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release) -elseif(APPLE) - set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -ftree-vectorize -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -Xclang -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -Wno-inconsistent-missing-override -Wno-return-type") - set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG") - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -Ofast -DNDEBUG -ftree-vectorize") - if (NOT PYBIND) - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -Ofast") - if (NOT PORTABLE) - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -mtune=native") - endif() - else() - # -Ofast is not supported in a python extension module - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -fPIC") - endif() -else() - set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mavx2 -mfma -msse2 -ftree-vectorize -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_AVX2 -fPIC") - if(USE_TCMALLOC) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free") - endif() - set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG") - if (NOT PYBIND) - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -Ofast") - if (NOT PORTABLE) - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -march=native -mtune=native") - endif() - else() - # -Ofast is not supported in a python extension module - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG") - endif() -endif() - -add_subdirectory(src) -if (NOT PYBIND) - add_subdirectory(apps) - add_subdirectory(apps/utils) -endif() - -if (UNIT_TEST) - enable_testing() - add_subdirectory(tests) -endif() - -if (MSVC) - message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" - "Alternatively, use MSBuild to build:\n\n" - "msbuild.exe ${PROJECT_NAME}.sln /m /nologo /t:Build /p:Configuration=\"Release\" /property:Platform=\"x64\"\n") -endif() - -if (RESTAPI) - if (MSVC) - set(DISKANN_CPPRESTSDK "${DISKANN_MSVC_PACKAGES}/cpprestsdk.v142/build/native") - # expected path for apt packaged intel mkl installs - link_libraries("${DISKANN_CPPRESTSDK}/x64/lib/cpprest142_2_10.lib") - include_directories("${DISKANN_CPPRESTSDK}/include") - endif() - add_subdirectory(apps/restapi) -endif() - -include(clang-format.cmake) - -if(PYBIND) - add_subdirectory(python) - - install(TARGETS _diskannpy - DESTINATION leann_backend_diskann - COMPONENT python_modules - ) - -endif() -############################################################################### -# PROTOBUF SECTION - Corrected to use CONFIG mode explicitly -############################################################################### -set(Protobuf_USE_STATIC_LIBS OFF) - -find_package(ZLIB REQUIRED) - -find_package(Protobuf REQUIRED) - -message(STATUS "Protobuf found: ${Protobuf_VERSION}") -message(STATUS "Protobuf include dirs: ${Protobuf_INCLUDE_DIRS}") -message(STATUS "Protobuf libraries: ${Protobuf_LIBRARIES}") -message(STATUS "Protobuf protoc executable: ${Protobuf_PROTOC_EXECUTABLE}") - -include_directories(${Protobuf_INCLUDE_DIRS}) - -set(PROTO_FILE "${CMAKE_CURRENT_SOURCE_DIR}/../embedding.proto") -protobuf_generate_cpp(PROTO_SRCS PROTO_HDRS ${PROTO_FILE}) -set(generated_proto_sources ${PROTO_SRCS}) - - -add_library(proto_embeddings STATIC ${generated_proto_sources}) -target_link_libraries(proto_embeddings PUBLIC protobuf::libprotobuf) -target_include_directories(proto_embeddings PUBLIC - ${CMAKE_CURRENT_BINARY_DIR} - ${Protobuf_INCLUDE_DIRS} -) - -target_link_libraries(diskann PRIVATE proto_embeddings protobuf::libprotobuf) -target_include_directories(diskann PRIVATE - ${CMAKE_CURRENT_BINARY_DIR} - ${Protobuf_INCLUDE_DIRS} -) - -target_link_libraries(diskann_s PRIVATE proto_embeddings protobuf::libprotobuf) -target_include_directories(diskann_s PRIVATE - ${CMAKE_CURRENT_BINARY_DIR} - ${Protobuf_INCLUDE_DIRS} -) - - -############################################################################### -# ZEROMQ SECTION - REQUIRED -############################################################################### - -find_package(ZeroMQ QUIET) -if(NOT ZeroMQ_FOUND) - find_path(ZeroMQ_INCLUDE_DIR zmq.h) - find_library(ZeroMQ_LIBRARY zmq) - if(ZeroMQ_INCLUDE_DIR AND ZeroMQ_LIBRARY) - set(ZeroMQ_FOUND TRUE) - endif() -endif() - -if(ZeroMQ_FOUND) - message(STATUS "Found ZeroMQ: ${ZeroMQ_LIBRARY}") - include_directories(${ZeroMQ_INCLUDE_DIR}) - target_link_libraries(diskann PRIVATE ${ZeroMQ_LIBRARY}) - target_link_libraries(diskann_s PRIVATE ${ZeroMQ_LIBRARY}) - add_definitions(-DUSE_ZEROMQ) -else() - message(FATAL_ERROR "ZeroMQ is required but not found. Please install ZeroMQ and try again.") -endif() - -target_link_libraries(diskann ${PYBIND11_LIBRARIES}) -target_link_libraries(diskann_s ${PYBIND11_LIBRARIES}) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/CMakeSettings.json b/packages/leann-backend-diskann/third_party/DiskANN/CMakeSettings.json deleted file mode 100644 index af5d7b5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/CMakeSettings.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "configurations": [ - { - "name": "x64-Release", - "generator": "Ninja", - "configurationType": "Release", - "inheritEnvironments": [ "msvc_x64" ], - "buildRoot": "${projectDir}\\out\\build\\${name}", - "installRoot": "${projectDir}\\out\\install\\${name}", - "cmakeCommandArgs": "", - "buildCommandArgs": "", - "ctestCommandArgs": "" - }, - { - "name": "WSL-GCC-Release", - "generator": "Ninja", - "configurationType": "RelWithDebInfo", - "buildRoot": "${projectDir}\\out\\build\\${name}", - "installRoot": "${projectDir}\\out\\install\\${name}", - "cmakeExecutable": "cmake", - "cmakeCommandArgs": "", - "buildCommandArgs": "", - "ctestCommandArgs": "", - "inheritEnvironments": [ "linux_x64" ], - "wslPath": "${defaultWSLPath}" - } - ] -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/CODE_OF_CONDUCT.md b/packages/leann-backend-diskann/third_party/DiskANN/CODE_OF_CONDUCT.md deleted file mode 100644 index f9ba8cf..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/CODE_OF_CONDUCT.md +++ /dev/null @@ -1,9 +0,0 @@ -# Microsoft Open Source Code of Conduct - -This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). - -Resources: - -- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/) -- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) -- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns diff --git a/packages/leann-backend-diskann/third_party/DiskANN/CONTRIBUTING.md b/packages/leann-backend-diskann/third_party/DiskANN/CONTRIBUTING.md deleted file mode 100644 index dcbf795..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/CONTRIBUTING.md +++ /dev/null @@ -1,9 +0,0 @@ -# Contributing - -This project welcomes contributions and suggestions. Most contributions require you to agree to a -Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us -the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com. - -When you submit a pull request, a CLA bot will automatically determine whether you need to provide -a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions -provided by the bot. You will only need to do this once across all repos using our CLA. diff --git a/packages/leann-backend-diskann/third_party/DiskANN/Dockerfile b/packages/leann-backend-diskann/third_party/DiskANN/Dockerfile deleted file mode 100644 index ea1979f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/Dockerfile +++ /dev/null @@ -1,17 +0,0 @@ -#Copyright(c) Microsoft Corporation.All rights reserved. -#Licensed under the MIT license. - -FROM ubuntu:jammy - -RUN apt update -RUN apt install -y software-properties-common -RUN add-apt-repository -y ppa:git-core/ppa -RUN apt update -RUN DEBIAN_FRONTEND=noninteractive apt install -y git make cmake g++ libaio-dev libgoogle-perftools-dev libunwind-dev clang-format libboost-dev libboost-program-options-dev libmkl-full-dev libcpprest-dev python3.10 - -WORKDIR /app -RUN git clone https://github.com/microsoft/DiskANN.git -WORKDIR /app/DiskANN -RUN mkdir build -RUN cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -RUN cmake --build build -- -j diff --git a/packages/leann-backend-diskann/third_party/DiskANN/DockerfileDev b/packages/leann-backend-diskann/third_party/DiskANN/DockerfileDev deleted file mode 100644 index 0e95e40..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/DockerfileDev +++ /dev/null @@ -1,17 +0,0 @@ -#Copyright(c) Microsoft Corporation.All rights reserved. -#Licensed under the MIT license. - -FROM ubuntu:jammy - -RUN apt update -RUN apt install -y software-properties-common -RUN add-apt-repository -y ppa:git-core/ppa -RUN apt update -RUN DEBIAN_FRONTEND=noninteractive apt install -y git make cmake g++ libaio-dev libgoogle-perftools-dev libunwind-dev clang-format libboost-dev libboost-program-options-dev libboost-test-dev libmkl-full-dev libcpprest-dev python3.10 - -WORKDIR /app -RUN git clone https://github.com/microsoft/DiskANN.git -WORKDIR /app/DiskANN -RUN mkdir build -RUN cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DUNIT_TEST=True -RUN cmake --build build -- -j diff --git a/packages/leann-backend-diskann/third_party/DiskANN/LICENSE b/packages/leann-backend-diskann/third_party/DiskANN/LICENSE deleted file mode 100644 index b7a909e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/LICENSE +++ /dev/null @@ -1,23 +0,0 @@ - DiskANN - - MIT License - - Copyright (c) Microsoft Corporation. - - Permission is hereby granted, free of charge, to any person obtaining a copy - of this software and associated documentation files (the "Software"), to deal - in the Software without restriction, including without limitation the rights - to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - copies of the Software, and to permit persons to whom the Software is - furnished to do so, subject to the following conditions: - - The above copyright notice and this permission notice shall be included in all - copies or substantial portions of the Software. - - THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - SOFTWARE diff --git a/packages/leann-backend-diskann/third_party/DiskANN/MANIFEST.in b/packages/leann-backend-diskann/third_party/DiskANN/MANIFEST.in deleted file mode 100644 index 0735c27..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/MANIFEST.in +++ /dev/null @@ -1,12 +0,0 @@ -include MANIFEST.in -include *.txt -include *.md -include setup.py -include pyproject.toml -include *.cmake -recursive-include gperftools * -recursive-include include * -recursive-include python * -recursive-include windows * -prune python/tests -recursive-include src * diff --git a/packages/leann-backend-diskann/third_party/DiskANN/README.md b/packages/leann-backend-diskann/third_party/DiskANN/README.md deleted file mode 100644 index 44f4c27..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/README.md +++ /dev/null @@ -1,135 +0,0 @@ -# DiskANN - -[![DiskANN Main](https://github.com/microsoft/DiskANN/actions/workflows/push-test.yml/badge.svg?branch=main)](https://github.com/microsoft/DiskANN/actions/workflows/push-test.yml) -[![PyPI version](https://img.shields.io/pypi/v/diskannpy.svg)](https://pypi.org/project/diskannpy/) -[![Downloads shield](https://pepy.tech/badge/diskannpy)](https://pepy.tech/project/diskannpy) -[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) - -[![DiskANN Paper](https://img.shields.io/badge/Paper-NeurIPS%3A_DiskANN-blue)](https://papers.nips.cc/paper/9527-rand-nsg-fast-accurate-billion-point-nearest-neighbor-search-on-a-single-node.pdf) -[![DiskANN Paper](https://img.shields.io/badge/Paper-Arxiv%3A_Fresh--DiskANN-blue)](https://arxiv.org/abs/2105.09613) -[![DiskANN Paper](https://img.shields.io/badge/Paper-Filtered--DiskANN-blue)](https://harsha-simhadri.org/pubs/Filtered-DiskANN23.pdf) - - -DiskANN is a suite of scalable, accurate and cost-effective approximate nearest neighbor search algorithms for large-scale vector search that support real-time changes and simple filters. -This code is based on ideas from the [DiskANN](https://papers.nips.cc/paper/9527-rand-nsg-fast-accurate-billion-point-nearest-neighbor-search-on-a-single-node.pdf), [Fresh-DiskANN](https://arxiv.org/abs/2105.09613) and the [Filtered-DiskANN](https://harsha-simhadri.org/pubs/Filtered-DiskANN23.pdf) papers with further improvements. -This code forked off from [code for NSG](https://github.com/ZJULearning/nsg) algorithm. - -This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). -For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or -contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. - -See [guidelines](CONTRIBUTING.md) for contributing to this project. - -## Linux build: - -Install the following packages through apt-get - -```bash -sudo apt install make cmake g++ libaio-dev libgoogle-perftools-dev clang-format libboost-all-dev -``` - -### Install Intel MKL -#### Ubuntu 20.04 or newer -```bash -sudo apt install libmkl-full-dev -``` - -#### Earlier versions of Ubuntu -Install Intel MKL either by downloading the [oneAPI MKL installer](https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html) or using [apt](https://software.intel.com/en-us/articles/installing-intel-free-libs-and-python-apt-repo) (we tested with build 2019.4-070 and 2022.1.2.146). - -``` -# OneAPI MKL Installer -wget https://registrationcenter-download.intel.com/akdlm/irc_nas/18487/l_BaseKit_p_2022.1.2.146.sh -sudo sh l_BaseKit_p_2022.1.2.146.sh -a --components intel.oneapi.lin.mkl.devel --action install --eula accept -s -``` - -### Build -```bash -mkdir build && cd build && cmake -DCMAKE_BUILD_TYPE=Release .. && make -j -``` - -## Windows build: - -The Windows version has been tested with Enterprise editions of Visual Studio 2022, 2019 and 2017. It should work with the Community and Professional editions as well without any changes. - -**Prerequisites:** - -* CMake 3.15+ (available in VisualStudio 2019+ or from https://cmake.org) -* NuGet.exe (install from https://www.nuget.org/downloads) - * The build script will use NuGet to get MKL, OpenMP and Boost packages. -* DiskANN git repository checked out together with submodules. To check out submodules after git clone: -``` -git submodule init -git submodule update -``` - -* Environment variables: - * [optional] If you would like to override the Boost library listed in windows/packages.config.in, set BOOST_ROOT to your Boost folder. - -**Build steps:** -* Open the "x64 Native Tools Command Prompt for VS 2019" (or corresponding version) and change to DiskANN folder -* Create a "build" directory inside it -* Change to the "build" directory and run -``` -cmake .. -``` -OR for Visual Studio 2017 and earlier: -``` -\cmake .. -``` -**This will create a diskann.sln solution**. Now you can: - -- Open it from VisualStudio and build either Release or Debug configuration. -- `\cmake --build build` -- Use MSBuild: -``` -msbuild.exe diskann.sln /m /nologo /t:Build /p:Configuration="Release" /property:Platform="x64" -``` - -* This will also build gperftools submodule for libtcmalloc_minimal dependency. -* Generated binaries are stored in the x64/Release or x64/Debug directories. - -## macOS Build - -### Prerequisites -* Apple Silicon. The code should still work on Intel-based Macs, but there are no guarantees. -* macOS >= 12.0 -* XCode Command Line Tools (install with `xcode-select --install`) -* [homebrew](https://brew.sh/) - -### Install Required Packages -```zsh -brew install cmake -brew install boost -brew install gperftools -brew install libomp -``` - -### Build DiskANN -```zsh -# same as ubuntu instructions -mkdir build && cd build && cmake -DCMAKE_BUILD_TYPE=Release .. && make -j -``` - -## Usage: - -Please see the following pages on using the compiled code: - -- [Commandline interface for building and search SSD based indices](workflows/SSD_index.md) -- [Commandline interface for building and search in memory indices](workflows/in_memory_index.md) -- [Commandline examples for using in-memory streaming indices](workflows/dynamic_index.md) -- [Commandline interface for building and search in memory indices with label data and filters](workflows/filtered_in_memory.md) -- [Commandline interface for building and search SSD based indices with label data and filters](workflows/filtered_ssd_index.md) -- [diskannpy - DiskANN as a python extension module](python/README.md) - -Please cite this software in your work as: - -``` -@misc{diskann-github, - author = {Simhadri, Harsha Vardhan and Krishnaswamy, Ravishankar and Srinivasa, Gopal and Subramanya, Suhas Jayaram and Antonijevic, Andrija and Pryce, Dax and Kaczynski, David and Williams, Shane and Gollapudi, Siddarth and Sivashankar, Varun and Karia, Neel and Singh, Aditi and Jaiswal, Shikhar and Mahapatro, Neelam and Adams, Philip and Tower, Bryan and Patel, Yash}}, - title = {{DiskANN: Graph-structured Indices for Scalable, Fast, Fresh and Filtered Approximate Nearest Neighbor Search}}, - url = {https://github.com/Microsoft/DiskANN}, - version = {0.6.1}, - year = {2023} -} -``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/SECURITY.md b/packages/leann-backend-diskann/third_party/DiskANN/SECURITY.md deleted file mode 100644 index f7b8998..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/SECURITY.md +++ /dev/null @@ -1,41 +0,0 @@ - - -## Security - -Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet), [Xamarin](https://github.com/xamarin), and [our GitHub organizations](https://opensource.microsoft.com/). - -If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://docs.microsoft.com/en-us/previous-versions/tn-archive/cc751383(v=technet.10)), please report it to us as described below. - -## Reporting Security Issues - -**Please do not report security vulnerabilities through public GitHub issues.** - -Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://msrc.microsoft.com/create-report). - -If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://www.microsoft.com/en-us/msrc/pgp-key-msrc). - -You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc). - -Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue: - - * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.) - * Full paths of source file(s) related to the manifestation of the issue - * The location of the affected source code (tag/branch/commit or direct URL) - * Any special configuration required to reproduce the issue - * Step-by-step instructions to reproduce the issue - * Proof-of-concept or exploit code (if possible) - * Impact of the issue, including how an attacker might exploit the issue - -This information will help us triage your report more quickly. - -If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://microsoft.com/msrc/bounty) page for more details about our active programs. - -## Preferred Languages - -We prefer all communications to be in English. - -## Policy - -Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://www.microsoft.com/en-us/msrc/cvd). - - \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/apps/CMakeLists.txt deleted file mode 100644 index e42c0b6..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/CMakeLists.txt +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_COMPILE_WARNING_AS_ERROR ON) - -add_executable(build_memory_index build_memory_index.cpp) -target_link_libraries(build_memory_index ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(build_stitched_index build_stitched_index.cpp) -target_link_libraries(build_stitched_index ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(search_memory_index search_memory_index.cpp) -target_link_libraries(search_memory_index ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(build_disk_index build_disk_index.cpp) -target_link_libraries(build_disk_index ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} ${DISKANN_ASYNC_LIB} Boost::program_options) - -add_executable(search_disk_index search_disk_index.cpp) -target_link_libraries(search_disk_index ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(range_search_disk_index range_search_disk_index.cpp) -target_link_libraries(range_search_disk_index ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(test_streaming_scenario test_streaming_scenario.cpp) -target_link_libraries(test_streaming_scenario ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -add_executable(test_insert_deletes_consolidate test_insert_deletes_consolidate.cpp) -target_link_libraries(test_insert_deletes_consolidate ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options) - -if (NOT MSVC) - install(TARGETS build_memory_index - build_stitched_index - search_memory_index - build_disk_index - search_disk_index - range_search_disk_index - test_streaming_scenario - test_insert_deletes_consolidate - RUNTIME - ) -endif() diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_disk_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/build_disk_index.cpp deleted file mode 100644 index f48b617..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_disk_index.cpp +++ /dev/null @@ -1,191 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include - -#include "utils.h" -#include "disk_utils.h" -#include "math_utils.h" -#include "index.h" -#include "partition.h" -#include "program_options_utils.hpp" - -namespace po = boost::program_options; - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix, codebook_prefix, label_file, universal_label, - label_type; - uint32_t num_threads, R, L, disk_PQ, build_PQ, QD, Lf, filter_threshold; - float B, M; - bool append_reorder_data = false; - bool use_opq = false; - - po::options_description desc{ - program_options_utils::make_program_description("build_disk_index", "Build a disk-based index.")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("data_path", po::value(&data_path)->required(), - program_options_utils::INPUT_DATA_PATH); - required_configs.add_options()("search_DRAM_budget,B", po::value(&B)->required(), - "DRAM budget in GB for searching the index to set the " - "compressed level for data while search happens"); - required_configs.add_options()("build_DRAM_budget,M", po::value(&M)->required(), - "DRAM budget in GB for building the index"); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("max_degree,R", po::value(&R)->default_value(64), - program_options_utils::MAX_BUILD_DEGREE); - optional_configs.add_options()("Lbuild,L", po::value(&L)->default_value(100), - program_options_utils::GRAPH_BUILD_COMPLEXITY); - optional_configs.add_options()("QD", po::value(&QD)->default_value(0), - " Quantized Dimension for compression"); - optional_configs.add_options()("codebook_prefix", po::value(&codebook_prefix)->default_value(""), - "Path prefix for pre-trained codebook"); - optional_configs.add_options()("PQ_disk_bytes", po::value(&disk_PQ)->default_value(0), - "Number of bytes to which vectors should be compressed " - "on SSD; 0 for no compression"); - optional_configs.add_options()("append_reorder_data", po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD."); - optional_configs.add_options()("build_PQ_bytes", po::value(&build_PQ)->default_value(0), - program_options_utils::BUIlD_GRAPH_PQ_BYTES); - optional_configs.add_options()("use_opq", po::bool_switch()->default_value(false), - program_options_utils::USE_OPQ); - optional_configs.add_options()("label_file", po::value(&label_file)->default_value(""), - program_options_utils::LABEL_FILE); - optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), - program_options_utils::UNIVERSAL_LABEL); - optional_configs.add_options()("FilteredLbuild", po::value(&Lf)->default_value(0), - program_options_utils::FILTERED_LBUILD); - optional_configs.add_options()("filter_threshold,F", po::value(&filter_threshold)->default_value(0), - "Threshold to break up the existing nodes to generate new graph " - "internally where each node has a maximum F labels."); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - program_options_utils::LABEL_TYPE_DESCRIPTION); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (vm["append_reorder_data"].as()) - append_reorder_data = true; - if (vm["use_opq"].as()) - use_opq = true; - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - bool use_filters = (label_file != "") ? true : false; - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else if (dist_fn == std::string("cosine")) - metric = diskann::Metric::COSINE; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } - - if (append_reorder_data) - { - if (disk_PQ == 0) - { - std::cout << "Error: It is not necessary to append data for reordering " - "when vectors are not compressed on disk." - << std::endl; - return -1; - } - if (data_type != std::string("float")) - { - std::cout << "Error: Appending data for reordering currently only " - "supported for float data type." - << std::endl; - return -1; - } - } - - std::string params = std::string(std::to_string(R)) + " " + std::string(std::to_string(L)) + " " + - std::string(std::to_string(B)) + " " + std::string(std::to_string(M)) + " " + - std::string(std::to_string(num_threads)) + " " + std::string(std::to_string(disk_PQ)) + " " + - std::string(std::to_string(append_reorder_data)) + " " + - std::string(std::to_string(build_PQ)) + " " + std::string(std::to_string(QD)); - - try - { - if (label_file != "" && label_type == "ushort") - { - if (data_type == std::string("int8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, - use_filters, label_file, universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, - use_filters, label_file, universal_label, filter_threshold, Lf); - else - { - diskann::cerr << "Error. Unsupported data type" << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("int8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else - { - diskann::cerr << "Error. Unsupported data type" << std::endl; - return -1; - } - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_memory_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/build_memory_index.cpp deleted file mode 100644 index 544e42d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_memory_index.cpp +++ /dev/null @@ -1,164 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include - -#include "index.h" -#include "utils.h" -#include "program_options_utils.hpp" - -#ifndef _WINDOWS -#include -#include -#else -#include -#endif - -#include "memory_mapper.h" -#include "ann_exception.h" -#include "index_factory.h" - -namespace po = boost::program_options; - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type; - uint32_t num_threads, R, L, Lf, build_PQ_bytes; - float alpha; - bool use_pq_build, use_opq; - - po::options_description desc{ - program_options_utils::make_program_description("build_memory_index", "Build a memory-based DiskANN index.")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("data_path", po::value(&data_path)->required(), - program_options_utils::INPUT_DATA_PATH); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("max_degree,R", po::value(&R)->default_value(64), - program_options_utils::MAX_BUILD_DEGREE); - optional_configs.add_options()("Lbuild,L", po::value(&L)->default_value(100), - program_options_utils::GRAPH_BUILD_COMPLEXITY); - optional_configs.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - program_options_utils::GRAPH_BUILD_ALPHA); - optional_configs.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), - program_options_utils::BUIlD_GRAPH_PQ_BYTES); - optional_configs.add_options()("use_opq", po::bool_switch()->default_value(false), - program_options_utils::USE_OPQ); - optional_configs.add_options()("label_file", po::value(&label_file)->default_value(""), - program_options_utils::LABEL_FILE); - optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), - program_options_utils::UNIVERSAL_LABEL); - - optional_configs.add_options()("FilteredLbuild", po::value(&Lf)->default_value(0), - program_options_utils::FILTERED_LBUILD); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - program_options_utils::LABEL_TYPE_DESCRIPTION); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - use_pq_build = (build_PQ_bytes > 0); - use_opq = vm["use_opq"].as(); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - - try - { - diskann::cout << "Starting index build with R: " << R << " Lbuild: " << L << " alpha: " << alpha - << " #threads: " << num_threads << std::endl; - - size_t data_num, data_dim; - diskann::get_bin_metadata(data_path, data_num, data_dim); - - auto index_build_params = diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - auto filter_params = diskann::IndexFilterParamsBuilder() - .with_universal_label(universal_label) - .with_label_file(label_file) - .with_save_path_prefix(index_path_prefix) - .build(); - auto config = diskann::IndexConfigBuilder() - .with_metric(metric) - .with_dimension(data_dim) - .with_max_points(data_num) - .with_data_load_store_strategy(diskann::DataStoreStrategy::MEMORY) - .with_graph_load_store_strategy(diskann::GraphStoreStrategy::MEMORY) - .with_data_type(data_type) - .with_label_type(label_type) - .is_dynamic_index(false) - .with_index_write_params(index_build_params) - .is_enable_tags(false) - .is_use_opq(use_opq) - .is_pq_dist_build(use_pq_build) - .with_num_pq_chunks(build_PQ_bytes) - .build(); - - auto index_factory = diskann::IndexFactory(config); - auto index = index_factory.create_instance(); - index->build(data_path, data_num, filter_params); - index->save(index_path_prefix.c_str()); - index.reset(); - return 0; - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_stitched_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/build_stitched_index.cpp deleted file mode 100644 index 60e38c1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/build_stitched_index.cpp +++ /dev/null @@ -1,441 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include "filter_utils.h" -#include -#ifndef _WINDOWS -#include -#endif - -#include "index.h" -#include "memory_mapper.h" -#include "parameters.h" -#include "utils.h" -#include "program_options_utils.hpp" - -namespace po = boost::program_options; -typedef std::tuple>, uint64_t> stitch_indices_return_values; - -/* - * Inline function to display progress bar. - */ -inline void print_progress(double percentage) -{ - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); -} - -/* - * Inline function to generate a random integer in a range. - */ -inline size_t random(size_t range_from, size_t range_to) -{ - std::random_device rand_dev; - std::mt19937 generator(rand_dev()); - std::uniform_int_distribution distr(range_from, range_to); - return distr(generator); -} - -/* - * function to handle command line parsing. - * - * Arguments are merely the inputs from the command line. - */ -void handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, path &final_index_path_prefix, - path &label_data_path, std::string &universal_label, uint32_t &num_threads, uint32_t &R, uint32_t &L, - uint32_t &stitched_R, float &alpha) -{ - po::options_description desc{ - program_options_utils::make_program_description("build_stitched_index", "Build a stitched DiskANN index.")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("index_path_prefix", - po::value(&final_index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("data_path", po::value(&input_data_path)->required(), - program_options_utils::INPUT_DATA_PATH); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("max_degree,R", po::value(&R)->default_value(64), - program_options_utils::MAX_BUILD_DEGREE); - optional_configs.add_options()("Lbuild,L", po::value(&L)->default_value(100), - program_options_utils::GRAPH_BUILD_COMPLEXITY); - optional_configs.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - program_options_utils::GRAPH_BUILD_ALPHA); - optional_configs.add_options()("label_file", po::value(&label_data_path)->default_value(""), - program_options_utils::LABEL_FILE); - optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), - program_options_utils::UNIVERSAL_LABEL); - optional_configs.add_options()("stitched_R", po::value(&stitched_R)->default_value(100), - "Degree to prune final graph down to"); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - exit(0); - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - throw; - } -} - -/* - * Custom index save to write the in-memory index to disk. - * Also writes required files for diskANN API - - * 1. labels_to_medoids - * 2. universal_label - * 3. data (redundant for static indices) - * 4. labels (redundant for static indices) - */ -void save_full_index(path final_index_path_prefix, path input_data_path, uint64_t final_index_size, - std::vector> stitched_graph, - tsl::robin_map entry_points, std::string universal_label, - path label_data_path) -{ - // aux. file 1 - auto saving_index_timer = std::chrono::high_resolution_clock::now(); - std::ifstream original_label_data_stream; - original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_label_data_stream.open(label_data_path, std::ios::binary); - std::ofstream new_label_data_stream; - new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_label_data_stream.open(final_index_path_prefix + "_labels.txt", std::ios::binary); - new_label_data_stream << original_label_data_stream.rdbuf(); - original_label_data_stream.close(); - new_label_data_stream.close(); - - // aux. file 2 - std::ifstream original_input_data_stream; - original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_input_data_stream.open(input_data_path, std::ios::binary); - std::ofstream new_input_data_stream; - new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary); - new_input_data_stream << original_input_data_stream.rdbuf(); - original_input_data_stream.close(); - new_input_data_stream.close(); - - // aux. file 3 - std::ofstream labels_to_medoids_writer; - labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit); - labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt"); - for (auto iter : entry_points) - labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl; - labels_to_medoids_writer.close(); - - // aux. file 4 (only if we're using a universal label) - if (universal_label != "") - { - std::ofstream universal_label_writer; - universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit); - universal_label_writer.open(final_index_path_prefix + "_universal_label.txt"); - universal_label_writer << universal_label << std::endl; - universal_label_writer.close(); - } - - // main index - uint64_t index_num_frozen_points = 0, index_num_edges = 0; - uint32_t index_max_observed_degree = 0, index_entry_point = 0; - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - for (auto &point_neighbors : stitched_graph) - { - index_max_observed_degree = std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); - } - - std::ofstream stitched_graph_writer; - stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); - stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); - - stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); - stitched_graph_writer.write((char *)&index_max_observed_degree, sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_num_frozen_points, sizeof(uint64_t)); - - size_t bytes_written = METADATA; - for (uint32_t node_point = 0; node_point < stitched_graph.size(); node_point++) - { - uint32_t current_node_num_neighbors = (uint32_t)stitched_graph[node_point].size(); - std::vector current_node_neighbors = stitched_graph[node_point]; - stitched_graph_writer.write((char *)¤t_node_num_neighbors, sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); - for (const auto ¤t_node_neighbor : current_node_neighbors) - { - stitched_graph_writer.write((char *)¤t_node_neighbor, sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); - } - index_num_edges += current_node_num_neighbors; - } - - if (bytes_written != final_index_size) - { - std::cerr << "Error: written bytes does not match allocated space" << std::endl; - throw; - } - - stitched_graph_writer.close(); - - std::chrono::duration saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer; - std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl; - std::cout << "Stitched graph average degree: " << ((float)index_num_edges) / ((float)(stitched_graph.size())) - << std::endl; - std::cout << "Stitched graph max degree: " << index_max_observed_degree << std::endl << std::endl; -} - -/* - * Unions the per-label graph indices together via the following policy: - * - any two nodes can only have at most one edge between them - - * - * Returns the "stitched" graph and its expected file size. - */ -template -stitch_indices_return_values stitch_label_indices( - path final_index_path_prefix, uint32_t total_number_of_points, label_set all_labels, - tsl::robin_map labels_to_number_of_points, - tsl::robin_map &label_entry_points, - tsl::robin_map> label_id_to_orig_id_map) -{ - size_t final_index_size = 0; - std::vector> stitched_graph(total_number_of_points); - - auto stitching_index_timer = std::chrono::high_resolution_clock::now(); - for (const auto &lbl : all_labels) - { - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - std::vector> curr_label_index; - uint64_t curr_label_index_size; - uint32_t curr_label_entry_point; - - std::tie(curr_label_index, curr_label_index_size) = - diskann::load_label_index(curr_label_index_path, labels_to_number_of_points[lbl]); - curr_label_entry_point = (uint32_t)random(0, curr_label_index.size()); - label_entry_points[lbl] = label_id_to_orig_id_map[lbl][curr_label_entry_point]; - - for (uint32_t node_point = 0; node_point < curr_label_index.size(); node_point++) - { - uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; - for (auto &node_neighbor : curr_label_index[node_point]) - { - uint32_t original_neighbor_id = label_id_to_orig_id_map[lbl][node_neighbor]; - std::vector curr_point_neighbors = stitched_graph[original_point_id]; - if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == - curr_point_neighbors.end()) - { - stitched_graph[original_point_id].push_back(original_neighbor_id); - final_index_size += sizeof(uint32_t); - } - } - } - } - - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); - - std::chrono::duration stitching_index_time = - std::chrono::high_resolution_clock::now() - stitching_index_timer; - std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl; - - return std::make_tuple(stitched_graph, final_index_size); -} - -/* - * Applies the prune_neighbors function from src/index.cpp to - * every node in the stitched graph. - * - * This is an optional step, hence the saving of both the full - * and pruned graph. - */ -template -void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, path input_data_path, - std::vector> stitched_graph, uint32_t stitched_R, - tsl::robin_map label_entry_points, std::string universal_label, - path label_data_path, uint32_t num_threads) -{ - size_t dimension, number_of_label_points; - auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr); - auto std_cout_buffer = std::cout.rdbuf(nullptr); - auto pruning_index_timer = std::chrono::high_resolution_clock::now(); - - diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); - - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, nullptr, nullptr, 0, false, false, - false, false, 0, false); - - // not searching this index, set search_l to 0 - index.load(full_index_path_prefix.c_str(), num_threads, 1); - - std::cout << "parsing labels" << std::endl; - - index.prune_all_neighbors(stitched_R, 750, 1.2); - index.save((final_index_path_prefix).c_str()); - - diskann::cout.rdbuf(diskann_cout_buffer); - std::cout.rdbuf(std_cout_buffer); - std::chrono::duration pruning_index_time = std::chrono::high_resolution_clock::now() - pruning_index_timer; - std::cout << "pruning performed in " << pruning_index_time.count() << " seconds\n" << std::endl; -} - -/* - * Delete all temporary artifacts. - * In the process of creating the stitched index, some temporary artifacts are - * created: - * 1. the separate bin files for each labels' points - * 2. the separate diskANN indices built for each label - * 3. the '.data' file created while generating the indices - */ -void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) -{ - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - path curr_label_index_path_data(curr_label_index_path + ".data"); - - if (std::remove(curr_label_index_path.c_str()) != 0) - throw; - if (std::remove(curr_label_input_data_path.c_str()) != 0) - throw; - if (std::remove(curr_label_index_path_data.c_str()) != 0) - throw; - } -} - -int main(int argc, char **argv) -{ - // 1. handle cmdline inputs - std::string data_type; - path input_data_path, final_index_path_prefix, label_data_path; - std::string universal_label; - uint32_t num_threads, R, L, stitched_R; - float alpha; - - auto index_timer = std::chrono::high_resolution_clock::now(); - handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, - num_threads, R, L, stitched_R, alpha); - - path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt"; - path labels_map_file = final_index_path_prefix + "_labels_map.txt"; - - convert_labels_string_to_int(label_data_path, labels_file_to_use, labels_map_file, universal_label); - - // 2. parse label file and create necessary data structures - std::vector point_ids_to_labels; - tsl::robin_map labels_to_number_of_points; - label_set all_labels; - - std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = - diskann::parse_label_file(labels_file_to_use, universal_label); - - // 3. for each label, make a separate data file - tsl::robin_map> label_id_to_orig_id_map; - uint32_t total_number_of_points = (uint32_t)point_ids_to_labels.size(); - -#ifndef _WINDOWS - if (data_type == "uint8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else - throw; -#else - if (data_type == "uint8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else - throw; -#endif - - // 4. for each created data file, create a vanilla diskANN index - if (data_type == "uint8") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else if (data_type == "int8") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else if (data_type == "float") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else - throw; - - // 5. "stitch" the indices together - std::vector> stitched_graph; - tsl::robin_map label_entry_points; - uint64_t stitched_graph_size; - - if (data_type == "uint8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else if (data_type == "int8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else if (data_type == "float") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else - throw; - path full_index_path_prefix = final_index_path_prefix + "_full"; - // 5a. save the stitched graph to disk - save_full_index(full_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, - universal_label, labels_file_to_use); - - // 6. run a prune on the stitched index, and save to disk - if (data_type == "uint8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else if (data_type == "int8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else if (data_type == "float") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else - throw; - - std::chrono::duration index_time = std::chrono::high_resolution_clock::now() - index_timer; - std::cout << "pruned/stitched graph generated in " << index_time.count() << " seconds" << std::endl; - - clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/python/README.md b/packages/leann-backend-diskann/third_party/DiskANN/apps/python/README.md deleted file mode 100644 index 2b0bc35..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/python/README.md +++ /dev/null @@ -1,46 +0,0 @@ - - -# Integration Tests -The following tests use Python to prepare, run, verify, and tear down the rest api services. - -We do make use of the built-in `unittest` library, but that's only to take advantage of test reporting purposes. - -These are decidedly **not** _unit_ tests. These are end to end integration tests. - -## Caveats -This has only been tested or built for Linux, though we have written platform agnostic Python for the smoke test -(i.e. using `os.path.join`, etc) - -It has been tested on Python 3.9 and 3.10, but should work on Python 3.6+. - -## How to Run - -First, build the DiskANN RestAPI code; see $REPOSITORY_ROOT/workflows/rest_api.md for detailed instructions. - -```bash -cd tests/python -python3 -m venv venv -source venv/bin/activate -pip install -r requirements.txt - -export DISKANN_BUILD_DIR=/path/to/your/diskann/build -python -m unittest -``` - -## Smoke Test Failed, Now What? -The smoke test written takes advantage of temporary directories that are only valid during the -lifetime of the test. The contents of these directories include: -- Randomized vectors (first in tsv, then bin form) used to build the PQFlashIndex -- The PQFlashIndex files - -It is useful to keep these around. By setting some environment variables, you can control whether an ephemeral, -temporary directory is used (and deleted on test completion), or left as an exercise for the developer to -clean up. - -The valid environment variables are: -- `DISKANN_REST_TEST_WORKING_DIR` (example: `$USER/DiskANNRestTest`) - - If this is specified, it **must exist** and **must be writeable**. Any existing files will be clobbered. -- `DISKANN_REST_SERVER` (example: `http://127.0.0.1:10067`) - - Note that if this is set, no data will be generated, nor will a server be started; it is presumed you have done - all the work in creating and starting the rest server prior to running the test and just submits requests against it. diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/python/restapi/__init__.py b/packages/leann-backend-diskann/third_party/DiskANN/apps/python/restapi/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/python/restapi/disk_ann_util.py b/packages/leann-backend-diskann/third_party/DiskANN/apps/python/restapi/disk_ann_util.py deleted file mode 100644 index ec89310..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/python/restapi/disk_ann_util.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import numpy as np -import os -import subprocess - - -def output_vectors( - diskann_build_path: str, - temporary_file_path: str, - vectors: np.ndarray, - timeout: int = 60 -) -> str: - vectors_as_tsv_path = os.path.join(temporary_file_path, "vectors.tsv") - with open(vectors_as_tsv_path, "w") as vectors_tsv_out: - for vector in vectors: - as_str = "\t".join((str(component) for component in vector)) - print(as_str, file=vectors_tsv_out) - # there is probably a clever way to have numpy write out C++ friendly floats, so feel free to remove this in - # favor of something more sane later - vectors_as_bin_path = os.path.join(temporary_file_path, "vectors.bin") - tsv_to_bin_path = os.path.join(diskann_build_path, "apps", "utils", "tsv_to_bin") - - number_of_points, dimensions = vectors.shape - args = [ - tsv_to_bin_path, - "float", - vectors_as_tsv_path, - vectors_as_bin_path, - str(dimensions), - str(number_of_points) - ] - completed = subprocess.run(args, timeout=timeout) - if completed.returncode != 0: - raise Exception(f"Unable to convert tsv to binary using tsv_to_bin, completed_process: {completed}") - return vectors_as_bin_path - - -def build_ssd_index( - diskann_build_path: str, - temporary_file_path: str, - vectors: np.ndarray, - per_process_timeout: int = 60 # this may not be long enough if you're doing something larger -): - vectors_as_bin_path = output_vectors(diskann_build_path, temporary_file_path, vectors, timeout=per_process_timeout) - - ssd_builder_path = os.path.join(diskann_build_path, "apps", "build_disk_index") - args = [ - ssd_builder_path, - "--data_type", "float", - "--dist_fn", "l2", - "--data_path", vectors_as_bin_path, - "--index_path_prefix", os.path.join(temporary_file_path, "smoke_test"), - "-R", "64", - "-L", "100", - "--search_DRAM_budget", "1", - "--build_DRAM_budget", "1", - "--num_threads", "1", - "--PQ_disk_bytes", "0" - ] - completed = subprocess.run(args, timeout=per_process_timeout) - - if completed.returncode != 0: - command_run = " ".join(args) - raise Exception(f"Unable to build a disk index with the command: '{command_run}'\ncompleted_process: {completed}\nstdout: {completed.stdout}\nstderr: {completed.stderr}") - # index is now built inside of temporary_file_path diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/range_search_disk_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/range_search_disk_index.cpp deleted file mode 100644 index 3167572..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/range_search_disk_index.cpp +++ /dev/null @@ -1,379 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include - -#include "index.h" -#include "disk_utils.h" -#include "math_utils.h" -#include "memory_mapper.h" -#include "pq_flash_index.h" -#include "partition.h" -#include "timer.h" -#include "program_options_utils.hpp" - -#ifndef _WINDOWS -#include -#include -#include -#include "linux_aligned_file_reader.h" -#else -#ifdef USE_BING_INFRA -#include "bing_aligned_file_reader.h" -#else -#include "windows_aligned_file_reader.h" -#endif -#endif - -namespace po = boost::program_options; - -#define WARMUP false - -void print_stats(std::string category, std::vector percentiles, std::vector results) -{ - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; -} - -template -int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, const std::string &query_file, - std::string >_file, const uint32_t num_threads, const float search_range, - const uint32_t beamwidth, const uint32_t num_nodes_to_cache, const std::vector &Lvec) -{ - std::string pq_prefix = index_path_prefix + "_pq"; - std::string disk_index_file = index_path_prefix + "_disk.index"; - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::endl; - else - diskann::cout << " beamwidth: " << beamwidth << std::endl; - - // load query bin - T *query = nullptr; - std::vector> groundtruth_ids; - size_t query_num, query_dim, query_aligned_dim, gt_num; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && file_exists(gt_file)) - { - diskann::load_range_truthset(gt_file, groundtruth_ids, - gt_num); // use for range search type of truthset - // diskann::prune_truthset_for_range(gt_file, search_range, - // groundtruth_ids, gt_num); // use for traditional truthset - if (gt_num != query_num) - { - diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; - return -1; - } - calc_recall_flag = true; - } - - std::shared_ptr reader = nullptr; -#ifdef _WINDOWS -#ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); -#else - reader.reset(new diskann::BingAlignedFileReader()); -#endif -#else - reader.reset(new LinuxAlignedFileReader()); -#endif - - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) - { - return res; - } - // cache bfs levels - std::vector node_list; - diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - // _pFlashIndex->generate_cache_list_from_sample_queries( - // warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, - // node_list); - _pFlashIndex->load_cache_list(node_list); - node_list.clear(); - node_list.shrink_to_fit(); - - omp_set_num_threads(num_threads); - - uint64_t warmup_L = 20; - uint64_t warmup_num = 0, warmup_dim = 0, warmup_aligned_dim = 0; - T *warmup = nullptr; - - if (WARMUP) - { - if (file_exists(warmup_query_file)) - { - diskann::load_aligned_bin(warmup_query_file, warmup, warmup_num, warmup_dim, warmup_aligned_dim); - } - else - { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } - } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); - -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)warmup_num; i++) - { - _pFlashIndex->cached_beam_search(warmup + (i * warmup_aligned_dim), 1, warmup_L, - warmup_result_ids_64.data() + (i * 1), - warmup_result_dists.data() + (i * 1), 4); - } - diskann::cout << "..done" << std::endl; - } - - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@rng=" + std::to_string(search_range); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) - << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) - << "CPU (s)"; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall_string << std::endl; - } - else - diskann::cout << std::endl; - diskann::cout << "===============================================================" - "===========================================" - << std::endl; - - std::vector>> query_result_ids(Lvec.size()); - - uint32_t optimized_beamwidth = 2; - uint32_t max_list_size = 10000; - - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint32_t L = Lvec[test_id]; - - if (beamwidth <= 0) - { - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); - } - else - optimized_beamwidth = beamwidth; - - query_result_ids[test_id].clear(); - query_result_ids[test_id].resize(query_num); - - diskann::QueryStats *stats = new diskann::QueryStats[query_num]; - - auto s = std::chrono::high_resolution_clock::now(); -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - std::vector indices; - std::vector distances; - uint32_t res_count = - _pFlashIndex->range_search(query + (i * query_aligned_dim), search_range, L, max_list_size, indices, - distances, optimized_beamwidth, stats + i); - query_result_ids[test_id][i].reserve(res_count); - query_result_ids[test_id][i].resize(res_count); - for (uint32_t idx = 0; idx < res_count; idx++) - query_result_ids[test_id][i][idx] = (uint32_t)indices[idx]; - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - auto qps = (1.0 * query_num) / (1.0 * diff.count()); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - double mean_cpuus = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - double recall = 0; - double ratio_of_sums = 0; - if (calc_recall_flag) - { - recall = - diskann::calculate_range_search_recall((uint32_t)query_num, groundtruth_ids, query_result_ids[test_id]); - - uint32_t total_true_positive = 0; - uint32_t total_positive = 0; - for (uint32_t i = 0; i < query_num; i++) - { - total_true_positive += (uint32_t)query_result_ids[test_id][i].size(); - total_positive += (uint32_t)groundtruth_ids[i].size(); - } - - ratio_of_sums = (1.0 * total_true_positive) / (1.0 * total_positive); - } - - diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps - << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_cpuus; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall << "," << ratio_of_sums << std::endl; - } - else - diskann::cout << std::endl; - } - - diskann::cout << "Done searching. " << std::endl; - - diskann::aligned_free(query); - if (warmup != nullptr) - diskann::aligned_free(warmup); - return 0; -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file; - uint32_t num_threads, W, num_nodes_to_cache; - std::vector Lvec; - float range; - - po::options_description desc{program_options_utils::make_program_description( - "range_search_disk_index", "Searches disk DiskANN indexes using ranges")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("query_file", po::value(&query_file)->required(), - program_options_utils::QUERY_FILE_DESCRIPTION); - required_configs.add_options()("search_list,L", - po::value>(&Lvec)->multitoken()->required(), - program_options_utils::SEARCH_LIST_DESCRIPTION); - required_configs.add_options()("range_threshold,K", po::value(&range)->required(), - "Number of neighbors to be returned"); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - program_options_utils::GROUND_TRUTH_FILE_DESCRIPTION); - optional_configs.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - program_options_utils::NUMBER_OF_NODES_TO_CACHE); - optional_configs.add_options()("beamwidth,W", po::value(&W)->default_value(2), - program_options_utils::BEAMWIDTH); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - - if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) - { - std::cout << "Currently support only floating point data for Inner Product." << std::endl; - return -1; - } - - try - { - if (data_type == std::string("float")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("int8")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("uint8")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/CMakeLists.txt deleted file mode 100644 index c73b427..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/CMakeLists.txt +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -set(CMAKE_CXX_STANDARD 17) - -add_executable(inmem_server inmem_server.cpp) -if(MSVC) - target_link_options(inmem_server PRIVATE /MACHINE:x64) - target_link_libraries(inmem_server debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib Boost::program_options) - target_link_libraries(inmem_server optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib Boost::program_options) -else() - target_link_libraries(inmem_server ${PROJECT_NAME} aio -ltcmalloc -lboost_system -lcrypto -lssl -lcpprest Boost::program_options) -endif() - -add_executable(ssd_server ssd_server.cpp) -if(MSVC) - target_link_options(ssd_server PRIVATE /MACHINE:x64) - target_link_libraries(ssd_server debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib Boost::program_options) - target_link_libraries(ssd_server optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib Boost::program_options) -else() - target_link_libraries(ssd_server ${PROJECT_NAME} aio -ltcmalloc -lboost_system -lcrypto -lssl -lcpprest Boost::program_options) -endif() - -add_executable(multiple_ssdindex_server multiple_ssdindex_server.cpp) -if(MSVC) - target_link_options(multiple_ssdindex_server PRIVATE /MACHINE:x64) - target_link_libraries(multiple_ssdindex_server debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib Boost::program_options) - target_link_libraries(multiple_ssdindex_server optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib Boost::program_options) -else() - target_link_libraries(multiple_ssdindex_server ${PROJECT_NAME} aio -ltcmalloc -lboost_system -lcrypto -lssl -lcpprest Boost::program_options) -endif() - -add_executable(client client.cpp) -if(MSVC) - target_link_options(client PRIVATE /MACHINE:x64) - target_link_libraries(client debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib Boost::program_options) - target_link_libraries(client optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib Boost::program_options) -else() - target_link_libraries(client ${PROJECT_NAME} -lboost_system -lcrypto -lssl -lcpprest Boost::program_options) -endif() \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/client.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/client.cpp deleted file mode 100644 index fdf4414..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/client.cpp +++ /dev/null @@ -1,124 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include - -#include -#include - -using namespace web; -using namespace web::http; -using namespace web::http::client; - -using namespace diskann; -namespace po = boost::program_options; - -template -void query_loop(const std::string &ip_addr_port, const std::string &query_file, const unsigned nq, const unsigned Ls, - const unsigned k_value) -{ - web::http::client::http_client client(U(ip_addr_port)); - - T *data; - size_t npts = 1, ndims = 128, rounded_dim = 128; - diskann::load_aligned_bin(query_file, data, npts, ndims, rounded_dim); - - for (unsigned i = 0; i < nq; ++i) - { - T *vec = data + i * rounded_dim; - web::http::http_request http_query(methods::POST); - web::json::value queryJson = web::json::value::object(); - queryJson[QUERY_ID_KEY] = i; - queryJson[K_KEY] = k_value; - queryJson[L_KEY] = Ls; - for (size_t i = 0; i < ndims; ++i) - { - queryJson[VECTOR_KEY][i] = web::json::value::number(vec[i]); - } - http_query.set_body(queryJson); - - client.request(http_query) - .then([](web::http::http_response response) -> pplx::task { - if (response.status_code() == status_codes::OK) - { - return response.extract_string(); - } - std::cerr << "Query failed" << std::endl; - return pplx::task_from_result(utility::string_t()); - }) - .then([](pplx::task previousTask) { - try - { - std::cout << previousTask.get() << std::endl; - } - catch (http_exception const &e) - { - std::wcout << e.what() << std::endl; - } - }) - .wait(); - } -} - -int main(int argc, char *argv[]) -{ - std::string data_type, query_file, address; - uint32_t num_queries; - uint32_t l_search, k_value; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the queries to search"); - desc.add_options()("num_queries,Q", po::value(&num_queries)->required(), - "Number of queries to search"); - desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); - desc.add_options()("k_value,K", po::value(&k_value)->default_value(10), "Value of K (default 10)"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; - } - - if (data_type == std::string("float")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else if (data_type == std::string("int8")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else if (data_type == std::string("uint8")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else - { - std::cerr << "Unsupported type " << argv[2] << std::endl; - return -1; - } - - return 0; -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/inmem_server.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/inmem_server.cpp deleted file mode 100644 index 11da541..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/inmem_server.cpp +++ /dev/null @@ -1,138 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include - -#include - -using namespace diskann; -namespace po = boost::program_options; - -std::unique_ptr g_httpServer(nullptr); -std::vector> g_inMemorySearch; - -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); - - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - - g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch, typestring)); - std::cout << "Created a server object" << std::endl; - - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; -} - -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); -} - -int main(int argc, char *argv[]) -{ - std::string data_type, index_file, data_file, address, dist_fn, tags_file; - uint32_t num_threads; - uint32_t l_search; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("data_file", po::value(&data_file)->required(), - "File containing the data found in the index"); - desc.add_options()("index_path_prefix", po::value(&index_file)->required(), - "Path prefix for saving index file components"); - desc.add_options()("num_threads,T", po::value(&num_threads)->required(), - "Number of threads used for building index"); - desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; - } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } - - if (data_type == std::string("float")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("int8")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("uint8")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else - { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - } - - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/main.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/main.cpp deleted file mode 100644 index cb48d67..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/main.cpp +++ /dev/null @@ -1,83 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include - -std::unique_ptr g_httpServer(nullptr); -std::unique_ptr g_inMemorySearch(nullptr); - -void setup(const utility::string_t &address) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); - - std::wcout << L"Attempting to start server on " << uri.to_string() << std::endl; - - g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch)); - g_httpServer->open().wait(); - - ucout << U"Listening for requests on: " << address << std::endl; -} - -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); -} - -void loadIndex(const char *indexFile, const char *baseFile, const char *idsFile) -{ - auto nsgSearch = new diskann::InMemorySearch(baseFile, indexFile, idsFile, diskann::L2); - g_inMemorySearch = std::unique_ptr(nsgSearch); -} - -std::wstring getHostingAddress(const char *hostNameAndPort) -{ - wchar_t buffer[4096]; - mbstowcs_s(nullptr, buffer, sizeof(buffer) / sizeof(buffer[0]), hostNameAndPort, - sizeof(buffer) / sizeof(buffer[0])); - return std::wstring(buffer); -} - -int main(int argc, char *argv[]) -{ - if (argc != 5) - { - std::cout << "Usage: nsg_server " - " " - << std::endl; - exit(1); - } - - auto address = getHostingAddress(argv[1]); - loadIndex(argv[2], argv[3], argv[4]); - while (1) - { - try - { - setup(address); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/multiple_ssdindex_server.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/multiple_ssdindex_server.cpp deleted file mode 100644 index 89cb06f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/multiple_ssdindex_server.cpp +++ /dev/null @@ -1,182 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -using namespace diskann; -namespace po = boost::program_options; - -std::unique_ptr g_httpServer(nullptr); -std::vector> g_ssdSearch; - -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); - - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - - g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; - - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; -} - -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); -} - -int main(int argc, char *argv[]) -{ - std::string data_type, index_prefix_paths, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("index_prefix_paths", po::value(&index_prefix_paths)->required(), - "Path prefix for loading index file components"); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } - - std::vector> index_tag_paths; - std::ifstream index_in(index_prefix_paths); - if (!index_in.is_open()) - { - std::cerr << "Could not open " << index_prefix_paths << std::endl; - exit(-1); - } - std::ifstream tags_in(tags_file); - if (!tags_in.is_open()) - { - std::cerr << "Could not open " << tags_file << std::endl; - exit(-1); - } - std::string prefix, tagfile; - while (std::getline(index_in, prefix)) - { - if (std::getline(tags_in, tagfile)) - { - index_tag_paths.push_back(std::make_pair(prefix, tagfile)); - } - else - { - std::cerr << "The number of tags specified does not match the number of " - "indices specified" - << std::endl; - exit(-1); - } - } - index_in.close(); - tags_in.close(); - - if (data_type == std::string("float")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - } - else if (data_type == std::string("int8")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - } - else if (data_type == std::string("uint8")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - } - else - { - std::cerr << "Unsupported data type " << data_type << std::endl; - exit(-1); - } - - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/ssd_server.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/ssd_server.cpp deleted file mode 100644 index d179973..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/restapi/ssd_server.cpp +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -using namespace diskann; -namespace po = boost::program_options; - -std::unique_ptr g_httpServer(nullptr); -std::vector> g_ssdSearch; - -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); - - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - - g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; - - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; -} - -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); -} - -int main(int argc, char *argv[]) -{ - std::string data_type, index_path_prefix, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for loading index file components"); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } - - if (data_type == std::string("float")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("int8")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("uint8")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else - { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - exit(-1); - } - - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/search_disk_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/search_disk_index.cpp deleted file mode 100644 index 6b0793d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/search_disk_index.cpp +++ /dev/null @@ -1,499 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "common_includes.h" -#include - -#include "index.h" -#include "disk_utils.h" -#include "math_utils.h" -#include "memory_mapper.h" -#include "partition.h" -#include "pq_flash_index.h" -#include "timer.h" -#include "percentile_stats.h" -#include "program_options_utils.hpp" - -#ifndef _WINDOWS -#include -#include -#include -#include "linux_aligned_file_reader.h" -#else -#ifdef USE_BING_INFRA -#include "bing_aligned_file_reader.h" -#else -#include "windows_aligned_file_reader.h" -#endif -#endif - -#define WARMUP false - -namespace po = boost::program_options; - -void print_stats(std::string category, std::vector percentiles, std::vector results) -{ - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; -} - -template -int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, - const std::string &result_output_prefix, const std::string &query_file, std::string >_file, - const uint32_t num_threads, const uint32_t recall_at, const uint32_t beamwidth, - const uint32_t num_nodes_to_cache, const uint32_t search_io_limit, - const std::vector &Lvec, const float fail_if_recall_below, - const std::vector &query_filters, const bool use_reorder_data = false) -{ - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::flush; - else - diskann::cout << " beamwidth: " << beamwidth << std::flush; - if (search_io_limit == std::numeric_limits::max()) - diskann::cout << "." << std::endl; - else - diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; - - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - // load query bin - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - bool filtered_search = false; - if (!query_filters.empty()) - { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) - { - std::cout << "Error. Mismatch in number of queries and size of query " - "filters file" - << std::endl; - return -1; // To return -1 or some other error handling? - } - } - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && gt_file != std::string("NULL") && file_exists(gt_file)) - { - diskann::load_truthset(gt_file, gt_ids, gt_dists, gt_num, gt_dim); - if (gt_num != query_num) - { - diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; - } - calc_recall_flag = true; - } - - std::shared_ptr reader = nullptr; -#ifdef _WINDOWS -#ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); -#else - reader.reset(new diskann::BingAlignedFileReader()); -#endif -#else - reader.reset(new LinuxAlignedFileReader()); -#endif - - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) - { - return res; - } - - std::vector node_list; - diskann::cout << "Caching " << num_nodes_to_cache << " nodes around medoid(s)" << std::endl; - _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - // if (num_nodes_to_cache > 0) - // _pFlashIndex->generate_cache_list_from_sample_queries(warmup_query_file, 15, 6, num_nodes_to_cache, - // num_threads, node_list); - _pFlashIndex->load_cache_list(node_list); - node_list.clear(); - node_list.shrink_to_fit(); - - omp_set_num_threads(num_threads); - - uint64_t warmup_L = 20; - uint64_t warmup_num = 0, warmup_dim = 0, warmup_aligned_dim = 0; - T *warmup = nullptr; - - if (WARMUP) - { - if (file_exists(warmup_query_file)) - { - diskann::load_aligned_bin(warmup_query_file, warmup, warmup_num, warmup_dim, warmup_aligned_dim); - } - else - { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } - } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); - -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)warmup_num; i++) - { - _pFlashIndex->cached_beam_search(warmup + (i * warmup_aligned_dim), 1, warmup_L, - warmup_result_ids_64.data() + (i * 1), - warmup_result_dists.data() + (i * 1), 4); - } - diskann::cout << "..done" << std::endl; - } - - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@" + std::to_string(recall_at); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) - << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) - << "Mean IO (us)" << std::setw(16) << "CPU (s)"; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall_string << std::endl; - } - else - diskann::cout << std::endl; - diskann::cout << "==================================================================" - "=================================================================" - << std::endl; - - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); - - uint32_t optimized_beamwidth = 2; - - double best_recall = 0.0; - - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint32_t L = Lvec[test_id]; - - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; - } - - if (beamwidth <= 0) - { - diskann::cout << "Tuning beamwidth.." << std::endl; - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); - } - else - optimized_beamwidth = beamwidth; - - query_result_ids[test_id].resize(recall_at * query_num); - query_result_dists[test_id].resize(recall_at * query_num); - - auto stats = new diskann::QueryStats[query_num]; - - std::vector query_result_ids_64(recall_at * query_num); - auto s = std::chrono::high_resolution_clock::now(); - -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - if (!filtered_search) - { - _pFlashIndex->cached_beam_search(query + (i * query_aligned_dim), recall_at, L, - query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), - optimized_beamwidth, use_reorder_data, stats + i); - } - else - { - LabelT label_for_search; - if (query_filters.size() == 1) - { // one label for all queries - label_for_search = _pFlashIndex->get_converted_label(query_filters[0]); - } - else - { // one label for each query - label_for_search = _pFlashIndex->get_converted_label(query_filters[i]); - } - _pFlashIndex->cached_beam_search( - query + (i * query_aligned_dim), recall_at, L, query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), optimized_beamwidth, true, label_for_search, - use_reorder_data, stats + i); - } - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - double qps = (1.0 * query_num) / (1.0 * diff.count()); - - diskann::convert_types(query_result_ids_64.data(), query_result_ids[test_id].data(), - query_num, recall_at); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - auto mean_cpuus = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - auto mean_io_us = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.io_us; }); - - double recall = 0; - if (calc_recall_flag) - { - recall = diskann::calculate_recall((uint32_t)query_num, gt_ids, gt_dists, (uint32_t)gt_dim, - query_result_ids[test_id].data(), recall_at, recall_at); - best_recall = std::max(recall, best_recall); - } - - diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps - << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_io_us << std::setw(16) << mean_cpuus; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall << std::endl; - } - else - diskann::cout << std::endl; - delete[] stats; - } - - diskann::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) - { - if (L < recall_at) - continue; - - std::string cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); - - cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_dists_float.bin"; - diskann::save_bin(cur_result_path, query_result_dists[test_id++].data(), query_num, recall_at); - } - - diskann::aligned_free(query); - if (warmup != nullptr) - diskann::aligned_free(warmup); - return best_recall >= fail_if_recall_below ? 0 : -1; -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file, filter_label, - label_type, query_filters_file; - uint32_t num_threads, K, W, num_nodes_to_cache, search_io_limit; - std::vector Lvec; - bool use_reorder_data = false; - float fail_if_recall_below = 0.0f; - - po::options_description desc{ - program_options_utils::make_program_description("search_disk_index", "Searches on-disk DiskANN indexes")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("result_path", po::value(&result_path_prefix)->required(), - program_options_utils::RESULT_PATH_DESCRIPTION); - required_configs.add_options()("query_file", po::value(&query_file)->required(), - program_options_utils::QUERY_FILE_DESCRIPTION); - required_configs.add_options()("recall_at,K", po::value(&K)->required(), - program_options_utils::NUMBER_OF_RESULTS_DESCRIPTION); - required_configs.add_options()("search_list,L", - po::value>(&Lvec)->multitoken()->required(), - program_options_utils::SEARCH_LIST_DESCRIPTION); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - program_options_utils::GROUND_TRUTH_FILE_DESCRIPTION); - optional_configs.add_options()("beamwidth,W", po::value(&W)->default_value(2), - program_options_utils::BEAMWIDTH); - optional_configs.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - program_options_utils::NUMBER_OF_NODES_TO_CACHE); - optional_configs.add_options()( - "search_io_limit", - po::value(&search_io_limit)->default_value(std::numeric_limits::max()), - "Max #IOs for search. Default value: uint32::max()"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("use_reorder_data", po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD. Default value: false"); - optional_configs.add_options()("filter_label", - po::value(&filter_label)->default_value(std::string("")), - program_options_utils::FILTER_LABEL_DESCRIPTION); - optional_configs.add_options()("query_filters_file", - po::value(&query_filters_file)->default_value(std::string("")), - program_options_utils::FILTERS_FILE_DESCRIPTION); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - program_options_utils::LABEL_TYPE_DESCRIPTION); - optional_configs.add_options()("fail_if_recall_below", - po::value(&fail_if_recall_below)->default_value(0.0f), - program_options_utils::FAIL_IF_RECALL_BELOW); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (vm["use_reorder_data"].as()) - use_reorder_data = true; - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - - if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) - { - std::cout << "Currently support only floating point data for Inner Product." << std::endl; - return -1; - } - - if (use_reorder_data && data_type != std::string("float")) - { - std::cout << "Error: Reorder data for reordering currently only " - "supported for float data type." - << std::endl; - return -1; - } - - if (filter_label != "" && query_filters_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; - } - - std::vector query_filters; - if (filter_label != "") - { - query_filters.push_back(filter_label); - } - else if (query_filters_file != "") - { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } - - try - { - if (!query_filters.empty() && label_type == "ushort") - { - if (data_type == std::string("float")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("float")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/search_memory_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/search_memory_index.cpp deleted file mode 100644 index 1a9acc2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/search_memory_index.cpp +++ /dev/null @@ -1,477 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifndef _WINDOWS -#include -#include -#include -#include -#endif - -#include "index.h" -#include "memory_mapper.h" -#include "utils.h" -#include "program_options_utils.hpp" -#include "index_factory.h" - -namespace po = boost::program_options; - -template -int search_memory_index(diskann::Metric &metric, const std::string &index_path, const std::string &result_path_prefix, - const std::string &query_file, const std::string &truthset_file, const uint32_t num_threads, - const uint32_t recall_at, const bool print_all_recalls, const std::vector &Lvec, - const bool dynamic, const bool tags, const bool show_qps_per_thread, - const std::vector &query_filters, const float fail_if_recall_below) -{ - using TagT = uint32_t; - // Load the query file - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - bool calc_recall_flag = false; - if (truthset_file != std::string("null") && file_exists(truthset_file)) - { - diskann::load_truthset(truthset_file, gt_ids, gt_dists, gt_num, gt_dim); - if (gt_num != query_num) - { - std::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; - } - calc_recall_flag = true; - } - else - { - diskann::cout << " Truthset file " << truthset_file << " not found. Not computing recall." << std::endl; - } - - bool filtered_search = false; - if (!query_filters.empty()) - { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) - { - std::cout << "Error. Mismatch in number of queries and size of query " - "filters file" - << std::endl; - return -1; // To return -1 or some other error handling? - } - } - - const size_t num_frozen_pts = diskann::get_graph_num_frozen_points(index_path); - - auto config = diskann::IndexConfigBuilder() - .with_metric(metric) - .with_dimension(query_dim) - .with_max_points(0) - .with_data_load_store_strategy(diskann::DataStoreStrategy::MEMORY) - .with_graph_load_store_strategy(diskann::GraphStoreStrategy::MEMORY) - .with_data_type(diskann_type_to_name()) - .with_label_type(diskann_type_to_name()) - .with_tag_type(diskann_type_to_name()) - .is_dynamic_index(dynamic) - .is_enable_tags(tags) - .is_concurrent_consolidate(false) - .is_pq_dist_build(false) - .is_use_opq(false) - .with_num_pq_chunks(0) - .with_num_frozen_pts(num_frozen_pts) - .build(); - - auto index_factory = diskann::IndexFactory(config); - auto index = index_factory.create_instance(); - index->load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end()))); - std::cout << "Index loaded" << std::endl; - - if (metric == diskann::FAST_L2) - index->optimize_index_layout(); - - std::cout << "Using " << num_threads << " threads to search" << std::endl; - std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - std::cout.precision(2); - const std::string qps_title = show_qps_per_thread ? "QPS/thread" : "QPS"; - uint32_t table_width = 0; - if (tags) - { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(20) << "Mean Latency (mus)" - << std::setw(15) << "99.9 Latency"; - table_width += 4 + 12 + 20 + 15; - } - else - { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(18) << "Avg dist cmps" - << std::setw(20) << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency"; - table_width += 4 + 12 + 18 + 20 + 15; - } - uint32_t recalls_to_print = 0; - const uint32_t first_recall = print_all_recalls ? 1 : recall_at; - if (calc_recall_flag) - { - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) - { - std::cout << std::setw(12) << ("Recall@" + std::to_string(curr_recall)); - } - recalls_to_print = recall_at + 1 - first_recall; - table_width += recalls_to_print * 12; - } - std::cout << std::endl; - std::cout << std::string(table_width, '=') << std::endl; - - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); - std::vector latency_stats(query_num, 0); - std::vector cmp_stats; - if (not tags || filtered_search) - { - cmp_stats = std::vector(query_num, 0); - } - - std::vector query_result_tags; - if (tags) - { - query_result_tags.resize(recall_at * query_num); - } - - double best_recall = 0.0; - - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint32_t L = Lvec[test_id]; - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; - } - - query_result_ids[test_id].resize(recall_at * query_num); - query_result_dists[test_id].resize(recall_at * query_num); - std::vector res = std::vector(); - - auto s = std::chrono::high_resolution_clock::now(); - omp_set_num_threads(num_threads); -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - auto qs = std::chrono::high_resolution_clock::now(); - if (filtered_search && !tags) - { - std::string raw_filter = query_filters.size() == 1 ? query_filters[0] : query_filters[i]; - - auto retval = index->search_with_filters(query + i * query_aligned_dim, raw_filter, recall_at, L, - query_result_ids[test_id].data() + i * recall_at, - query_result_dists[test_id].data() + i * recall_at); - cmp_stats[i] = retval.second; - } - else if (metric == diskann::FAST_L2) - { - index->search_with_optimized_layout(query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at); - } - else if (tags) - { - if (!filtered_search) - { - index->search_with_tags(query + i * query_aligned_dim, recall_at, L, - query_result_tags.data() + i * recall_at, nullptr, res); - } - else - { - std::string raw_filter = query_filters.size() == 1 ? query_filters[0] : query_filters[i]; - - index->search_with_tags(query + i * query_aligned_dim, recall_at, L, - query_result_tags.data() + i * recall_at, nullptr, res, true, raw_filter); - } - - for (int64_t r = 0; r < (int64_t)recall_at; r++) - { - query_result_ids[test_id][recall_at * i + r] = query_result_tags[recall_at * i + r]; - } - } - else - { - cmp_stats[i] = index - ->search(query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at) - .second; - } - auto qe = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = qe - qs; - latency_stats[i] = (float)(diff.count() * 1000000); - } - std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; - - double displayed_qps = query_num / diff.count(); - - if (show_qps_per_thread) - displayed_qps /= num_threads; - - std::vector recalls; - if (calc_recall_flag) - { - recalls.reserve(recalls_to_print); - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) - { - recalls.push_back(diskann::calculate_recall((uint32_t)query_num, gt_ids, gt_dists, (uint32_t)gt_dim, - query_result_ids[test_id].data(), recall_at, curr_recall)); - } - } - - std::sort(latency_stats.begin(), latency_stats.end()); - double mean_latency = - std::accumulate(latency_stats.begin(), latency_stats.end(), 0.0) / static_cast(query_num); - - float avg_cmps = (float)std::accumulate(cmp_stats.begin(), cmp_stats.end(), 0) / (float)query_num; - - if (tags && !filtered_search) - { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(20) << (float)mean_latency - << std::setw(15) << (float)latency_stats[(uint64_t)(0.999 * query_num)]; - } - else - { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(18) << avg_cmps - << std::setw(20) << (float)mean_latency << std::setw(15) - << (float)latency_stats[(uint64_t)(0.999 * query_num)]; - } - for (double recall : recalls) - { - std::cout << std::setw(12) << recall; - best_recall = std::max(recall, best_recall); - } - std::cout << std::endl; - } - - std::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) - { - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; - } - std::string cur_result_path_prefix = result_path_prefix + "_" + std::to_string(L); - - std::string cur_result_path = cur_result_path_prefix + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); - - cur_result_path = cur_result_path_prefix + "_dists_float.bin"; - diskann::save_bin(cur_result_path, query_result_dists[test_id].data(), query_num, recall_at); - - test_id++; - } - - diskann::aligned_free(query); - return best_recall >= fail_if_recall_below ? 0 : -1; -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, index_path_prefix, result_path, query_file, gt_file, filter_label, label_type, - query_filters_file; - uint32_t num_threads, K; - std::vector Lvec; - bool print_all_recalls, dynamic, tags, show_qps_per_thread; - float fail_if_recall_below = 0.0f; - - po::options_description desc{ - program_options_utils::make_program_description("search_memory_index", "Searches in-memory DiskANN indexes")}; - try - { - desc.add_options()("help,h", "Print this information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("result_path", po::value(&result_path)->required(), - program_options_utils::RESULT_PATH_DESCRIPTION); - required_configs.add_options()("query_file", po::value(&query_file)->required(), - program_options_utils::QUERY_FILE_DESCRIPTION); - required_configs.add_options()("recall_at,K", po::value(&K)->required(), - program_options_utils::NUMBER_OF_RESULTS_DESCRIPTION); - required_configs.add_options()("search_list,L", - po::value>(&Lvec)->multitoken()->required(), - program_options_utils::SEARCH_LIST_DESCRIPTION); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("filter_label", - po::value(&filter_label)->default_value(std::string("")), - program_options_utils::FILTER_LABEL_DESCRIPTION); - optional_configs.add_options()("query_filters_file", - po::value(&query_filters_file)->default_value(std::string("")), - program_options_utils::FILTERS_FILE_DESCRIPTION); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - program_options_utils::LABEL_TYPE_DESCRIPTION); - optional_configs.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - program_options_utils::GROUND_TRUTH_FILE_DESCRIPTION); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()( - "dynamic", po::value(&dynamic)->default_value(false), - "Whether the index is dynamic. Dynamic indices must have associated tags. Default false."); - optional_configs.add_options()("tags", po::value(&tags)->default_value(false), - "Whether to search with external identifiers (tags). Default false."); - optional_configs.add_options()("fail_if_recall_below", - po::value(&fail_if_recall_below)->default_value(0.0f), - program_options_utils::FAIL_IF_RECALL_BELOW); - - // Output controls - po::options_description output_controls("Output controls"); - output_controls.add_options()("print_all_recalls", po::bool_switch(&print_all_recalls), - "Print recalls at all positions, from 1 up to specified " - "recall_at value"); - output_controls.add_options()("print_qps_per_thread", po::bool_switch(&show_qps_per_thread), - "Print overall QPS divided by the number of threads in " - "the output table"); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs).add(output_controls); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if ((dist_fn == std::string("mips")) && (data_type == std::string("float"))) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else if ((dist_fn == std::string("fast_l2")) && (data_type == std::string("float"))) - { - metric = diskann::Metric::FAST_L2; - } - else - { - std::cout << "Unsupported distance function. Currently only l2/ cosine are " - "supported in general, and mips/fast_l2 only for floating " - "point data." - << std::endl; - return -1; - } - - if (dynamic && not tags) - { - std::cerr << "Tags must be enabled while searching dynamically built indices" << std::endl; - return -1; - } - - if (fail_if_recall_below < 0.0 || fail_if_recall_below >= 100.0) - { - std::cerr << "fail_if_recall_below parameter must be between 0 and 100%" << std::endl; - return -1; - } - - if (filter_label != "" && query_filters_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; - } - - std::vector query_filters; - if (filter_label != "") - { - query_filters.push_back(filter_label); - } - else if (query_filters_file != "") - { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } - - try - { - if (!query_filters.empty() && label_type == "ushort") - { - if (data_type == std::string("int8")) - { - return search_memory_index( - metric, index_path_prefix, result_path, query_file, gt_file, num_threads, K, print_all_recalls, - Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); - } - else if (data_type == std::string("uint8")) - { - return search_memory_index( - metric, index_path_prefix, result_path, query_file, gt_file, num_threads, K, print_all_recalls, - Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); - } - else if (data_type == std::string("float")) - { - return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else - { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("int8")) - { - return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else if (data_type == std::string("uint8")) - { - return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else if (data_type == std::string("float")) - { - return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else - { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - return -1; - } - } - } - catch (std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/test_insert_deletes_consolidate.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/test_insert_deletes_consolidate.cpp deleted file mode 100644 index 97aed18..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/test_insert_deletes_consolidate.cpp +++ /dev/null @@ -1,536 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "utils.h" -#include "filter_utils.h" -#include "program_options_utils.hpp" -#include "index_factory.h" - -#ifndef _WINDOWS -#include -#include -#include -#endif - -#include "memory_mapper.h" - -namespace po = boost::program_options; - -// load_aligned_bin modified to read pieces of the file, but using ifstream -// instead of cached_ifstream. -template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) -{ - diskann::Timer timer; - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)dim_i32; - - size_t expected_actual_file_size = npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != expected_actual_file_size) - { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << actual_file_size << " while expected size is " - << expected_actual_file_size << " npts = " << npts << " dim = " << dim << " size of = " << sizeof(T) - << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (offset_points + points_to_read > npts) - { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read - << " points, but have only " << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); - - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Read " << points_to_read << " points using non-cached reads in " << elapsedSeconds << std::endl; -} - -std::string get_save_filename(const std::string &save_path, size_t points_to_skip, size_t points_deleted, - size_t last_point_threshold) -{ - std::string final_path = save_path; - if (points_to_skip > 0) - { - final_path += "skip" + std::to_string(points_to_skip) + "-"; - } - - final_path += "del" + std::to_string(points_deleted) + "-"; - final_path += std::to_string(last_point_threshold); - return final_path; -} - -template -void insert_till_next_checkpoint(diskann::AbstractIndex &index, size_t start, size_t end, int32_t thread_count, T *data, - size_t aligned_dim, std::vector> &location_to_labels) -{ - diskann::Timer insert_timer; -#pragma omp parallel for num_threads(thread_count) schedule(dynamic) - for (int64_t j = start; j < (int64_t)end; j++) - { - if (!location_to_labels.empty()) - { - index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j), - location_to_labels[j - start]); - } - else - { - index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)); - } - } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds - << " points/second overall, " << (end - start) / elapsedSeconds / thread_count << " per thread)\n "; -} - -template -void delete_from_beginning(diskann::AbstractIndex &index, diskann::IndexWriteParameters &delete_params, - size_t points_to_skip, size_t points_to_delete_from_beginning) -{ - try - { - std::cout << std::endl - << "Lazy deleting points " << points_to_skip << " to " - << points_to_skip + points_to_delete_from_beginning << "... "; - for (size_t i = points_to_skip; i < points_to_skip + points_to_delete_from_beginning; ++i) - index.lazy_delete(static_cast(i + 1)); // Since tags are data location + 1 - std::cout << "done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "rate: (" << points_to_delete_from_beginning / report._time << " points/second overall, " - << points_to_delete_from_beginning / report._time / delete_params.num_threads << " per thread)" - << std::endl; - } - catch (std::system_error &e) - { - std::cout << "Exception caught in deletion thread: " << e.what() << std::endl; - } -} - -template -void build_incremental_index(const std::string &data_path, diskann::IndexWriteParameters ¶ms, size_t points_to_skip, - size_t max_points_to_insert, size_t beginning_index_size, float start_point_norm, - uint32_t num_start_pts, size_t points_per_checkpoint, size_t checkpoints_per_snapshot, - const std::string &save_path, size_t points_to_delete_from_beginning, - size_t start_deletes_after, bool concurrent, const std::string &label_file, - const std::string &universal_label) -{ - size_t dim, aligned_dim; - size_t num_points; - diskann::get_bin_metadata(data_path, num_points, dim); - aligned_dim = ROUND_UP(dim, 8); - bool has_labels = label_file != ""; - using TagT = uint32_t; - using LabelT = uint32_t; - - size_t current_point_offset = points_to_skip; - const size_t last_point_threshold = points_to_skip + max_points_to_insert; - - bool enable_tags = true; - using TagT = uint32_t; - auto index_search_params = diskann::IndexSearchParams(params.search_list_size, params.num_threads); - diskann::IndexConfig index_config = diskann::IndexConfigBuilder() - .with_metric(diskann::L2) - .with_dimension(dim) - .with_max_points(max_points_to_insert) - .is_dynamic_index(true) - .with_index_write_params(params) - .with_index_search_params(index_search_params) - .with_data_type(diskann_type_to_name()) - .with_tag_type(diskann_type_to_name()) - .with_label_type(diskann_type_to_name()) - .with_data_load_store_strategy(diskann::DataStoreStrategy::MEMORY) - .with_graph_load_store_strategy(diskann::GraphStoreStrategy::MEMORY) - .is_enable_tags(enable_tags) - .is_filtered(has_labels) - .with_num_frozen_pts(num_start_pts) - .is_concurrent_consolidate(concurrent) - .build(); - - diskann::IndexFactory index_factory = diskann::IndexFactory(index_config); - auto index = index_factory.create_instance(); - - if (universal_label != "") - { - LabelT u_label = 0; - index->set_universal_label(u_label); - } - - if (points_to_skip > num_points) - { - throw diskann::ANNException("Asked to skip more points than in data file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (max_points_to_insert == 0) - { - max_points_to_insert = num_points; - } - - if (points_to_skip + max_points_to_insert > num_points) - { - max_points_to_insert = num_points - points_to_skip; - std::cerr << "WARNING: Reducing max_points_to_insert to " << max_points_to_insert - << " points since the data file has only that many" << std::endl; - } - - if (beginning_index_size > max_points_to_insert) - { - beginning_index_size = max_points_to_insert; - std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size - << " points since the data file has only that many" << std::endl; - } - if (checkpoints_per_snapshot > 0 && beginning_index_size > points_per_checkpoint) - { - beginning_index_size = points_per_checkpoint; - std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size << std::endl; - } - - T *data = nullptr; - diskann::alloc_aligned( - (void **)&data, std::max(points_per_checkpoint, beginning_index_size) * aligned_dim * sizeof(T), 8 * sizeof(T)); - - std::vector tags(beginning_index_size); - std::iota(tags.begin(), tags.end(), 1 + static_cast(current_point_offset)); - - load_aligned_bin_part(data_path, data, current_point_offset, beginning_index_size); - std::cout << "load aligned bin succeeded" << std::endl; - diskann::Timer timer; - - if (beginning_index_size > 0) - { - index->build(data, beginning_index_size, tags); - } - else - { - index->set_start_points_at_random(static_cast(start_point_norm)); - } - - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Initial non-incremental index build time for " << beginning_index_size << " points took " - << elapsedSeconds << " seconds (" << beginning_index_size / elapsedSeconds << " points/second)\n "; - - current_point_offset += beginning_index_size; - - if (points_to_delete_from_beginning > max_points_to_insert) - { - points_to_delete_from_beginning = static_cast(max_points_to_insert); - std::cerr << "WARNING: Reducing points to delete from beginning to " << points_to_delete_from_beginning - << " points since the data file has only that many" << std::endl; - } - - std::vector> location_to_labels; - if (concurrent) - { - // handle labels - const auto save_path_inc = get_save_filename(save_path + ".after-concurrent-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - std::string labels_file_to_use = save_path_inc + "_label_formatted.txt"; - std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt"; - if (has_labels) - { - convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label); - auto parse_result = diskann::parse_formatted_label_file(labels_file_to_use); - location_to_labels = std::get<0>(parse_result); - } - - int32_t sub_threads = (params.num_threads + 1) / 2; - bool delete_launched = false; - std::future delete_task; - - diskann::Timer timer; - - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, current_point_offset += points_per_checkpoint) - { - const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(*index, start, end, sub_threads, data, aligned_dim, - location_to_labels); - }); - insert_task.wait(); - - if (!delete_launched && end >= start_deletes_after && - end >= points_to_skip + points_to_delete_from_beginning) - { - delete_launched = true; - diskann::IndexWriteParameters delete_params = - diskann::IndexWriteParametersBuilder(params).with_num_threads(sub_threads).build(); - - delete_task = std::async(std::launch::async, [&]() { - delete_from_beginning(*index, delete_params, points_to_skip, - points_to_delete_from_beginning); - }); - } - } - delete_task.wait(); - - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - index->save(save_path_inc.c_str(), true); - } - else - { - const auto save_path_inc = get_save_filename(save_path + ".after-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - std::string labels_file_to_use = save_path_inc + "_label_formatted.txt"; - std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt"; - if (has_labels) - { - convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label); - auto parse_result = diskann::parse_formatted_label_file(labels_file_to_use); - location_to_labels = std::get<0>(parse_result); - } - - size_t last_snapshot_points_threshold = 0; - size_t num_checkpoints_till_snapshot = checkpoints_per_snapshot; - - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, current_point_offset += points_per_checkpoint) - { - const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(*index, start, end, (int32_t)params.num_threads, data, - aligned_dim, location_to_labels); - - if (checkpoints_per_snapshot > 0 && --num_checkpoints_till_snapshot == 0) - { - diskann::Timer save_timer; - - const auto save_path_inc = - get_save_filename(save_path + ".inc-", points_to_skip, points_to_delete_from_beginning, end); - index->save(save_path_inc.c_str(), false); - const double elapsedSeconds = save_timer.elapsed() / 1000000.0; - const size_t points_saved = end - points_to_skip; - - std::cout << "Saved " << points_saved << " points in " << elapsedSeconds << " seconds (" - << points_saved / elapsedSeconds << " points/second)\n"; - - num_checkpoints_till_snapshot = checkpoints_per_snapshot; - last_snapshot_points_threshold = end; - } - - std::cout << "Number of points in the index post insertion " << end << std::endl; - } - - if (checkpoints_per_snapshot > 0 && last_snapshot_points_threshold != last_point_threshold) - { - const auto save_path_inc = get_save_filename(save_path + ".inc-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - // index.save(save_path_inc.c_str(), false); - } - - if (points_to_delete_from_beginning > 0) - { - delete_from_beginning(*index, params, points_to_skip, points_to_delete_from_beginning); - } - - index->save(save_path_inc.c_str(), true); - } - - diskann::aligned_free(data); -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix; - uint32_t num_threads, R, L, num_start_pts; - float alpha, start_point_norm; - size_t points_to_skip, max_points_to_insert, beginning_index_size, points_per_checkpoint, checkpoints_per_snapshot, - points_to_delete_from_beginning, start_deletes_after; - bool concurrent; - - // label options - std::string label_file, label_type, universal_label; - std::uint32_t Lf, unique_labels_supported; - - po::options_description desc{program_options_utils::make_program_description("test_insert_deletes_consolidate", - "Test insert deletes & consolidate")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("data_path", po::value(&data_path)->required(), - program_options_utils::INPUT_DATA_PATH); - required_configs.add_options()("points_to_skip", po::value(&points_to_skip)->required(), - "Skip these first set of points from file"); - required_configs.add_options()("beginning_index_size", po::value(&beginning_index_size)->required(), - "Batch build will be called on these set of points"); - required_configs.add_options()("points_per_checkpoint", po::value(&points_per_checkpoint)->required(), - "Insertions are done in batches of points_per_checkpoint"); - required_configs.add_options()("checkpoints_per_snapshot", - po::value(&checkpoints_per_snapshot)->required(), - "Save the index to disk every few checkpoints"); - required_configs.add_options()("points_to_delete_from_beginning", - po::value(&points_to_delete_from_beginning)->required(), ""); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - program_options_utils::NUMBER_THREADS_DESCRIPTION); - optional_configs.add_options()("max_degree,R", po::value(&R)->default_value(64), - program_options_utils::MAX_BUILD_DEGREE); - optional_configs.add_options()("Lbuild,L", po::value(&L)->default_value(100), - program_options_utils::GRAPH_BUILD_COMPLEXITY); - optional_configs.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - program_options_utils::GRAPH_BUILD_ALPHA); - optional_configs.add_options()("max_points_to_insert", - po::value(&max_points_to_insert)->default_value(0), - "These number of points from the file are inserted after " - "points_to_skip"); - optional_configs.add_options()("do_concurrent", po::value(&concurrent)->default_value(false), ""); - optional_configs.add_options()("start_deletes_after", - po::value(&start_deletes_after)->default_value(0), ""); - optional_configs.add_options()("start_point_norm", po::value(&start_point_norm)->default_value(0), - "Set the start point to a random point on a sphere of this radius"); - - // optional params for filters - optional_configs.add_options()("label_file", po::value(&label_file)->default_value(""), - "Input label file in txt format for Filtered Index search. " - "The file should contain comma separated filters for each node " - "with each line corresponding to a graph node"); - optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with labels_file"); - optional_configs.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - optional_configs.add_options()("unique_labels_supported", - po::value(&unique_labels_supported)->default_value(0), - "Number of unique labels supported by the dynamic index."); - - optional_configs.add_options()( - "num_start_points", - po::value(&num_start_pts)->default_value(diskann::defaults::NUM_FROZEN_POINTS_DYNAMIC), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (beginning_index_size == 0) - if (start_point_norm == 0) - { - std::cout << "When beginning_index_size is 0, use a start " - "point with " - "appropriate norm" - << std::endl; - return -1; - } - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - bool has_labels = false; - if (!label_file.empty() || label_file != "") - { - has_labels = true; - } - - if (num_start_pts < unique_labels_supported) - { - num_start_pts = unique_labels_supported; - } - - try - { - diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(500) - .with_alpha(alpha) - .with_num_threads(num_threads) - .with_filter_list_size(Lf) - .build(); - - if (data_type == std::string("int8")) - build_incremental_index( - data_path, params, points_to_skip, max_points_to_insert, beginning_index_size, start_point_norm, - num_start_pts, points_per_checkpoint, checkpoints_per_snapshot, index_path_prefix, - points_to_delete_from_beginning, start_deletes_after, concurrent, label_file, universal_label); - else if (data_type == std::string("uint8")) - build_incremental_index( - data_path, params, points_to_skip, max_points_to_insert, beginning_index_size, start_point_norm, - num_start_pts, points_per_checkpoint, checkpoints_per_snapshot, index_path_prefix, - points_to_delete_from_beginning, start_deletes_after, concurrent, label_file, universal_label); - else if (data_type == std::string("float")) - build_incremental_index(data_path, params, points_to_skip, max_points_to_insert, - beginning_index_size, start_point_norm, num_start_pts, points_per_checkpoint, - checkpoints_per_snapshot, index_path_prefix, points_to_delete_from_beginning, - start_deletes_after, concurrent, label_file, universal_label); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - } - catch (const std::exception &e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } - catch (...) - { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); - } - - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/test_streaming_scenario.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/test_streaming_scenario.cpp deleted file mode 100644 index 5a43a69..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/test_streaming_scenario.cpp +++ /dev/null @@ -1,523 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "utils.h" -#include "filter_utils.h" -#include "program_options_utils.hpp" - -#ifndef _WINDOWS -#include -#include -#include -#endif - -#include "memory_mapper.h" - -namespace po = boost::program_options; - -// load_aligned_bin modified to read pieces of the file, but using ifstream -// instead of cached_ifstream. -template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)dim_i32; - - size_t expected_actual_file_size = npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != expected_actual_file_size) - { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << actual_file_size << " while expected size is " - << expected_actual_file_size << " npts = " << npts << " dim = " << dim << " size of = " << sizeof(T) - << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (offset_points + points_to_read > npts) - { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read - << " points, but have only " << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); -} - -std::string get_save_filename(const std::string &save_path, size_t active_window, size_t consolidate_interval, - size_t max_points_to_insert) -{ - std::string final_path = save_path; - final_path += "act" + std::to_string(active_window) + "-"; - final_path += "cons" + std::to_string(consolidate_interval) + "-"; - final_path += "max" + std::to_string(max_points_to_insert); - return final_path; -} - -template -void insert_next_batch(diskann::AbstractIndex &index, size_t start, size_t end, size_t insert_threads, T *data, - size_t aligned_dim, std::vector> &pts_to_labels) -{ - try - { - diskann::Timer insert_timer; - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - size_t num_failed = 0; -#pragma omp parallel for num_threads((int32_t)insert_threads) schedule(dynamic) reduction(+ : num_failed) - for (int64_t j = start; j < (int64_t)end; j++) - { - int insert_result = -1; - if (pts_to_labels.size() > 0) - { - insert_result = index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j), - pts_to_labels[j - start]); - } - else - { - insert_result = index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)); - } - - if (insert_result != 0) - { - std::cerr << "Insert failed " << j << std::endl; - num_failed++; - } - } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds - << " points/second overall, " << (end - start) / elapsedSeconds / insert_threads << " per thread)" - << std::endl; - if (num_failed > 0) - std::cout << num_failed << " of " << end - start << "inserts failed" << std::endl; - } - catch (std::system_error &e) - { - std::cout << "Exiting after catching exception in insertion task: " << e.what() << std::endl; - exit(-1); - } -} - -template -void delete_and_consolidate(diskann::AbstractIndex &index, diskann::IndexWriteParameters &delete_params, size_t start, - size_t end) -{ - try - { - std::cout << std::endl << "Lazy deleting points " << start << " to " << end << "... "; - for (size_t i = start; i < end; ++i) - index.lazy_delete(static_cast(1 + i)); - std::cout << "lazy delete done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - while (report._status != diskann::consolidation_report::status_code::SUCCESS) - { - int wait_time = 5; - if (report._status == diskann::consolidation_report::status_code::LOCK_FAIL) - { - diskann::cerr << "Unable to acquire consolidate delete lock after " - << "deleting points " << start << " to " << end << ". Will retry in " << wait_time - << "seconds." << std::endl; - } - else if (report._status == diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR) - { - diskann::cerr << "Inconsistent counts in data structure. " - << "Will retry in " << wait_time << "seconds." << std::endl; - } - else - { - std::cerr << "Exiting after unknown error in consolidate delete" << std::endl; - exit(-1); - } - std::this_thread::sleep_for(std::chrono::seconds(wait_time)); - report = index.consolidate_deletes(delete_params); - } - auto points_processed = report._active_points + report._slots_released; - auto deletion_rate = points_processed / report._time; - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "Deletion rate: " << deletion_rate << "/sec " - << "Deletion rate: " << deletion_rate / delete_params.num_threads << "/thread/sec " << std::endl; - } - catch (std::system_error &e) - { - std::cerr << "Exiting after catching exception in deletion task: " << e.what() << std::endl; - exit(-1); - } -} - -template -void build_incremental_index(const std::string &data_path, const uint32_t L, const uint32_t R, const float alpha, - const uint32_t insert_threads, const uint32_t consolidate_threads, - size_t max_points_to_insert, size_t active_window, size_t consolidate_interval, - const float start_point_norm, uint32_t num_start_pts, const std::string &save_path, - const std::string &label_file, const std::string &universal_label, const uint32_t Lf) -{ - const uint32_t C = 500; - const bool saturate_graph = false; - bool has_labels = label_file != ""; - - diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_threads(insert_threads) - .with_filter_list_size(Lf) - .build(); - - auto index_search_params = diskann::IndexSearchParams(L, insert_threads); - diskann::IndexWriteParameters delete_params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_threads(consolidate_threads) - .with_filter_list_size(Lf) - .build(); - - size_t dim, aligned_dim; - size_t num_points; - - std::vector> pts_to_labels; - - const auto save_path_inc = - get_save_filename(save_path + ".after-streaming-", active_window, consolidate_interval, max_points_to_insert); - std::string labels_file_to_use = save_path_inc + "_label_formatted.txt"; - std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt"; - if (has_labels) - { - convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label); - auto parse_result = diskann::parse_formatted_label_file(labels_file_to_use); - pts_to_labels = std::get<0>(parse_result); - } - - diskann::get_bin_metadata(data_path, num_points, dim); - diskann::cout << "metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" - << std::endl; - aligned_dim = ROUND_UP(dim, 8); - auto index_config = diskann::IndexConfigBuilder() - .with_metric(diskann::L2) - .with_dimension(dim) - .with_max_points(active_window + 4 * consolidate_interval) - .is_dynamic_index(true) - .is_enable_tags(true) - .is_use_opq(false) - .is_filtered(has_labels) - .with_num_pq_chunks(0) - .is_pq_dist_build(false) - .with_num_frozen_pts(num_start_pts) - .with_tag_type(diskann_type_to_name()) - .with_label_type(diskann_type_to_name()) - .with_data_type(diskann_type_to_name()) - .with_index_write_params(params) - .with_index_search_params(index_search_params) - .with_data_load_store_strategy(diskann::DataStoreStrategy::MEMORY) - .with_graph_load_store_strategy(diskann::GraphStoreStrategy::MEMORY) - .build(); - - diskann::IndexFactory index_factory = diskann::IndexFactory(index_config); - auto index = index_factory.create_instance(); - - if (universal_label != "") - { - LabelT u_label = 0; - index->set_universal_label(u_label); - } - - if (max_points_to_insert == 0) - { - max_points_to_insert = num_points; - } - - if (num_points < max_points_to_insert) - throw diskann::ANNException(std::string("num_points(") + std::to_string(num_points) + - ") < max_points_to_insert(" + std::to_string(max_points_to_insert) + ")", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (max_points_to_insert < active_window + consolidate_interval) - throw diskann::ANNException("ERROR: max_points_to_insert < " - "active_window + consolidate_interval", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (consolidate_interval < max_points_to_insert / 1000) - throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, __FUNCSIG__, __FILE__, __LINE__); - - index->set_start_points_at_random(static_cast(start_point_norm)); - - T *data = nullptr; - diskann::alloc_aligned((void **)&data, std::max(consolidate_interval, active_window) * aligned_dim * sizeof(T), - 8 * sizeof(T)); - - std::vector tags(max_points_to_insert); - std::iota(tags.begin(), tags.end(), static_cast(0)); - - diskann::Timer timer; - - std::vector> delete_tasks; - - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, 0, active_window); - insert_next_batch(*index, (size_t)0, active_window, params.num_threads, data, aligned_dim, - pts_to_labels); - }); - insert_task.wait(); - - for (size_t start = active_window; start + consolidate_interval <= max_points_to_insert; - start += consolidate_interval) - { - auto end = std::min(start + consolidate_interval, max_points_to_insert); - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_next_batch(*index, start, end, params.num_threads, data, aligned_dim, - pts_to_labels); - }); - insert_task.wait(); - - if (delete_tasks.size() > 0) - delete_tasks[delete_tasks.size() - 1].wait(); - if (start >= active_window + consolidate_interval) - { - auto start_del = start - active_window - consolidate_interval; - auto end_del = start - active_window; - - delete_tasks.emplace_back(std::async(std::launch::async, [&]() { - delete_and_consolidate(*index, delete_params, (size_t)start_del, (size_t)end_del); - })); - } - } - if (delete_tasks.size() > 0) - delete_tasks[delete_tasks.size() - 1].wait(); - - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - - index->save(save_path_inc.c_str(), true); - - diskann::aligned_free(data); -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type; - uint32_t insert_threads, consolidate_threads, R, L, num_start_pts, Lf, unique_labels_supported; - float alpha, start_point_norm; - size_t max_points_to_insert, active_window, consolidate_interval; - - po::options_description desc{program_options_utils::make_program_description("test_streaming_scenario", - "Test insert deletes & consolidate")}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - - // Required parameters - po::options_description required_configs("Required"); - required_configs.add_options()("data_type", po::value(&data_type)->required(), - program_options_utils::DATA_TYPE_DESCRIPTION); - required_configs.add_options()("dist_fn", po::value(&dist_fn)->required(), - program_options_utils::DISTANCE_FUNCTION_DESCRIPTION); - required_configs.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - program_options_utils::INDEX_PATH_PREFIX_DESCRIPTION); - required_configs.add_options()("data_path", po::value(&data_path)->required(), - program_options_utils::INPUT_DATA_PATH); - required_configs.add_options()("active_window", po::value(&active_window)->required(), - "Program maintains an index over an active window of " - "this size that slides through the data"); - required_configs.add_options()("consolidate_interval", po::value(&consolidate_interval)->required(), - "The program simultaneously adds this number of points to the " - "right of " - "the window while deleting the same number from the left"); - required_configs.add_options()("start_point_norm", po::value(&start_point_norm)->required(), - "Set the start point to a random point on a sphere of this radius"); - - // Optional parameters - po::options_description optional_configs("Optional"); - optional_configs.add_options()("max_degree,R", po::value(&R)->default_value(64), - program_options_utils::MAX_BUILD_DEGREE); - optional_configs.add_options()("Lbuild,L", po::value(&L)->default_value(100), - program_options_utils::GRAPH_BUILD_COMPLEXITY); - optional_configs.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - program_options_utils::GRAPH_BUILD_ALPHA); - optional_configs.add_options()("insert_threads", - po::value(&insert_threads)->default_value(omp_get_num_procs() / 2), - "Number of threads used for inserting into the index (defaults to " - "omp_get_num_procs()/2)"); - optional_configs.add_options()( - "consolidate_threads", po::value(&consolidate_threads)->default_value(omp_get_num_procs() / 2), - "Number of threads used for consolidating deletes to " - "the index (defaults to omp_get_num_procs()/2)"); - optional_configs.add_options()("max_points_to_insert", - po::value(&max_points_to_insert)->default_value(0), - "The number of points from the file that the program streams " - "over "); - optional_configs.add_options()( - "num_start_points", - po::value(&num_start_pts)->default_value(diskann::defaults::NUM_FROZEN_POINTS_DYNAMIC), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - optional_configs.add_options()("label_file", po::value(&label_file)->default_value(""), - "Input label file in txt format for Filtered Index search. " - "The file should contain comma separated filters for each node " - "with each line corresponding to a graph node"); - optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with labels_file"); - optional_configs.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - optional_configs.add_options()("unique_labels_supported", - po::value(&unique_labels_supported)->default_value(0), - "Number of unique labels supported by the dynamic index."); - - // Merge required and optional parameters - desc.add(required_configs).add(optional_configs); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - // Validate arguments - if (start_point_norm == 0) - { - std::cout << "When beginning_index_size is 0, use a start point with " - "appropriate norm" - << std::endl; - return -1; - } - - if (label_type != std::string("ushort") && label_type != std::string("uint")) - { - std::cerr << "Invalid label type. Supported types are uint and ushort" << std::endl; - return -1; - } - - if (data_type != std::string("int8") && data_type != std::string("uint8") && data_type != std::string("float")) - { - std::cerr << "Invalid data type. Supported types are int8, uint8 and float" << std::endl; - return -1; - } - - // TODO: Are additional distance functions supported? - if (dist_fn != std::string("l2") && dist_fn != std::string("mips")) - { - std::cerr << "Invalid distance function. Supported functions are l2 and mips" << std::endl; - return -1; - } - - if (num_start_pts < unique_labels_supported) - { - num_start_pts = unique_labels_supported; - } - - try - { - if (data_type == std::string("uint8")) - { - if (label_type == std::string("ushort")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - else if (label_type == std::string("uint")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - } - else if (data_type == std::string("int8")) - { - if (label_type == std::string("ushort")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - else if (label_type == std::string("uint")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - } - else if (data_type == std::string("float")) - { - if (label_type == std::string("ushort")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - else if (label_type == std::string("uint")) - { - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, max_points_to_insert, active_window, - consolidate_interval, start_point_norm, num_start_pts, index_path_prefix, label_file, - universal_label, Lf); - } - } - } - catch (const std::exception &e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } - catch (...) - { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); - } - - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/CMakeLists.txt deleted file mode 100644 index 3b8cf22..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/CMakeLists.txt +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_COMPILE_WARNING_AS_ERROR ON) - - -add_executable(fvecs_to_bin fvecs_to_bin.cpp) - -add_executable(fvecs_to_bvecs fvecs_to_bvecs.cpp) - -add_executable(rand_data_gen rand_data_gen.cpp) -target_link_libraries(rand_data_gen ${PROJECT_NAME} Boost::program_options) - -add_executable(float_bin_to_int8 float_bin_to_int8.cpp) - -add_executable(ivecs_to_bin ivecs_to_bin.cpp) - -add_executable(count_bfs_levels count_bfs_levels.cpp) -target_link_libraries(count_bfs_levels ${PROJECT_NAME} Boost::program_options) - -add_executable(tsv_to_bin tsv_to_bin.cpp) - -add_executable(bin_to_tsv bin_to_tsv.cpp) - -add_executable(int8_to_float int8_to_float.cpp) -target_link_libraries(int8_to_float ${PROJECT_NAME}) - -add_executable(int8_to_float_scale int8_to_float_scale.cpp) -target_link_libraries(int8_to_float_scale ${PROJECT_NAME}) - -add_executable(uint8_to_float uint8_to_float.cpp) -target_link_libraries(uint8_to_float ${PROJECT_NAME}) - -add_executable(uint32_to_uint8 uint32_to_uint8.cpp) -target_link_libraries(uint32_to_uint8 ${PROJECT_NAME}) - -add_executable(vector_analysis vector_analysis.cpp) -target_link_libraries(vector_analysis ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -add_executable(gen_random_slice gen_random_slice.cpp) -target_link_libraries(gen_random_slice ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -add_executable(simulate_aggregate_recall simulate_aggregate_recall.cpp) - -add_executable(calculate_recall calculate_recall.cpp) -target_link_libraries(calculate_recall ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -# Compute ground truth thing outside of DiskANN main source that depends on MKL. -add_executable(compute_groundtruth compute_groundtruth.cpp) -target_include_directories(compute_groundtruth PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES}) -target_link_libraries(compute_groundtruth ${PROJECT_NAME} ${DISKANN_MKL_LINK_LIBRARIES} ${DISKANN_ASYNC_LIB} Boost::program_options) - -add_executable(compute_groundtruth_for_filters compute_groundtruth_for_filters.cpp) -target_include_directories(compute_groundtruth_for_filters PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES}) -target_link_libraries(compute_groundtruth_for_filters ${PROJECT_NAME} ${DISKANN_MKL_LINK_LIBRARIES} ${DISKANN_ASYNC_LIB} Boost::program_options) - - -add_executable(generate_pq generate_pq.cpp) -target_link_libraries(generate_pq ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - - -add_executable(partition_data partition_data.cpp) -target_link_libraries(partition_data ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -add_executable(partition_with_ram_budget partition_with_ram_budget.cpp) -target_link_libraries(partition_with_ram_budget ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -add_executable(merge_shards merge_shards.cpp) -target_link_libraries(merge_shards ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} ${DISKANN_ASYNC_LIB}) - -add_executable(create_disk_layout create_disk_layout.cpp) -target_link_libraries(create_disk_layout ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS}) - -add_executable(generate_synthetic_labels generate_synthetic_labels.cpp) -target_link_libraries(generate_synthetic_labels ${PROJECT_NAME} Boost::program_options) - -add_executable(stats_label_data stats_label_data.cpp) -target_link_libraries(stats_label_data ${PROJECT_NAME} Boost::program_options) - -if (NOT MSVC) - include(GNUInstallDirs) - install(TARGETS fvecs_to_bin - fvecs_to_bvecs - rand_data_gen - float_bin_to_int8 - ivecs_to_bin - count_bfs_levels - tsv_to_bin - bin_to_tsv - int8_to_float - int8_to_float_scale - uint8_to_float - uint32_to_uint8 - vector_analysis - gen_random_slice - simulate_aggregate_recall - calculate_recall - compute_groundtruth - compute_groundtruth_for_filters - generate_pq - partition_data - partition_with_ram_budget - merge_shards - create_disk_layout - generate_synthetic_labels - stats_label_data - RUNTIME - ) -endif() \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_fvecs.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_fvecs.cpp deleted file mode 100644 index e9a6a8e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_fvecs.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "util.h" - -void block_convert(std::ifstream &writr, std::ofstream &readr, float *read_buf, float *write_buf, uint64_t npts, - uint64_t ndims) -{ - writr.write((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(unsigned))); -#pragma omp parallel for - for (uint64_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(float)); - } - readr.read((char *)write_buf, npts * ndims * sizeof(float)); -} - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_bin output_fvecs" << std::endl; - exit(-1); - } - std::ifstream readr(argv[1], std::ios::binary); - int npts_s32; - int ndims_s32; - readr.read((char *)&npts_s32, sizeof(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - size_t npts = npts_s32; - size_t ndims = ndims_s32; - uint32_t ndims_u32 = (uint32_t)ndims_s32; - // uint64_t fsize = writr.tellg(); - readr.seekg(0, std::ios::beg); - - unsigned ndims_u32; - writr.write((char *)&ndims_u32, sizeof(unsigned)); - writr.seekg(0, std::ios::beg); - uint64_t ndims = (uint64_t)ndims_u32; - uint64_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - uint64_t blk_size = 131072; - uint64_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - - std::ofstream writr(argv[2], std::ios::binary); - float *read_buf = new float[npts * (ndims + 1)]; - float *write_buf = new float[npts * ndims]; - for (uint64_t i = 0; i < nblks; i++) - { - uint64_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writr, readr, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writr.close(); - readr.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_tsv.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_tsv.cpp deleted file mode 100644 index 7851bef..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/bin_to_tsv.cpp +++ /dev/null @@ -1,69 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -template -void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, size_t npts, size_t ndims) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(float)); - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - writer << read_buf[d + i * ndims]; - if (d < ndims - 1) - writer << "\t"; - else - writer << "\n"; - } - } -} - -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " input_bin output_tsv" << std::endl; - exit(-1); - } - std::string type_string(argv[1]); - if ((type_string != std::string("float")) && (type_string != std::string("int8")) && - (type_string != std::string("uin8"))) - { - std::cerr << "Error: type not supported. Use float/int8/uint8" << std::endl; - } - - std::ifstream reader(argv[2], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[3]); - char *read_buf = new char[blk_size * ndims * 4]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (type_string == std::string("float")) - block_convert(writer, reader, (float *)read_buf, cblk_size, ndims); - else if (type_string == std::string("int8")) - block_convert(writer, reader, (int8_t *)read_buf, cblk_size, ndims); - else if (type_string == std::string("uint8")) - block_convert(writer, reader, (uint8_t *)read_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - - writer.close(); - reader.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/calculate_recall.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/calculate_recall.cpp deleted file mode 100644 index dc76252..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/calculate_recall.cpp +++ /dev/null @@ -1,55 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include - -#include "utils.h" -#include "disk_utils.h" - -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " " << std::endl; - return -1; - } - uint32_t *gold_std = NULL; - float *gs_dist = nullptr; - uint32_t *our_results = NULL; - float *or_dist = nullptr; - size_t points_num, points_num_gs, points_num_or; - size_t dim_gs; - size_t dim_or; - diskann::load_truthset(argv[1], gold_std, gs_dist, points_num_gs, dim_gs); - diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); - - if (points_num_gs != points_num_or) - { - std::cout << "Error. Number of queries mismatch in ground truth and " - "our results" - << std::endl; - return -1; - } - points_num = points_num_gs; - - uint32_t recall_at = std::atoi(argv[3]); - - if ((dim_or < recall_at) || (recall_at > dim_gs)) - { - std::cout << "ground truth has size " << dim_gs << "; our set has " << dim_or << " points. Asking for recall " - << recall_at << std::endl; - return -1; - } - std::cout << "Calculating recall@" << recall_at << std::endl; - double recall_val = diskann::calculate_recall((uint32_t)points_num, gold_std, gs_dist, (uint32_t)dim_gs, - our_results, (uint32_t)dim_or, (uint32_t)recall_at); - - // double avg_recall = (recall*1.0)/(points_num*1.0); - std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth.cpp deleted file mode 100644 index da32fd7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth.cpp +++ /dev/null @@ -1,574 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef _WINDOWS -#include -#else -#include -#endif -#include "filter_utils.h" -#include "utils.h" - -// WORKS FOR UPTO 2 BILLION POINTS (as we use INT INSTEAD OF UNSIGNED) - -#define PARTSIZE 10000000 -#define ALIGNMENT 512 - -// custom types (for readability) -typedef tsl::robin_set label_set; -typedef std::string path; - -namespace po = boost::program_options; - -template T div_round_up(const T numerator, const T denominator) -{ - return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); -} - -using pairIF = std::pair; -struct cmpmaxstruct -{ - bool operator()(const pairIF &l, const pairIF &r) - { - return l.second < r.second; - }; -}; - -using maxPQIFCS = std::priority_queue, cmpmaxstruct>; - -template T *aligned_malloc(const size_t n, const size_t alignment) -{ -#ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); -#else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); -#endif -} - -inline bool custom_dist(const std::pair &a, const std::pair &b) -{ - return a.second < b.second; -} - -void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const uint64_t dim) -{ - assert(points_l2sq != NULL); -#pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = cblas_sdot((int64_t)dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, - matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); -} - -void distsq_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 -{ - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, - ones_vec, nqueries, (float)1.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, - queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); - if (ones_vec_alloc) - delete[] ones_vec; -} - -void inner_prod_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 -{ - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - - if (ones_vec_alloc) - delete[] ones_vec; -} - -void exact_knn(const size_t dim, const size_t k, - size_t *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major -{ - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) - { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; -#pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) - { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } - -#pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) - { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } - } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); - } - - std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " - << dim << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) - { - int64_t q_b = b * q_batch_size; - int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) - { - distsq_to_points(dim, dist_matrix, npoints, points, points_l2sq, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim, queries_l2sq + q_b); - } - else - { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); - } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; - -#pragma omp parallel for schedule(dynamic, 16) - for (long long q = q_b; q < q_e; q++) - { - maxPQIFCS point_dist; - for (size_t p = 0; p < k; p++) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) - { - if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) - point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) - { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted(dist_closest_points + (ptrdiff_t)q * (ptrdiff_t)k, - dist_closest_points + (ptrdiff_t)(q + 1) * (ptrdiff_t)k)); - } - std::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e << ")" << std::endl; - } - - delete[] dist_matrix; - - delete[] points_l2sq; - delete[] queries_l2sq; - - if (metric == diskann::Metric::COSINE) - { - delete[] points; - delete[] queries; - } -} - -template inline int get_num_parts(const char *filename) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - std::cout << "#pts = " << npts_i32 << ", #dims = " << ndims_i32 << std::endl; - reader.close(); - uint32_t num_parts = - (npts_i32 % PARTSIZE) == 0 ? npts_i32 / PARTSIZE : (uint32_t)std::floor(npts_i32 / PARTSIZE) + 1; - std::cout << "Number of parts: " << num_parts << std::endl; - return num_parts; -} - -template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, int part_num) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims << ", size = " << npts * ndims * sizeof(T) << "B" - << std::endl; - - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - T *data_T = new T[npts * ndims]; - reader.read((char *)data_T, sizeof(T) * npts * ndims); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts * ndims, ALIGNMENT); -#pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (int64_t j = 0; j < (int64_t)ndims; j++) - { - float cur_val_float = (float)data_T[i * ndims + j]; - std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, sizeof(float)); - } - } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; -} - -template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) -{ - std::ofstream writer; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::cout << "Writing bin: " << filename << "\n"; - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - std::cout << "bin: #pts = " << npts << ", #dims = " << ndims - << ", size = " << npts * ndims * sizeof(T) + 2 * sizeof(int) << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; -} - -inline void save_groundtruth_as_one_file(const std::string filename, int32_t *data, float *distances, size_t npts, - size_t ndims) -{ - std::ofstream writer(filename, std::ios::binary | std::ios::out); - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - std::cout << "Saving truthset in one file (npts, dim, npts*dim id-matrix, " - "npts*dim dist-matrix) with npts = " - << npts << ", dim = " << ndims << ", size = " << 2 * npts * ndims * sizeof(uint32_t) + 2 * sizeof(int) - << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; -} - -template -std::vector>> processUnfilteredParts(const std::string &base_file, - size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, - std::vector &location_to_tag) -{ - float *base_data = nullptr; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - size_t *closest_points_part = new size_t[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - auto part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, - metric); - - for (size_t i = 0; i < nqueries; i++) - { - for (size_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } - } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); - } - return res; -}; - -template -int aux_main(const std::string &base_file, const std::string &query_file, const std::string >_file, size_t k, - const diskann::Metric &metric, const std::string &tags_file = std::string("")) -{ - size_t npoints, nqueries, dim; - - float *query_data; - - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results = - processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); - - for (size_t i = 0; i < nqueries; i++) - { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) - { - if (j == k) - break; - if (tags_enabled) - { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } - else - { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; - } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; - } - - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); - - return 0; -} - -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file; - uint64_t K; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("base_file", po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary " - "format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it " - "will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; - return -1; - } - - try - { - if (data_type == std::string("float")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth_for_filters.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth_for_filters.cpp deleted file mode 100644 index 52e5864..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/compute_groundtruth_for_filters.cpp +++ /dev/null @@ -1,919 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef _WINDOWS -#include -#else -#include -#endif - -#include "filter_utils.h" -#include "utils.h" - -// WORKS FOR UPTO 2 BILLION POINTS (as we use INT INSTEAD OF UNSIGNED) - -#define PARTSIZE 10000000 -#define ALIGNMENT 512 - -// custom types (for readability) -typedef tsl::robin_set label_set; -typedef std::string path; - -namespace po = boost::program_options; - -template T div_round_up(const T numerator, const T denominator) -{ - return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); -} - -using pairIF = std::pair; -struct cmpmaxstruct -{ - bool operator()(const pairIF &l, const pairIF &r) - { - return l.second < r.second; - }; -}; - -using maxPQIFCS = std::priority_queue, cmpmaxstruct>; - -template T *aligned_malloc(const size_t n, const size_t alignment) -{ -#ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); -#else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); -#endif -} - -inline bool custom_dist(const std::pair &a, const std::pair &b) -{ - return a.second < b.second; -} - -void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const uint64_t dim) -{ - assert(points_l2sq != NULL); -#pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = cblas_sdot((int64_t)dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, - matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); -} - -void distsq_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 -{ - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, - ones_vec, nqueries, (float)1.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, - queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); - if (ones_vec_alloc) - delete[] ones_vec; -} - -void inner_prod_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 -{ - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - - if (ones_vec_alloc) - delete[] ones_vec; -} - -void exact_knn(const size_t dim, const size_t k, - size_t *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major -{ - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) - { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; -#pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) - { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } - -#pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) - { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } - } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); - } - - std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " - << dim << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (uint64_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) - { - int64_t q_b = b * q_batch_size; - int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) - { - distsq_to_points(dim, dist_matrix, npoints, points, points_l2sq, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim, queries_l2sq + q_b); - } - else - { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); - } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; - -#pragma omp parallel for schedule(dynamic, 16) - for (long long q = q_b; q < q_e; q++) - { - maxPQIFCS point_dist; - for (size_t p = 0; p < k; p++) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) - { - if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) - point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) - { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted(dist_closest_points + (ptrdiff_t)q * (ptrdiff_t)k, - dist_closest_points + (ptrdiff_t)(q + 1) * (ptrdiff_t)k)); - } - std::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e << ")" << std::endl; - } - - delete[] dist_matrix; - - delete[] points_l2sq; - delete[] queries_l2sq; - - if (metric == diskann::Metric::COSINE) - { - delete[] points; - delete[] queries; - } -} - -template inline int get_num_parts(const char *filename) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - std::cout << "#pts = " << npts_i32 << ", #dims = " << ndims_i32 << std::endl; - reader.close(); - int num_parts = (npts_i32 % PARTSIZE) == 0 ? npts_i32 / PARTSIZE : (uint32_t)std::floor(npts_i32 / PARTSIZE) + 1; - std::cout << "Number of parts: " << num_parts << std::endl; - return num_parts; -} - -template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts_u64, size_t &ndims_u64, int part_num) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts_u64 = end_id - start_id; - ndims_u64 = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 - << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" << std::endl; - - reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - T *data_T = new T[npts_u64 * ndims_u64]; - reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); -#pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts_u64; i++) - { - for (int64_t j = 0; j < (int64_t)ndims_u64; j++) - { - float cur_val_float = (float)data_T[i * ndims_u64 + j]; - std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, sizeof(float)); - } - } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; -} - -template -inline std::vector load_filtered_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, - int part_num, const char *label_file, - const std::string &filter_label, - const std::string &universal_label, size_t &npoints_filt, - std::vector> &pts_to_labels) -{ - std::ifstream reader(filename, std::ios::binary); - if (reader.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); - } - - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - std::vector rev_map; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint32_t)ndims_i32; - uint64_t nptsuint64_t = (uint64_t)npts; - uint64_t ndimsuint64_t = (uint64_t)ndims; - npoints_filt = 0; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims - << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" << std::endl; - std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl; - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - - T *data_T = new T[nptsuint64_t * ndimsuint64_t]; - reader.read((char *)data_T, sizeof(T) * nptsuint64_t * ndimsuint64_t); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - - data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); - - for (int64_t i = 0; i < (int64_t)nptsuint64_t; i++) - { - if (std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), filter_label) != - pts_to_labels[start_id + i].end() || - std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), universal_label) != - pts_to_labels[start_id + i].end()) - { - rev_map.push_back(start_id + i); - for (int64_t j = 0; j < (int64_t)ndimsuint64_t; j++) - { - float cur_val_float = (float)data_T[i * ndimsuint64_t + j]; - std::memcpy((char *)(data + npoints_filt * ndimsuint64_t + j), (char *)&cur_val_float, sizeof(float)); - } - npoints_filt++; - } - } - delete[] data_T; - std::cout << "Finished converting part data to float.. identified " << npoints_filt - << " points matching the filter." << std::endl; - return rev_map; -} - -template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) -{ - std::ofstream writer; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::cout << "Writing bin: " << filename << "\n"; - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - std::cout << "bin: #pts = " << npts << ", #dims = " << ndims - << ", size = " << npts * ndims * sizeof(T) + 2 * sizeof(int) << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; -} - -inline void save_groundtruth_as_one_file(const std::string filename, int32_t *data, float *distances, size_t npts, - size_t ndims) -{ - std::ofstream writer(filename, std::ios::binary | std::ios::out); - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - std::cout << "Saving truthset in one file (npts, dim, npts*dim id-matrix, " - "npts*dim dist-matrix) with npts = " - << npts << ", dim = " << ndims << ", size = " << 2 * npts * ndims * sizeof(uint32_t) + 2 * sizeof(int) - << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; -} - -inline void parse_label_file_into_vec(size_t &line_cnt, const std::string &map_file, - std::vector> &pts_to_labels) -{ - std::ifstream infile(map_file); - std::string line, token; - std::set labels; - infile.clear(); - infile.seekg(0, std::ios::beg); - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - lbls.push_back(token); - labels.insert(token); - } - std::sort(lbls.begin(), lbls.end()); - pts_to_labels.push_back(lbls); - } - std::cout << "Identified " << labels.size() << " distinct label(s), and populated labels for " - << pts_to_labels.size() << " points" << std::endl; -} - -template -std::vector>> processUnfilteredParts(const std::string &base_file, - size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, - std::vector &location_to_tag) -{ - float *base_data = nullptr; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - size_t *closest_points_part = new size_t[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - auto part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, - metric); - - for (size_t i = 0; i < nqueries; i++) - { - for (uint64_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } - } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); - } - return res; -}; - -template -std::vector>> processFilteredParts( - const std::string &base_file, const std::string &label_file, const std::string &filter_label, - const std::string &universal_label, size_t &nqueries, size_t &npoints, size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, std::vector &location_to_tag) -{ - size_t npoints_filt = 0; - float *base_data = nullptr; - std::vector>> res(nqueries); - int num_parts = get_num_parts(base_file.c_str()); - - std::vector> pts_to_labels; - if (filter_label != "") - parse_label_file_into_vec(npoints, label_file, pts_to_labels); - - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - std::vector rev_map; - if (filter_label != "") - rev_map = load_filtered_bin_as_float(base_file.c_str(), base_data, npoints, dim, p, label_file.c_str(), - filter_label, universal_label, npoints_filt, pts_to_labels); - size_t *closest_points_part = new size_t[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - auto part_k = k < npoints_filt ? k : npoints_filt; - if (npoints_filt > 0) - { - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints_filt, base_data, nqueries, - query_data, metric); - } - - for (size_t i = 0; i < nqueries; i++) - { - for (uint64_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(rev_map[closest_points_part[i * part_k + j]]), - dist_closest_points_part[i * part_k + j])); - } - } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); - } - return res; -}; - -template -int aux_main(const std::string &base_file, const std::string &label_file, const std::string &query_file, - const std::string >_file, size_t k, const std::string &universal_label, const diskann::Metric &metric, - const std::string &filter_label, const std::string &tags_file = std::string("")) -{ - size_t npoints, nqueries, dim; - - float *query_data = nullptr; - - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results; - if (filter_label == "") - { - results = processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); - } - else - { - results = processFilteredParts(base_file, label_file, filter_label, universal_label, nqueries, npoints, dim, - k, query_data, metric, location_to_tag); - } - - for (size_t i = 0; i < nqueries; i++) - { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) - { - if (j == k) - break; - if (tags_enabled) - { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } - else - { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; - } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; - } - - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); - - return 0; -} - -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file, label_file, filter_label, - universal_label, filter_label_file; - uint64_t K; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("base_file", po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("label_file", po::value(&label_file)->default_value(""), - "Input labels file in txt format if present"); - desc.add_options()("filter_label", po::value(&filter_label)->default_value(""), - "Input filter label if doing filtered groundtruth"); - desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with label_file"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary " - "format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it " - "will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - desc.add_options()("filter_label_file", - po::value(&filter_label_file)->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; - } - - if (filter_label != "" && filter_label_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; - return -1; - } - - std::vector filter_labels; - if (filter_label != "") - { - filter_labels.push_back(filter_label); - } - else if (filter_label_file != "") - { - filter_labels = read_file_to_vector_of_strings(filter_label_file, false); - } - - // only when there is no filter label or 1 filter label for all queries - if (filter_labels.size() == 1) - { - try - { - if (data_type == std::string("float")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } - } - else - { // Each query has its own filter label - // Split up data and query bins into label specific ones - tsl::robin_map labels_to_number_of_points; - tsl::robin_map labels_to_number_of_queries; - - label_set all_labels; - for (size_t i = 0; i < filter_labels.size(); i++) - { - std::string label = filter_labels[i]; - all_labels.insert(label); - - if (labels_to_number_of_queries.find(label) == labels_to_number_of_queries.end()) - { - labels_to_number_of_queries[label] = 0; - } - labels_to_number_of_queries[label] += 1; - } - - size_t npoints; - std::vector> point_to_labels; - parse_label_file_into_vec(npoints, label_file, point_to_labels); - std::vector point_ids_to_labels(point_to_labels.size()); - std::vector query_ids_to_labels(filter_labels.size()); - - for (size_t i = 0; i < point_to_labels.size(); i++) - { - for (size_t j = 0; j < point_to_labels[i].size(); j++) - { - std::string label = point_to_labels[i][j]; - if (all_labels.find(label) != all_labels.end()) - { - point_ids_to_labels[i].insert(point_to_labels[i][j]); - if (labels_to_number_of_points.find(label) == labels_to_number_of_points.end()) - { - labels_to_number_of_points[label] = 0; - } - labels_to_number_of_points[label] += 1; - } - } - } - - for (size_t i = 0; i < filter_labels.size(); i++) - { - query_ids_to_labels[i].insert(filter_labels[i]); - } - - tsl::robin_map> label_id_to_orig_id; - tsl::robin_map> label_query_id_to_orig_id; - - if (data_type == std::string("float")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else if (data_type == std::string("int8")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else if (data_type == std::string("uint8")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else - { - diskann::cerr << "Invalid data type" << std::endl; - return -1; - } - - // Generate label specific ground truths - - try - { - for (const auto &label : all_labels) - { - std::string filtered_base_file = base_file + "_" + label; - std::string filtered_query_file = query_file + "_" + label; - std::string filtered_gt_file = gt_file + "_" + label; - if (data_type == std::string("float")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("int8")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("uint8")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } - - // Combine the label specific ground truths to produce a single GT file - - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t gt_num, gt_dim; - - std::vector> final_gt_ids; - std::vector> final_gt_dists; - - uint32_t query_num = 0; - for (const auto &lbl : all_labels) - { - query_num += labels_to_number_of_queries[lbl]; - } - - for (uint32_t i = 0; i < query_num; i++) - { - final_gt_ids.push_back(std::vector(K)); - final_gt_dists.push_back(std::vector(K)); - } - - for (const auto &lbl : all_labels) - { - std::string filtered_gt_file = gt_file + "_" + lbl; - load_truthset(filtered_gt_file, gt_ids, gt_dists, gt_num, gt_dim); - - for (uint32_t i = 0; i < labels_to_number_of_queries[lbl]; i++) - { - uint32_t orig_query_id = label_query_id_to_orig_id[lbl][i]; - for (uint64_t j = 0; j < K; j++) - { - final_gt_ids[orig_query_id][j] = label_id_to_orig_id[lbl][gt_ids[i * K + j]]; - final_gt_dists[orig_query_id][j] = gt_dists[i * K + j]; - } - } - } - - int32_t *closest_points = new int32_t[query_num * K]; - float *dist_closest_points = new float[query_num * K]; - - for (uint32_t i = 0; i < query_num; i++) - { - for (uint32_t j = 0; j < K; j++) - { - closest_points[i * K + j] = final_gt_ids[i][j]; - dist_closest_points[i * K + j] = final_gt_dists[i][j]; - } - } - - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, query_num, K); - - // cleanup artifacts - std::cout << "Cleaning up artifacts..." << std::endl; - tsl::robin_set paths_to_clean{gt_file, base_file, query_file}; - clean_up_artifacts(paths_to_clean, all_labels); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/count_bfs_levels.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/count_bfs_levels.cpp deleted file mode 100644 index 6dd2d62..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/count_bfs_levels.cpp +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifndef _WINDOWS -#include -#include -#include -#include -#endif - -#include "utils.h" -#include "index.h" -#include "memory_mapper.h" - -namespace po = boost::program_options; - -template void bfs_count(const std::string &index_path, uint32_t data_dims) -{ - using TagT = uint32_t; - using LabelT = uint32_t; - diskann::Index index(diskann::Metric::L2, data_dims, 0, nullptr, nullptr, 0, false, false, false, - false, 0, false); - std::cout << "Index class instantiated" << std::endl; - index.load(index_path.c_str(), 1, 100); - std::cout << "Index loaded" << std::endl; - index.count_nodes_at_bfs_levels(); -} - -int main(int argc, char **argv) -{ - std::string data_type, index_path_prefix; - uint32_t data_dims; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("data_dims", po::value(&data_dims)->required(), "Dimensionality of the data"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - try - { - if (data_type == std::string("int8")) - bfs_count(index_path_prefix, data_dims); - else if (data_type == std::string("uint8")) - bfs_count(index_path_prefix, data_dims); - if (data_type == std::string("float")) - bfs_count(index_path_prefix, data_dims); - } - catch (std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index BFS failed." << std::endl; - return -1; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/create_disk_layout.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/create_disk_layout.cpp deleted file mode 100644 index f494c12..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/create_disk_layout.cpp +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include - -#include "utils.h" -#include "disk_utils.h" -#include "cached_io.h" - -template int create_disk_layout(char **argv) -{ - std::string base_file(argv[2]); - std::string vamana_file(argv[3]); - std::string output_file(argv[4]); - diskann::create_disk_layout(base_file, vamana_file, output_file); - return 0; -} - -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << argv[0] - << " data_type data_bin " - "vamana_index_file output_diskann_index_file" - << std::endl; - exit(-1); - } - - int ret_val = -1; - if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argv); - else - { - std::cout << "unsupported type. use int8/uint8/float " << std::endl; - ret_val = -2; - } - return ret_val; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/float_bin_to_int8.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/float_bin_to_int8.cpp deleted file mode 100644 index 1982005..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/float_bin_to_int8.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -void block_convert(std::ofstream &writer, int8_t *write_buf, std::ifstream &reader, float *read_buf, size_t npts, - size_t ndims, float bias, float scale) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(float)); - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - write_buf[d + i * ndims] = (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); - } - } - writer.write((char *)write_buf, npts * ndims); -} - -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << "Usage: " << argv[0] << " input_bin output_tsv bias scale" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[blk_size * ndims]; - auto write_buf = new int8_t[blk_size * ndims]; - float bias = (float)atof(argv[3]); - float scale = (float)atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bin.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bin.cpp deleted file mode 100644 index 873ad3b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bin.cpp +++ /dev/null @@ -1,95 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -// Convert float types -void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *read_buf, float *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(float)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(float)); -} - -// Convert byte types -void block_convert_byte(std::ifstream &reader, std::ofstream &writer, uint8_t *read_buf, uint8_t *write_buf, - size_t npts, size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(uint8_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + sizeof(uint32_t))) + sizeof(uint32_t), - ndims * sizeof(uint8_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint8_t)); -} - -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " input_vecs output_bin" << std::endl; - exit(-1); - } - - int datasize = sizeof(float); - - if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) - { - datasize = sizeof(uint8_t); - } - else if (strcmp(argv[1], "float") != 0) - { - std::cout << "Error: type not supported. Use float/int8/uint8" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims * datasize) + sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - int32_t npts_s32 = (int32_t)npts; - int32_t ndims_s32 = (int32_t)ndims; - writer.write((char *)&npts_s32, sizeof(int32_t)); - writer.write((char *)&ndims_s32, sizeof(int32_t)); - - size_t chunknpts = std::min(npts, blk_size); - uint8_t *read_buf = new uint8_t[chunknpts * ((ndims * datasize) + sizeof(uint32_t))]; - uint8_t *write_buf = new uint8_t[chunknpts * ndims * datasize]; - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (datasize == sizeof(float)) - { - block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, cblk_size, ndims); - } - else - { - block_convert_byte(reader, writer, read_buf, write_buf, cblk_size, ndims); - } - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - reader.close(); - writer.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bvecs.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bvecs.cpp deleted file mode 100644 index f9c2aa7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/fvecs_to_bvecs.cpp +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -void block_convert(std::ifstream &reader, std::ofstream &writer, float *read_buf, uint8_t *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), sizeof(uint32_t)); - for (size_t d = 0; d < ndims; d++) - write_buf[i * (ndims + 4) + 4 + d] = (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; - } - writer.write((char *)write_buf, npts * (ndims * 1 + 4)); -} - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_fvecs output_bvecs(uint8)" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[npts * (ndims + 1)]; - auto write_buf = new uint8_t[npts * (ndims + 4)]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(reader, writer, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - reader.close(); - writer.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/gen_random_slice.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/gen_random_slice.cpp deleted file mode 100644 index a4cd96e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/gen_random_slice.cpp +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "partition.h" -#include "utils.h" - -#include -#include -#include -#include - -template int aux_main(char **argv) -{ - std::string base_file(argv[2]); - std::string output_prefix(argv[3]); - float sampling_rate = (float)(std::atof(argv[4])); - gen_random_slice(base_file, output_prefix, sampling_rate); - return 0; -} - -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } - - if (std::string(argv[1]) == std::string("float")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("int8")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - aux_main(argv); - } - else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_pq.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_pq.cpp deleted file mode 100644 index a881b11..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_pq.cpp +++ /dev/null @@ -1,70 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "math_utils.h" -#include "pq.h" -#include "partition.h" - -#define KMEANS_ITERS_FOR_PQ 15 - -template -bool generate_pq(const std::string &data_path, const std::string &index_prefix_path, const size_t num_pq_centers, - const size_t num_pq_chunks, const float sampling_rate, const bool opq) -{ - std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; - - // generates random sample and sets it to train_data and updates train_size - size_t train_size, train_dim; - float *train_data; - gen_random_slice(data_path, sampling_rate, train_data, train_size, train_dim); - std::cout << "For computing pivots, loaded sample data of size " << train_size << std::endl; - - if (opq) - { - diskann::generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, (uint32_t)num_pq_centers, - (uint32_t)num_pq_chunks, pq_pivots_path, true); - } - else - { - diskann::generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, (uint32_t)num_pq_centers, - (uint32_t)num_pq_chunks, KMEANS_ITERS_FOR_PQ, pq_pivots_path); - } - diskann::generate_pq_data_from_pivots(data_path, (uint32_t)num_pq_centers, (uint32_t)num_pq_chunks, - pq_pivots_path, pq_compressed_vectors_path, true); - - delete[] train_data; - - return 0; -} - -int main(int argc, char **argv) -{ - if (argc != 7) - { - std::cout << "Usage: \n" - << argv[0] - << " " - " " - " " - << std::endl; - } - else - { - const std::string data_path(argv[2]); - const std::string index_prefix_path(argv[3]); - const size_t num_pq_centers = 256; - const size_t num_pq_chunks = (size_t)atoi(argv[4]); - const float sampling_rate = (float)atof(argv[5]); - const bool opq = atoi(argv[6]) == 0 ? false : true; - - if (std::string(argv[1]) == std::string("float")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("int8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("uint8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else - std::cout << "Error. wrong file type" << std::endl; - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_synthetic_labels.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_synthetic_labels.cpp deleted file mode 100644 index 6741760..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/generate_synthetic_labels.cpp +++ /dev/null @@ -1,204 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include "utils.h" - -namespace po = boost::program_options; -class ZipfDistribution -{ - public: - ZipfDistribution(uint64_t num_points, uint32_t num_labels) - : num_labels(num_labels), num_points(num_points), - uniform_zero_to_one(std::uniform_real_distribution<>(0.0, 1.0)) - { - } - - std::unordered_map createDistributionMap() - { - std::unordered_map map; - uint32_t primary_label_freq = (uint32_t)ceil(num_points * distribution_factor); - for (uint32_t i{1}; i < num_labels + 1; i++) - { - map[i] = (uint32_t)ceil(primary_label_freq / i); - } - return map; - } - - int writeDistribution(std::ofstream &outfile) - { - auto distribution_map = createDistributionMap(); - for (uint32_t i{0}; i < num_points; i++) - { - bool label_written = false; - for (auto it = distribution_map.cbegin(); it != distribution_map.cend(); it++) - { - auto label_selection_probability = std::bernoulli_distribution(distribution_factor / (double)it->first); - if (label_selection_probability(rand_engine) && distribution_map[it->first] > 0) - { - if (label_written) - { - outfile << ','; - } - outfile << it->first; - label_written = true; - // remove label from map if we have used all labels - distribution_map[it->first] -= 1; - } - } - if (!label_written) - { - outfile << 0; - } - if (i < num_points - 1) - { - outfile << '\n'; - } - } - return 0; - } - - int writeDistribution(std::string filename) - { - std::ofstream outfile(filename); - if (!outfile.is_open()) - { - std::cerr << "Error: could not open output file " << filename << '\n'; - return -1; - } - writeDistribution(outfile); - outfile.close(); - } - - private: - const uint32_t num_labels; - const uint64_t num_points; - const double distribution_factor = 0.7; - std::knuth_b rand_engine; - const std::uniform_real_distribution uniform_zero_to_one; -}; - -int main(int argc, char **argv) -{ - std::string output_file, distribution_type; - uint32_t num_labels; - uint64_t num_points; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("output_file,O", po::value(&output_file)->required(), - "Filename for saving the label file"); - desc.add_options()("num_points,N", po::value(&num_points)->required(), "Number of points in dataset"); - desc.add_options()("num_labels,L", po::value(&num_labels)->required(), - "Number of unique labels, up to 5000"); - desc.add_options()("distribution_type,DT", po::value(&distribution_type)->default_value("random"), - "Distribution function for labels defaults " - "to random"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (num_labels > 5000) - { - std::cerr << "Error: num_labels must be 5000 or less" << '\n'; - return -1; - } - - if (num_points <= 0) - { - std::cerr << "Error: num_points must be greater than 0" << '\n'; - return -1; - } - - std::cout << "Generating synthetic labels for " << num_points << " points with " << num_labels << " unique labels" - << '\n'; - - try - { - std::ofstream outfile(output_file); - if (!outfile.is_open()) - { - std::cerr << "Error: could not open output file " << output_file << '\n'; - return -1; - } - - if (distribution_type == "zipf") - { - ZipfDistribution zipf(num_points, num_labels); - zipf.writeDistribution(outfile); - } - else if (distribution_type == "random") - { - for (size_t i = 0; i < num_points; i++) - { - bool label_written = false; - for (size_t j = 1; j <= num_labels; j++) - { - // 50% chance to assign each label - if (rand() > (RAND_MAX / 2)) - { - if (label_written) - { - outfile << ','; - } - outfile << j; - label_written = true; - } - } - if (!label_written) - { - outfile << 0; - } - if (i < num_points - 1) - { - outfile << '\n'; - } - } - } - else if (distribution_type == "one_per_point") - { - std::random_device rd; // obtain a random number from hardware - std::mt19937 gen(rd()); // seed the generator - std::uniform_int_distribution<> distr(0, num_labels); // define the range - - for (size_t i = 0; i < num_points; i++) - { - outfile << distr(gen); - if (i != num_points - 1) - outfile << '\n'; - } - } - if (outfile.is_open()) - { - outfile.close(); - } - - std::cout << "Labels written to " << output_file << '\n'; - } - catch (const std::exception &ex) - { - std::cerr << "Label generation failed: " << ex.what() << '\n'; - return -1; - } - - return 0; -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float.cpp deleted file mode 100644 index dcdfddc..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; - exit(-1); - } - - int8_t *input; - size_t npts, nd; - diskann::load_bin(argv[1], input, npts, nd); - float *output = new float[npts * nd]; - diskann::convert_types(input, output, npts, nd); - diskann::save_bin(argv[2], output, npts, nd); - delete[] output; - delete[] input; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float_scale.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float_scale.cpp deleted file mode 100644 index 19fbc6c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/int8_to_float_scale.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -void block_convert(std::ofstream &writer, float *write_buf, std::ifstream &reader, int8_t *read_buf, size_t npts, - size_t ndims, float bias, float scale) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - write_buf[d + i * ndims] = (((float)read_buf[d + i * ndims] - bias) * scale); - } - } - writer.write((char *)write_buf, npts * ndims * sizeof(float)); -} - -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << "Usage: " << argv[0] << " input-int8.bin output-float.bin bias scale" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new int8_t[blk_size * ndims]; - auto write_buf = new float[blk_size * ndims]; - float bias = (float)atof(argv[3]); - float scale = (float)atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/ivecs_to_bin.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/ivecs_to_bin.cpp deleted file mode 100644 index ea8a4a3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/ivecs_to_bin.cpp +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -void block_convert(std::ifstream &reader, std::ofstream &writer, uint32_t *read_buf, uint32_t *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(uint32_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(uint32_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); -} - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_ivecs output_bin" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - int npts_s32 = (int)npts; - int ndims_s32 = (int)ndims; - writer.write((char *)&npts_s32, sizeof(int)); - writer.write((char *)&ndims_s32, sizeof(int)); - uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; - uint32_t *write_buf = new uint32_t[npts * ndims]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(reader, writer, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - reader.close(); - writer.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/merge_shards.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/merge_shards.cpp deleted file mode 100644 index 106c15e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/merge_shards.cpp +++ /dev/null @@ -1,42 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "disk_utils.h" -#include "cached_io.h" -#include "utils.h" - -int main(int argc, char **argv) -{ - if (argc != 9) - { - std::cout << argv[0] - << " vamana_index_prefix[1] vamana_index_suffix[2] " - "idmaps_prefix[3] " - "idmaps_suffix[4] n_shards[5] max_degree[6] " - "output_vamana_path[7] " - "output_medoids_path[8]" - << std::endl; - exit(-1); - } - - std::string vamana_prefix(argv[1]); - std::string vamana_suffix(argv[2]); - std::string idmaps_prefix(argv[3]); - std::string idmaps_suffix(argv[4]); - uint64_t nshards = (uint64_t)std::atoi(argv[5]); - uint32_t max_degree = (uint64_t)std::atoi(argv[6]); - std::string output_index(argv[7]); - std::string output_medoids(argv[8]); - - return diskann::merge_shards(vamana_prefix, vamana_suffix, idmaps_prefix, idmaps_suffix, nshards, max_degree, - output_index, output_medoids); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_data.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_data.cpp deleted file mode 100644 index 2520f3f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_data.cpp +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include "cached_io.h" -#include "partition.h" - -// DEPRECATED: NEED TO REPROGRAM - -int main(int argc, char **argv) -{ - if (argc != 7) - { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } - - const std::string data_path(argv[2]); - const std::string prefix_path(argv[3]); - const float sampling_rate = (float)atof(argv[4]); - const size_t num_partitions = (size_t)std::atoi(argv[5]); - const size_t max_reps = 15; - const size_t k_index = (size_t)std::atoi(argv[6]); - - if (std::string(argv[1]) == std::string("float")) - partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); - else - std::cout << "unsupported data format. use float/int8/uint8" << std::endl; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_with_ram_budget.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_with_ram_budget.cpp deleted file mode 100644 index 937b68d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/partition_with_ram_budget.cpp +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include "cached_io.h" -#include "partition.h" - -// DEPRECATED: NEED TO REPROGRAM - -int main(int argc, char **argv) -{ - if (argc != 8) - { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } - - const std::string data_path(argv[2]); - const std::string prefix_path(argv[3]); - const float sampling_rate = (float)atof(argv[4]); - const double ram_budget = (double)std::atof(argv[5]); - const size_t graph_degree = (size_t)std::atoi(argv[6]); - const size_t k_index = (size_t)std::atoi(argv[7]); - - if (std::string(argv[1]) == std::string("float")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); - else - std::cout << "unsupported data format. use float/int8/uint8" << std::endl; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/rand_data_gen.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/rand_data_gen.cpp deleted file mode 100644 index e89ede8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/rand_data_gen.cpp +++ /dev/null @@ -1,237 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include - -#include "utils.h" - -namespace po = boost::program_options; - -int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, bool normalization, float norm, - float rand_scale) -{ - auto vec = new float[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - std::uniform_real_distribution<> unif_dis(1.0, rand_scale); - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - float scale = 1.0f; - if (rand_scale > 1.0f) - scale = (float)unif_dis(gen); - for (size_t d = 0; d < ndims; ++d) - vec[d] = scale * (float)normal_rand(gen); - if (normalization) - { - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - } - - writer.write((char *)vec, ndims * sizeof(float)); - } - - delete[] vec; - return 0; -} - -int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, float norm) -{ - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) - vec[d] = (float)normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - - for (size_t d = 0; d < ndims; ++d) - { - vec_T[d] = (int8_t)std::round(vec[d]); - } - - writer.write((char *)vec_T, ndims * sizeof(int8_t)); - } - - delete[] vec; - delete[] vec_T; - return 0; -} - -int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, float norm) -{ - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) - vec[d] = (float)normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - - for (size_t d = 0; d < ndims; ++d) - { - vec_T[d] = 128 + (int8_t)std::round(vec[d]); - } - - writer.write((char *)vec_T, ndims * sizeof(uint8_t)); - } - - delete[] vec; - delete[] vec_T; - return 0; -} - -int main(int argc, char **argv) -{ - std::string data_type, output_file; - size_t ndims, npts; - float norm, rand_scaling; - bool normalization = false; - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("output_file", po::value(&output_file)->required(), - "File name for saving the random vectors"); - desc.add_options()("ndims,D", po::value(&ndims)->required(), "Dimensoinality of the vector"); - desc.add_options()("npts,N", po::value(&npts)->required(), "Number of vectors"); - desc.add_options()("norm", po::value(&norm)->default_value(-1.0f), - "Norm of the vectors (if not specified, vectors are not normalized)"); - desc.add_options()("rand_scaling", po::value(&rand_scaling)->default_value(1.0f), - "Each vector will be scaled (if not explicitly normalized) by a factor randomly chosen from " - "[1, rand_scale]. Only applicable for floating point data"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; - } - - if (norm > 0.0) - { - normalization = true; - } - - if (rand_scaling < 1.0) - { - std::cout << "We will only scale the vector norms randomly in [1, value], so value must be >= 1." << std::endl; - return -1; - } - - if ((rand_scaling > 1.0) && (normalization == true)) - { - std::cout << "Data cannot be normalized and randomly scaled at same time. Use one or the other." << std::endl; - return -1; - } - - if (data_type == std::string("int8") || data_type == std::string("uint8")) - { - if (norm > 127) - { - std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " - "greater " - "than 127" - << std::endl; - return -1; - } - if (rand_scaling > 1.0) - { - std::cout << "Data scaling only supported for floating point data." << std::endl; - return -1; - } - } - - try - { - std::ofstream writer; - writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); - writer.open(output_file, std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - - int ret = 0; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (data_type == std::string("float")) - { - ret = block_write_float(writer, ndims, cblk_size, normalization, norm, rand_scaling); - } - else if (data_type == std::string("int8")) - { - ret = block_write_int8(writer, ndims, cblk_size, norm); - } - else if (data_type == std::string("uint8")) - { - ret = block_write_uint8(writer, ndims, cblk_size, norm); - } - if (ret == 0) - std::cout << "Block #" << i << " written" << std::endl; - else - { - writer.close(); - std::cout << "failed to write" << std::endl; - return -1; - } - } - writer.close(); - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } - - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/simulate_aggregate_recall.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/simulate_aggregate_recall.cpp deleted file mode 100644 index 73c4ea0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/simulate_aggregate_recall.cpp +++ /dev/null @@ -1,85 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include - -inline float aggregate_recall(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, uint32_t *count, - const std::vector &recalls) -{ - float found = 0; - for (uint32_t i = 0; i < npart; ++i) - { - size_t max_found = std::min(count[i], k); - found += recalls[max_found - 1] * max_found; - } - return found / (float)k_aggr; -} - -void simulate(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, const uint32_t nsim, - const std::vector &recalls) -{ - std::random_device r; - std::default_random_engine randeng(r()); - std::uniform_int_distribution uniform_dist(0, npart - 1); - - uint32_t *count = new uint32_t[npart]; - double aggr_recall = 0; - - for (uint32_t i = 0; i < nsim; ++i) - { - for (uint32_t p = 0; p < npart; ++p) - { - count[p] = 0; - } - for (uint32_t t = 0; t < k_aggr; ++t) - { - count[uniform_dist(randeng)]++; - } - aggr_recall += aggregate_recall(k_aggr, k, npart, count, recalls); - } - - std::cout << "Aggregate recall is " << aggr_recall / (double)nsim << std::endl; - delete[] count; -} - -int main(int argc, char **argv) -{ - if (argc < 6) - { - std::cout << argv[0] << " k_aggregate k_out npart nsim recall@1 recall@2 ... recall@k" << std::endl; - exit(-1); - } - - const uint32_t k_aggr = atoi(argv[1]); - const uint32_t k = atoi(argv[2]); - const uint32_t npart = atoi(argv[3]); - const uint32_t nsim = atoi(argv[4]); - - std::vector recalls; - for (int ctr = 5; ctr < argc; ctr++) - { - recalls.push_back((float)atof(argv[ctr])); - } - - if (recalls.size() != k) - { - std::cerr << "Please input k numbers for recall@1, recall@2 .. recall@k" << std::endl; - } - if (k_aggr > npart * k) - { - std::cerr << "k_aggr must be <= k * npart" << std::endl; - exit(-1); - } - if (nsim <= npart * k_aggr) - { - std::cerr << "Choose nsim > npart*k_aggr" << std::endl; - exit(-1); - } - - simulate(k_aggr, k, npart, nsim, recalls); - - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/stats_label_data.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/stats_label_data.cpp deleted file mode 100644 index 3342672..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/stats_label_data.cpp +++ /dev/null @@ -1,147 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "utils.h" - -#ifndef _WINDOWS -#include -#include -#include -#include -#else -#include -#endif -namespace po = boost::program_options; - -void stats_analysis(const std::string labels_file, std::string univeral_label, uint32_t density = 10) -{ - std::string token, line; - std::ifstream labels_stream(labels_file); - std::unordered_map label_counts; - std::string label_with_max_points; - uint32_t max_points = 0; - long long sum = 0; - long long point_cnt = 0; - float avg_labels_per_pt, mean_label_size; - - std::vector labels_per_point; - uint32_t dense_pts = 0; - if (labels_stream.is_open()) - { - while (getline(labels_stream, line)) - { - point_cnt++; - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - while (getline(iss, token, ',')) - { - lbl_cnt++; - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (label_counts.find(token) == label_counts.end()) - label_counts[token] = 0; - label_counts[token]++; - } - if (lbl_cnt >= density) - { - dense_pts++; - } - labels_per_point.emplace_back(lbl_cnt); - } - } - - std::cout << "fraction of dense points with >= " << density - << " labels = " << (float)dense_pts / (float)labels_per_point.size() << std::endl; - std::sort(labels_per_point.begin(), labels_per_point.end()); - - std::vector> label_count_vec; - - for (auto it = label_counts.begin(); it != label_counts.end(); it++) - { - auto &lbl = *it; - label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second)); - if (lbl.second > max_points) - { - max_points = lbl.second; - label_with_max_points = lbl.first; - } - sum += lbl.second; - } - - sort(label_count_vec.begin(), label_count_vec.end(), - [](const std::pair &lhs, const std::pair &rhs) { - return lhs.second < rhs.second; - }); - - for (float p = 0; p < 1; p += 0.05) - { - std::cout << "Percentile " << (100 * p) << "\t" << label_count_vec[(size_t)(p * label_count_vec.size())].first - << " with count=" << label_count_vec[(size_t)(p * label_count_vec.size())].second << std::endl; - } - - std::cout << "Most common label " - << "\t" << label_count_vec[label_count_vec.size() - 1].first - << " with count=" << label_count_vec[label_count_vec.size() - 1].second << std::endl; - if (label_count_vec.size() > 1) - std::cout << "Second common label " - << "\t" << label_count_vec[label_count_vec.size() - 2].first - << " with count=" << label_count_vec[label_count_vec.size() - 2].second << std::endl; - if (label_count_vec.size() > 2) - std::cout << "Third common label " - << "\t" << label_count_vec[label_count_vec.size() - 3].first - << " with count=" << label_count_vec[label_count_vec.size() - 3].second << std::endl; - avg_labels_per_pt = sum / (float)point_cnt; - mean_label_size = sum / (float)label_counts.size(); - std::cout << "Total number of points = " << point_cnt << ", number of labels = " << label_counts.size() - << std::endl; - std::cout << "Average number of labels per point = " << avg_labels_per_pt << std::endl; - std::cout << "Mean label size excluding 0 = " << mean_label_size << std::endl; - std::cout << "Most popular label is " << label_with_max_points << " with " << max_points << " pts" << std::endl; -} - -int main(int argc, char **argv) -{ - std::string labels_file, universal_label; - uint32_t density; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("labels_file", po::value(&labels_file)->required(), - "path to labels data file."); - desc.add_options()("universal_label", po::value(&universal_label)->required(), - "Universal label used in labels file."); - desc.add_options()("density", po::value(&density)->default_value(1), - "Number of labels each point in labels file, defaults to 1"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &e) - { - std::cerr << e.what() << '\n'; - return -1; - } - stats_analysis(labels_file, universal_label, density); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/tsv_to_bin.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/tsv_to_bin.cpp deleted file mode 100644 index c590a8f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/tsv_to_bin.cpp +++ /dev/null @@ -1,121 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -void block_convert_float(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new float[npts * (ndims + 1)]; - - auto cursor = read_buf; - float val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = val; - cursor++; - } - } - writer.write((char *)read_buf, npts * ndims * sizeof(float)); - delete[] read_buf; -} - -void block_convert_int8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new int8_t[npts * (ndims + 1)]; - - auto cursor = read_buf; - int val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = (int8_t)val; - cursor++; - } - } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; -} - -void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new uint8_t[npts * (ndims + 1)]; - - auto cursor = read_buf; - int val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = (uint8_t)val; - cursor++; - } - } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; -} - -int main(int argc, char **argv) -{ - if (argc != 6) - { - std::cout << argv[0] - << " input_filename.tsv output_filename.bin " - "dim num_pts>" - << std::endl; - exit(-1); - } - - if (std::string(argv[1]) != std::string("float") && std::string(argv[1]) != std::string("int8") && - std::string(argv[1]) != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - } - - size_t ndims = atoi(argv[4]); - size_t npts = atoi(argv[5]); - - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - // size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - reader.seekg(0, std::ios::beg); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (std::string(argv[1]) == std::string("float")) - { - block_convert_float(reader, writer, cblk_size, ndims); - } - else if (std::string(argv[1]) == std::string("int8")) - { - block_convert_int8(reader, writer, cblk_size, ndims); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - block_convert_uint8(reader, writer, cblk_size, ndims); - } - std::cout << "Block #" << i << " written" << std::endl; - } - - reader.close(); - writer.close(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint32_to_uint8.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint32_to_uint8.cpp deleted file mode 100644 index 87b6fb8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint32_to_uint8.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; - exit(-1); - } - - uint32_t *input; - size_t npts, nd; - diskann::load_bin(argv[1], input, npts, nd); - uint8_t *output = new uint8_t[npts * nd]; - diskann::convert_types(input, output, npts, nd); - diskann::save_bin(argv[2], output, npts, nd); - delete[] output; - delete[] input; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint8_to_float.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint8_to_float.cpp deleted file mode 100644 index 6415b7c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/uint8_to_float.cpp +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "utils.h" - -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_uint8_bin output_float_bin" << std::endl; - exit(-1); - } - - uint8_t *input; - size_t npts, nd; - diskann::load_bin(argv[1], input, npts, nd); - float *output = new float[npts * nd]; - diskann::convert_types(input, output, npts, nd); - diskann::save_bin(argv[2], output, npts, nd); - delete[] output; - delete[] input; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/vector_analysis.cpp b/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/vector_analysis.cpp deleted file mode 100644 index 009df6d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/apps/utils/vector_analysis.cpp +++ /dev/null @@ -1,163 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "partition.h" -#include "utils.h" - -template int analyze_norm(std::string base_file) -{ - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); -#pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - norms[i] = std::sqrt(norms[i]); - } - std::sort(norms.begin(), norms.end()); - for (int p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " << norms[(uint64_t)(std::floor((p / 100.0) * npts))] << std::endl; - std::cout << "percentile 100" - << ": " << norms[npts - 1] << std::endl; - delete[] data; - return 0; -} - -template int normalize_base(std::string base_file, std::string out_file) -{ - std::cout << "Normalizing base" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - // std::vector norms(npts, 0); -#pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - float pt_norm = 0; - for (size_t d = 0; d < ndims; d++) - pt_norm += data[i * ndims + d] * data[i * ndims + d]; - pt_norm = std::sqrt(pt_norm); - for (size_t d = 0; d < ndims; d++) - data[i * ndims + d] = static_cast(data[i * ndims + d] / pt_norm); - } - diskann::save_bin(out_file, data, npts, ndims); - delete[] data; - return 0; -} - -template int augment_base(std::string base_file, std::string out_file, bool prep_base = true) -{ - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); - float max_norm = 0; -#pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - max_norm = norms[i] > max_norm ? norms[i] : max_norm; - } - // std::sort(norms.begin(), norms.end()); - max_norm = std::sqrt(max_norm); - std::cout << "Max norm: " << max_norm << std::endl; - T *new_data; - size_t newdims = ndims + 1; - new_data = new T[npts * newdims]; - for (size_t i = 0; i < npts; i++) - { - if (prep_base) - { - for (size_t j = 0; j < ndims; j++) - { - new_data[i * newdims + j] = static_cast(data[i * ndims + j] / max_norm); - } - float diff = 1 - (norms[i] / (max_norm * max_norm)); - diff = diff <= 0 ? 0 : std::sqrt(diff); - new_data[i * newdims + ndims] = static_cast(diff); - if (diff <= 0) - { - std::cout << i << " has large max norm, investigate if needed. diff = " << diff << std::endl; - } - } - else - { - for (size_t j = 0; j < ndims; j++) - { - new_data[i * newdims + j] = static_cast(data[i * ndims + j] / std::sqrt(norms[i])); - } - new_data[i * newdims + ndims] = 0; - } - } - diskann::save_bin(out_file, new_data, npts, newdims); - delete[] new_data; - delete[] data; - return 0; -} - -template int aux_main(char **argv) -{ - std::string base_file(argv[2]); - uint32_t option = atoi(argv[3]); - if (option == 1) - analyze_norm(base_file); - else if (option == 2) - augment_base(base_file, std::string(argv[4]), true); - else if (option == 3) - augment_base(base_file, std::string(argv[4]), false); - else if (option == 4) - normalize_base(base_file, std::string(argv[4])); - return 0; -} - -int main(int argc, char **argv) -{ - if (argc < 4) - { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "[option: 1-norm analysis, 2-prep_base_for_mip, " - "3-prep_query_for_mip, 4-normalize-vecs] [out_file for " - "options 2/3/4]" - << std::endl; - exit(-1); - } - - if (std::string(argv[1]) == std::string("float")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("int8")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - aux_main(argv); - } - else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/clang-format.cmake b/packages/leann-backend-diskann/third_party/DiskANN/clang-format.cmake deleted file mode 100644 index 19bb3a8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/clang-format.cmake +++ /dev/null @@ -1,22 +0,0 @@ -if (NOT MSVC) - message(STATUS "Setting up `make format` and `make checkformat`") - # additional target to perform clang-format run, requires clang-format - # get all project files - file(GLOB_RECURSE ALL_SOURCE_FILES include/*.h include/*.hpp python/src/*.cpp src/*.cpp src/*.hpp apps/*.cpp apps/*.hpp) - - message(status ${ALL_SOURCE_FILES}) - - add_custom_target( - format - COMMAND /usr/bin/clang-format - -i - ${ALL_SOURCE_FILES} - ) - add_custom_target( - checkformat - COMMAND /usr/bin/clang-format - --Werror - --dry-run - ${ALL_SOURCE_FILES} - ) -endif() diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_data_store.h b/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_data_store.h deleted file mode 100644 index 89856f1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_data_store.h +++ /dev/null @@ -1,127 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include "types.h" -#include "windows_customizations.h" -#include "distance.h" - -namespace diskann -{ - -template class AbstractScratch; - -template class AbstractDataStore -{ - public: - AbstractDataStore(const location_t capacity, const size_t dim); - - virtual ~AbstractDataStore() = default; - - // Return number of points returned - virtual location_t load(const std::string &filename) = 0; - - // Why does store take num_pts? Since store only has capacity, but we allow - // resizing we can end up in a situation where the store has spare capacity. - // To optimize disk utilization, we pass the number of points that are "true" - // points, so that the store can discard the empty locations before saving. - virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - - DISKANN_DLLEXPORT virtual location_t capacity() const; - - DISKANN_DLLEXPORT virtual size_t get_dims() const; - - // Implementers can choose to return _dim if they are not - // concerned about memory alignment. - // Some distance metrics (like l2) need data vectors to be aligned, so we - // align the dimension by padding zeros. - virtual size_t get_aligned_dim() const = 0; - - // populate the store with vectors (either from a pointer or bin file), - // potentially after pre-processing the vectors if the metric deems so - // e.g., normalizing vectors for cosine distance over floating-point vectors - // useful for bulk or static index building. - virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; - virtual void populate_data(const std::string &filename, const size_t offset) = 0; - - // save the first num_pts many vectors back to bin file - // note: cannot undo the pre-processing done in populate data - virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) = 0; - - // Returns the updated capacity of the datastore. Clients should check - // if resize actually changed the capacity to new_num_points before - // proceeding with operations. See the code below: - // auto new_capcity = data_store->resize(new_num_points); - // if ( new_capacity >= new_num_points) { - // //PROCEED - // else - // //ERROR. - virtual location_t resize(const location_t new_num_points); - - // operations on vectors - // like populate_data function, but over one vector at a time useful for - // streaming setting - virtual void get_vector(const location_t i, data_t *dest) const = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; - virtual void prefetch_vector(const location_t loc) = 0; - - // internal shuffle operations to move around vectors - // will bulk-move all the vectors in [old_start_loc, old_start_loc + - // num_points) to [new_start_loc, new_start_loc + num_points) and set the old - // positions to zero vectors. - virtual void move_vectors(const location_t old_start_loc, const location_t new_start_loc, - const location_t num_points) = 0; - - // same as above, without resetting the vectors in [from_loc, from_loc + - // num_points) to zero - virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - - // With the PQ Data Store PR, we have also changed iterate_to_fixed_point to NOT take the query - // from the scratch object. Therefore every data store has to implement preprocess_query which - // at the least will be to copy the query into the scratch object. So making this pure virtual. - virtual void preprocess_query(const data_t *aligned_query, - AbstractScratch *query_scratch = nullptr) const = 0; - // distance functions. - virtual float get_distance(const data_t *query, const location_t loc) const = 0; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances, AbstractScratch *scratch_space = nullptr) const = 0; - // Specific overload for index.cpp. - virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const = 0; - virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - - // stats of the data stored in store - // Returns the point in the dataset that is closest to the mean of all points - // in the dataset - virtual location_t calculate_medoid() const = 0; - - // REFACTOR PQ TODO: Each data store knows about its distance function, so this is - // redundant. However, we don't have an OptmizedDataStore yet, and to preserve code - // compability, we are exposing this function. - virtual Distance *get_dist_fn() const = 0; - - // search helpers - // if the base data is aligned per the request of the metric, this will tell - // how to align the query vector in a consistent manner - virtual size_t get_alignment_factor() const = 0; - - protected: - // Expand the datastore to new_num_points. Returns the new capacity created, - // which should be == new_num_points in the normal case. Implementers can also - // return _capacity to indicate that there are not implementing this method. - virtual location_t expand(const location_t new_num_points) = 0; - - // Shrink the datastore to new_num_points. It is NOT an error if shrink - // doesn't reduce the capacity so callers need to check this correctly. See - // also for "default" implementation - virtual location_t shrink(const location_t new_num_points) = 0; - - location_t _capacity; - size_t _dim; -}; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_graph_store.h b/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_graph_store.h deleted file mode 100644 index 4d6906c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_graph_store.h +++ /dev/null @@ -1,68 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include "types.h" - -namespace diskann -{ - -class AbstractGraphStore -{ - public: - AbstractGraphStore(const size_t total_pts, const size_t reserve_graph_degree) - : _capacity(total_pts), _reserve_graph_degree(reserve_graph_degree) - { - } - - virtual ~AbstractGraphStore() = default; - - // returns tuple of - virtual std::tuple load(const std::string &index_path_prefix, - const size_t num_points) = 0; - virtual int store(const std::string &index_path_prefix, const size_t num_points, const size_t num_fz_points, - const uint32_t start) = 0; - - // not synchronised, user should use lock when necvessary. - virtual const std::vector &get_neighbours(const location_t i) const = 0; - virtual void add_neighbour(const location_t i, location_t neighbour_id) = 0; - virtual void clear_neighbours(const location_t i) = 0; - virtual void swap_neighbours(const location_t a, location_t b) = 0; - - virtual void set_neighbours(const location_t i, std::vector &neighbours) = 0; - - virtual size_t resize_graph(const size_t new_size) = 0; - virtual void clear_graph() = 0; - - virtual uint32_t get_max_observed_degree() = 0; - - // set during load - virtual size_t get_max_range_of_graph() = 0; - - // Total internal points _max_points + _num_frozen_points - size_t get_total_points() - { - return _capacity; - } - - protected: - // Internal function, changes total points when resize_graph is called. - void set_total_points(size_t new_capacity) - { - _capacity = new_capacity; - } - - size_t get_reserve_graph_degree() - { - return _reserve_graph_degree; - } - - private: - size_t _capacity; - size_t _reserve_graph_degree; -}; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_index.h b/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_index.h deleted file mode 100644 index 059866f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_index.h +++ /dev/null @@ -1,129 +0,0 @@ -#pragma once -#include "distance.h" -#include "parameters.h" -#include "utils.h" -#include "types.h" -#include "index_config.h" -#include "index_build_params.h" -#include - -namespace diskann -{ -struct consolidation_report -{ - enum status_code - { - SUCCESS = 0, - FAIL = 1, - LOCK_FAIL = 2, - INCONSISTENT_COUNT_ERROR = 3 - }; - status_code _status; - size_t _active_points, _max_points, _empty_slots, _slots_released, _delete_set_size, _num_calls_to_process_delete; - double _time; - - consolidation_report(status_code status, size_t active_points, size_t max_points, size_t empty_slots, - size_t slots_released, size_t delete_set_size, size_t num_calls_to_process_delete, - double time_secs) - : _status(status), _active_points(active_points), _max_points(max_points), _empty_slots(empty_slots), - _slots_released(slots_released), _delete_set_size(delete_set_size), - _num_calls_to_process_delete(num_calls_to_process_delete), _time(time_secs) - { - } -}; - -/* A templated independent class for intercation with Index. Uses Type Erasure to add virtual implemetation of methods -that can take any type(using std::any) and Provides a clean API that can be inherited by different type of Index. -*/ -class AbstractIndex -{ - public: - AbstractIndex() = default; - virtual ~AbstractIndex() = default; - - virtual void build(const std::string &data_file, const size_t num_points_to_load, - IndexFilterParams &build_params) = 0; - - template - void build(const data_type *data, const size_t num_points_to_load, const std::vector &tags); - - virtual void save(const char *filename, bool compact_before_save = false) = 0; - -#ifdef EXEC_ENV_OLS - virtual void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) = 0; -#else - virtual void load(const char *index_file, uint32_t num_threads, uint32_t search_l) = 0; -#endif - - // For FastL2 search on optimized layout - template - void search_with_optimized_layout(const data_type *query, size_t K, size_t L, uint32_t *indices); - - // Initialize space for res_vectors before calling. - template - size_t search_with_tags(const data_type *query, const uint64_t K, const uint32_t L, tag_type *tags, - float *distances, std::vector &res_vectors, bool use_filters = false, - const std::string filter_label = ""); - - // Added search overload that takes L as parameter, so that we - // can customize L on a per-query basis without tampering with "Parameters" - // IDtype is either uint32_t or uint64_t - template - std::pair search(const data_type *query, const size_t K, const uint32_t L, IDType *indices, - float *distances = nullptr); - - // Filter support search - // IndexType is either uint32_t or uint64_t - template - std::pair search_with_filters(const DataType &query, const std::string &raw_label, - const size_t K, const uint32_t L, IndexType *indices, - float *distances); - - // insert points with labels, labels should be present for filtered index - template - int insert_point(const data_type *point, const tag_type tag, const std::vector &labels); - - // insert point for unfiltered index build. do not use with filtered index - template int insert_point(const data_type *point, const tag_type tag); - - // delete point with tag, or return -1 if point can not be deleted - template int lazy_delete(const tag_type &tag); - - // batch delete tags and populates failed tags if unabke to delete given tags. - template - void lazy_delete(const std::vector &tags, std::vector &failed_tags); - - template void get_active_tags(tsl::robin_set &active_tags); - - template void set_start_points_at_random(data_type radius, uint32_t random_seed = 0); - - virtual consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters) = 0; - - virtual void optimize_index_layout() = 0; - - // memory should be allocated for vec before calling this function - template int get_vector_by_tag(tag_type &tag, data_type *vec); - - template void set_universal_label(const label_type universal_label); - - private: - virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) = 0; - virtual std::pair _search(const DataType &query, const size_t K, const uint32_t L, - std::any &indices, float *distances = nullptr) = 0; - virtual std::pair _search_with_filters(const DataType &query, const std::string &filter_label, - const size_t K, const uint32_t L, std::any &indices, - float *distances) = 0; - virtual int _insert_point(const DataType &data_point, const TagType tag, Labelvector &labels) = 0; - virtual int _insert_point(const DataType &data_point, const TagType tag) = 0; - virtual int _lazy_delete(const TagType &tag) = 0; - virtual void _lazy_delete(TagVector &tags, TagVector &failed_tags) = 0; - virtual void _get_active_tags(TagRobinSet &active_tags) = 0; - virtual void _set_start_points_at_random(DataType radius, uint32_t random_seed = 0) = 0; - virtual int _get_vector_by_tag(TagType &tag, DataType &vec) = 0; - virtual size_t _search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags, - float *distances, DataVector &res_vectors, bool use_filters = false, - const std::string filter_label = "") = 0; - virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) = 0; - virtual void _set_universal_label(const LabelType universal_label) = 0; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_scratch.h b/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_scratch.h deleted file mode 100644 index b42a836..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/abstract_scratch.h +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once -namespace diskann -{ - -template class PQScratch; - -// By somewhat more than a coincidence, it seems that both InMemQueryScratch -// and SSDQueryScratch have the aligned query and PQScratch objects. So we -// can put them in a neat hierarchy and keep PQScratch as a standalone class. -template class AbstractScratch -{ - public: - AbstractScratch() = default; - // This class does not take any responsibilty for memory management of - // its members. It is the responsibility of the derived classes to do so. - virtual ~AbstractScratch() = default; - - // Scratch objects should not be copied - AbstractScratch(const AbstractScratch &) = delete; - AbstractScratch &operator=(const AbstractScratch &) = delete; - - data_t *aligned_query_T() - { - return _aligned_query_T; - } - PQScratch *pq_scratch() - { - return _pq_scratch; - } - - protected: - data_t *_aligned_query_T = nullptr; - PQScratch *_pq_scratch = nullptr; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/aligned_file_reader.h b/packages/leann-backend-diskann/third_party/DiskANN/include/aligned_file_reader.h deleted file mode 100644 index 2e2716a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/aligned_file_reader.h +++ /dev/null @@ -1,138 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#define MAX_IO_DEPTH 128 - -#include -#include - -#ifdef __linux__ -#include -#include -#include -#include -typedef io_context_t IOContext; -#elif __APPLE__ -#include -#include -#include - -struct IOContext -{ - int fd; - dispatch_io_t channel; - dispatch_queue_t queue; - dispatch_group_t grp; -}; -#elif _WINDOWS -#include -#include -#include - -#ifndef USE_BING_INFRA -struct IOContext -{ - HANDLE fhandle = NULL; - HANDLE iocp = NULL; - std::vector reqs; -}; -#else -#include "IDiskPriorityIO.h" -#include -// TODO: Caller code is very callous about copying IOContext objects -// all over the place. MUST verify that it won't cause leaks/logical -// errors. -// Because of such callous copying, we have to use ptr->atomic instead -// of atomic, as atomic is not copyable. -struct IOContext -{ - enum Status - { - READ_WAIT = 0, - READ_SUCCESS, - READ_FAILED, - PROCESS_COMPLETE - }; - - std::shared_ptr m_pDiskIO = nullptr; - std::shared_ptr> m_pRequests; - std::shared_ptr> m_pRequestsStatus; - - // waitonaddress on this memory to wait for IO completion signal - // reader should signal this memory after IO completion - // TODO: WindowsAlignedFileReader can be modified to take advantage of this - // and can largely share code with the file reader for Bing. - mutable volatile long m_completeCount = 0; - - IOContext() - : m_pRequestsStatus(new std::vector()), m_pRequests(new std::vector()) - { - (*m_pRequestsStatus).reserve(MAX_IO_DEPTH); - (*m_pRequests).reserve(MAX_IO_DEPTH); - } -}; -#endif - -#endif - -#include -#include -#include -#include "tsl/robin_map.h" -#include "utils.h" - -// NOTE :: all 3 fields must be 512-aligned -struct AlignedRead -{ - uint64_t offset; // where to read from - uint64_t len; // how much to read - void *buf; // where to read into - - AlignedRead() : offset(0), len(0), buf(nullptr) - { - } - - AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) - { - assert(IS_512_ALIGNED(offset)); - assert(IS_512_ALIGNED(len)); - assert(IS_512_ALIGNED(buf)); - // assert(malloc_usable_size(buf) >= len); - } -}; - -class AlignedFileReader -{ - protected: - tsl::robin_map ctx_map; - std::mutex ctx_mut; - - public: - // returns the thread-specific context - // returns (io_context_t)(-1) if thread is not registered - virtual IOContext &get_ctx() = 0; - - virtual ~AlignedFileReader(){}; - - // register thread-id for a context - virtual void register_thread() = 0; - // de-register thread-id for a context - virtual void deregister_thread() = 0; - virtual void deregister_all_threads() = 0; - - // Open & close ops - // Blocking calls - virtual void open(const std::string &fname) = 0; - virtual void close() = 0; - - // process batch of aligned requests in parallel - // NOTE :: blocking call - virtual void read(std::vector &read_reqs, IOContext &ctx, bool async = false) = 0; - -#ifdef USE_BING_INFRA - // wait for completion of one request in a batch of requests - virtual void wait(IOContext &ctx, int &completedIndex) = 0; -#endif -}; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/ann_exception.h b/packages/leann-backend-diskann/third_party/DiskANN/include/ann_exception.h deleted file mode 100644 index 55f069b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/ann_exception.h +++ /dev/null @@ -1,35 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include "windows_customizations.h" -#include - -#ifndef _WINDOWS -#define __FUNCSIG__ __PRETTY_FUNCTION__ -#endif - -namespace diskann -{ - -class ANNException : public std::runtime_error -{ - public: - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode); - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum); - - private: - int _errorCode; -}; - -class FileException : public ANNException -{ - public: - DISKANN_DLLEXPORT FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum); -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/any_wrappers.h b/packages/leann-backend-diskann/third_party/DiskANN/include/any_wrappers.h deleted file mode 100644 index da9005c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/any_wrappers.h +++ /dev/null @@ -1,53 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include -#include "tsl/robin_set.h" - -namespace AnyWrapper -{ - -/* - * Base Struct to hold refrence to the data. - * Note: No memory mamagement, caller need to keep object alive. - */ -struct AnyReference -{ - template AnyReference(Ty &reference) : _data(&reference) - { - } - - template Ty &get() - { - auto ptr = std::any_cast(_data); - return *ptr; - } - - private: - std::any _data; -}; -struct AnyRobinSet : public AnyReference -{ - template AnyRobinSet(const tsl::robin_set &robin_set) : AnyReference(robin_set) - { - } - template AnyRobinSet(tsl::robin_set &robin_set) : AnyReference(robin_set) - { - } -}; - -struct AnyVector : public AnyReference -{ - template AnyVector(const std::vector &vector) : AnyReference(vector) - { - } - template AnyVector(std::vector &vector) : AnyReference(vector) - { - } -}; -} // namespace AnyWrapper diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/apple_aligned_file_reader.h b/packages/leann-backend-diskann/third_party/DiskANN/include/apple_aligned_file_reader.h deleted file mode 100644 index 160e1ea..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/apple_aligned_file_reader.h +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once -#ifdef __APPLE__ -#include "aligned_file_reader.h" - -class AppleAlignedFileReader : public AlignedFileReader -{ - private: - uint64_t file_sz; - FileHandle file_desc; - - public: - AppleAlignedFileReader(); - ~AppleAlignedFileReader(); - - IOContext &get_ctx(); - - void register_thread(); - void deregister_thread(); - void deregister_all_threads(); - - void open(const std::string &fname); - void close(); - - void read(std::vector &read_reqs, IOContext &ctx, bool async = false); -}; -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/boost_dynamic_bitset_fwd.h b/packages/leann-backend-diskann/third_party/DiskANN/include/boost_dynamic_bitset_fwd.h deleted file mode 100644 index 5aebb2b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/boost_dynamic_bitset_fwd.h +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -namespace boost -{ -#ifndef BOOST_DYNAMIC_BITSET_FWD_HPP -template > class dynamic_bitset; -#endif -} // namespace boost diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/cached_io.h b/packages/leann-backend-diskann/third_party/DiskANN/include/cached_io.h deleted file mode 100644 index daef2f2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/cached_io.h +++ /dev/null @@ -1,217 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include - -#include "logger.h" -#include "ann_exception.h" - -// sequential cached reads -class cached_ifstream -{ - public: - cached_ifstream() - { - } - cached_ifstream(const std::string &filename, uint64_t cacheSize) : cache_size(cacheSize), cur_off(0) - { - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - this->open(filename, cache_size); - } - ~cached_ifstream() - { - delete[] cache_buf; - reader.close(); - } - - void open(const std::string &filename, uint64_t cacheSize) - { - this->cur_off = 0; - - try - { - reader.open(filename, std::ios::binary | std::ios::ate); - fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - assert(reader.is_open()); - assert(cacheSize > 0); - cacheSize = (std::min)(cacheSize, fsize); - this->cache_size = cacheSize; - cache_buf = new char[cacheSize]; - reader.read(cache_buf, cacheSize); - diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize << ", cache_size: " << cacheSize - << std::endl; - } - catch (std::system_error &e) - { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - - size_t get_file_size() - { - return fsize; - } - - void read(char *read_buf, uint64_t n_bytes) - { - assert(cache_buf != nullptr); - assert(read_buf != nullptr); - - if (n_bytes <= (cache_size - cur_off)) - { - // case 1: cache contains all data - memcpy(read_buf, cache_buf + cur_off, n_bytes); - cur_off += n_bytes; - } - else - { - // case 2: cache contains some data - uint64_t cached_bytes = cache_size - cur_off; - if (n_bytes - cached_bytes > fsize - reader.tellg()) - { - std::stringstream stream; - stream << "Reading beyond end of file" << std::endl; - stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes << " fsize: " << fsize - << " current pos:" << reader.tellg() << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - memcpy(read_buf, cache_buf + cur_off, cached_bytes); - - // go to disk and fetch more data - reader.read(read_buf + cached_bytes, n_bytes - cached_bytes); - // reset cur off - cur_off = cache_size; - - uint64_t size_left = fsize - reader.tellg(); - - if (size_left >= cache_size) - { - reader.read(cache_buf, cache_size); - cur_off = 0; - } - // note that if size_left < cache_size, then cur_off = cache_size, - // so subsequent reads will all be directly from file - } - } - - private: - // underlying ifstream - std::ifstream reader; - // # bytes to cache in one shot read - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - // file size - uint64_t fsize = 0; -}; - -// sequential cached writes -class cached_ofstream -{ - public: - cached_ofstream(const std::string &filename, uint64_t cache_size) : cache_size(cache_size), cur_off(0) - { - writer.exceptions(std::ifstream::failbit | std::ifstream::badbit); - try - { - writer.open(filename, std::ios::binary); - assert(writer.is_open()); - assert(cache_size > 0); - cache_buf = new char[cache_size]; - diskann::cout << "Opened: " << filename.c_str() << ", cache_size: " << cache_size << std::endl; - } - catch (std::system_error &e) - { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - - ~cached_ofstream() - { - this->close(); - } - - void close() - { - // dump any remaining data in memory - if (cur_off > 0) - { - this->flush_cache(); - } - - if (cache_buf != nullptr) - { - delete[] cache_buf; - cache_buf = nullptr; - } - - if (writer.is_open()) - writer.close(); - diskann::cout << "Finished writing " << fsize << "B" << std::endl; - } - - size_t get_file_size() - { - return fsize; - } - // writes n_bytes from write_buf to the underlying ofstream/cache - void write(char *write_buf, uint64_t n_bytes) - { - assert(cache_buf != nullptr); - if (n_bytes <= (cache_size - cur_off)) - { - // case 1: cache can take all data - memcpy(cache_buf + cur_off, write_buf, n_bytes); - cur_off += n_bytes; - } - else - { - // case 2: cache cant take all data - // go to disk and write existing cache data - writer.write(cache_buf, cur_off); - fsize += cur_off; - // write the new data to disk - writer.write(write_buf, n_bytes); - fsize += n_bytes; - // memset all cache data and reset cur_off - memset(cache_buf, 0, cache_size); - cur_off = 0; - } - } - - void flush_cache() - { - assert(cache_buf != nullptr); - writer.write(cache_buf, cur_off); - fsize += cur_off; - memset(cache_buf, 0, cache_size); - cur_off = 0; - } - - void reset() - { - flush_cache(); - writer.seekp(0); - } - - private: - // underlying ofstream - std::ofstream writer; - // # bytes to cache for one shot write - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - - // file size - uint64_t fsize = 0; -}; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/common_includes.h b/packages/leann-backend-diskann/third_party/DiskANN/include/common_includes.h deleted file mode 100644 index e1a51bd..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/common_includes.h +++ /dev/null @@ -1,27 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/concurrent_queue.h b/packages/leann-backend-diskann/third_party/DiskANN/include/concurrent_queue.h deleted file mode 100644 index 1e57bbf..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/concurrent_queue.h +++ /dev/null @@ -1,132 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include -#include -#include -#include -#include - -namespace diskann -{ - -template class ConcurrentQueue -{ - typedef std::chrono::microseconds chrono_us_t; - typedef std::unique_lock mutex_locker; - - std::queue q; - std::mutex mut; - std::mutex push_mut; - std::mutex pop_mut; - std::condition_variable push_cv; - std::condition_variable pop_cv; - T null_T; - - public: - ConcurrentQueue() - { - } - - ConcurrentQueue(T nullT) - { - this->null_T = nullT; - } - - ~ConcurrentQueue() - { - this->push_cv.notify_all(); - this->pop_cv.notify_all(); - } - - // queue stats - uint64_t size() - { - mutex_locker lk(this->mut); - uint64_t ret = q.size(); - lk.unlock(); - return ret; - } - - bool empty() - { - return (this->size() == 0); - } - - // PUSH BACK - void push(T &new_val) - { - mutex_locker lk(this->mut); - this->q.push(new_val); - lk.unlock(); - } - - template void insert(Iterator iter_begin, Iterator iter_end) - { - mutex_locker lk(this->mut); - for (Iterator it = iter_begin; it != iter_end; it++) - { - this->q.push(*it); - } - lk.unlock(); - } - - // POP FRONT - T pop() - { - mutex_locker lk(this->mut); - if (this->q.empty()) - { - lk.unlock(); - return this->null_T; - } - else - { - T ret = this->q.front(); - this->q.pop(); - // diskann::cout << "thread_id: " << std::this_thread::get_id() << - // ", ctx: " - // << ret.ctx << "\n"; - lk.unlock(); - return ret; - } - } - - // register for notifications - void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10}) - { - mutex_locker lk(this->push_mut); - this->push_cv.wait_for(lk, wait_time); - lk.unlock(); - } - - void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10}) - { - mutex_locker lk(this->pop_mut); - this->pop_cv.wait_for(lk, wait_time); - lk.unlock(); - } - - // just notify functions - void push_notify_one() - { - this->push_cv.notify_one(); - } - void push_notify_all() - { - this->push_cv.notify_all(); - } - void pop_notify_one() - { - this->pop_cv.notify_one(); - } - void pop_notify_all() - { - this->pop_cv.notify_all(); - } -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/cosine_similarity.h b/packages/leann-backend-diskann/third_party/DiskANN/include/cosine_similarity.h deleted file mode 100644 index 539a8b0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/cosine_similarity.h +++ /dev/null @@ -1,285 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#ifndef __APPLE__ -#include -#include -#include -#include "simd_utils.h" -#endif - -extern bool Avx2SupportedCPU; - -#ifdef _WINDOWS -// SIMD implementation of Cosine similarity. Taken from hnsw library. - -/** - * Non-metric Space Library - * - * Authors: Bilegsaikhan Naidan (https://github.com/bileg), Leonid Boytsov - * (http://boytsov.info). With contributions from Lawrence Cayton - * (http://lcayton.com/) and others. - * - * For the complete list of contributors and further details see: - * https://github.com/searchivarius/NonMetricSpaceLib - * - * Copyright (c) 2014 - * - * This code is released under the - * Apache License Version 2.0 http://www.apache.org/licenses/. - * - */ - -namespace diskann -{ - -using namespace std; - -#define PORTABLE_ALIGN16 __declspec(align(16)) - -static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, uint32_t qty) -{ - if (Avx2SupportedCPU) - { - __m256 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm256_setzero_ps(); - while (qty >= 32) - { - __m256i rx = _mm256_load_si256((__m256i *)pVect1), ry = _mm256_load_si256((__m256i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry)); - pVect1 += 32; - pVect2 += 32; - qty -= 32; - } - while (qty > 0) - { - __m128i rx = _mm_load_si128((__m128i *)pVect1), ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; - } - cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos); - p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len); - float denominator = max(numeric_limits::min() * 2, sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) * - sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4])); - float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator; - - return max(float(-1), min(float(1), cosine)); - } - - __m128 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm_setzero_ps(); - __m128i rx, ry; - while (qty >= 16) - { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry)); - pVect1 += 16; - pVect2 += 16; - qty -= 16; - } - while (qty > 0) - { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; - } - cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos); - p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len); - float norm1 = p1Len.m128_f32[0]; - float norm2 = p2Len.m128_f32[0]; - - static const float eps = numeric_limits::min() * 2; - - if (norm1 < eps) - { /* - * This shouldn't normally happen for this space, but - * if it does, we don't want to get NANs - */ - if (norm2 < eps) - { - return 1; - } - return 0; - } - /* - * Sometimes due to rounding errors, we get values > 1 or < -1. - * This throws off other functions that use scalar product, e.g., acos - */ - return max(float(-1), min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2))); -} - -static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, uint32_t qty) -{ - // Didn't get significant performance gain compared with 128bit version. - static const float eps = numeric_limits::min() * 2; - - if (Avx2SupportedCPU) - { - uint32_t qty8 = qty / 8; - - const float *pEnd1 = pVect1 + 8 * qty8; - const float *pEnd2 = pVect1 + qty; - - __m256 v1, v2; - __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0); - __m256 sum_square1 = sum_prod; - __m256 sum_square2 = sum_prod; - - while (pVect1 < pEnd1) - { - v1 = _mm256_loadu_ps(pVect1); - pVect1 += 8; - v2 = _mm256_loadu_ps(pVect2); - pVect2 += 8; - sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2)); - sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1)); - sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2)); - } - - float PORTABLE_ALIGN16 TmpResProd[8]; - float PORTABLE_ALIGN16 TmpResSquare1[8]; - float PORTABLE_ALIGN16 TmpResSquare2[8]; - - _mm256_store_ps(TmpResProd, sum_prod); - _mm256_store_ps(TmpResSquare1, sum_square1); - _mm256_store_ps(TmpResSquare2, sum_square2); - - float sum = 0.0f; - float norm1 = 0.0f; - float norm2 = 0.0f; - for (uint32_t i = 0; i < 8; ++i) - { - sum += TmpResProd[i]; - norm1 += TmpResSquare1[i]; - norm2 += TmpResSquare2[i]; - } - - while (pVect1 < pEnd2) - { - sum += (*pVect1) * (*pVect2); - norm1 += (*pVect1) * (*pVect1); - norm2 += (*pVect2) * (*pVect2); - - ++pVect1; - ++pVect2; - } - - if (norm1 < eps) - { - return norm2 < eps ? 1.0f : 0.0f; - } - - return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); - } - - __m128 v1, v2; - __m128 sum_prod = _mm_set1_ps(0); - __m128 sum_square1 = sum_prod; - __m128 sum_square2 = sum_prod; - - while (qty >= 4) - { - v1 = _mm_loadu_ps(pVect1); - pVect1 += 4; - v2 = _mm_loadu_ps(pVect2); - pVect2 += 4; - sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2)); - sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1)); - sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2)); - - qty -= 4; - } - - float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + sum_prod.m128_f32[2] + sum_prod.m128_f32[3]; - float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + sum_square1.m128_f32[2] + sum_square1.m128_f32[3]; - float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + sum_square2.m128_f32[2] + sum_square2.m128_f32[3]; - - if (norm1 < eps) - { - return norm2 < eps ? 1.0f : 0.0f; - } - - return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); -} - -static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, uint32_t qty) -{ - return NormScalarProductSIMD(pVect1, pVect2, qty); -} - -template static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) -{ - return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty)); -} - -// static template float CosineSimilarity2<__int8>(const __int8* pVect1, -// const __int8* pVect2, size_t qty); - -// static template float CosineSimilarity2(const float* pVect1, -// const float* pVect2, size_t qty); - -template static void CosineSimilarityNormalize(T *pVector, uint32_t qty) -{ - T sum = 0; - for (uint32_t i = 0; i < qty; ++i) - { - sum += pVector[i] * pVector[i]; - } - sum = 1 / sqrt(sum); - if (sum == 0) - { - sum = numeric_limits::min(); - } - for (uint32_t i = 0; i < qty; ++i) - { - pVector[i] *= sum; - } -} - -// template static void CosineSimilarityNormalize(float* pVector, -// size_t qty); -// template static void CosineSimilarityNormalize(double* pVector, -// size_t qty); - -template <> void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int8 type vector, you can not use cosine distance!"); -} - -template <> void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int16 type vector, you can not use cosine distance!"); -} - -template <> void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int type vector, you can not use cosine distance!"); -} -} // namespace diskann -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/defaults.h b/packages/leann-backend-diskann/third_party/DiskANN/include/defaults.h deleted file mode 100644 index ef1750f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/defaults.h +++ /dev/null @@ -1,34 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include - -namespace diskann -{ -namespace defaults -{ -const float ALPHA = 1.2f; -const uint32_t NUM_THREADS = 0; -const uint32_t MAX_OCCLUSION_SIZE = 750; -const bool HAS_LABELS = false; -const uint32_t FILTER_LIST_SIZE = 0; -const uint32_t NUM_FROZEN_POINTS_STATIC = 0; -const uint32_t NUM_FROZEN_POINTS_DYNAMIC = 1; - -// In-mem index related limits -const float GRAPH_SLACK_FACTOR = 1.3f; - -// SSD Index related limits -const uint64_t MAX_GRAPH_DEGREE = 512; -const uint64_t SECTOR_LEN = 4096; -const uint64_t MAX_N_SECTOR_READS = 128; - -// following constants should always be specified, but are useful as a -// sensible default at cli / python boundaries -const uint32_t MAX_DEGREE = 64; -const uint32_t BUILD_LIST_SIZE = 100; -const uint32_t SATURATE_GRAPH = false; -const uint32_t SEARCH_LIST_SIZE = 100; -} // namespace defaults -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/disk_utils.h b/packages/leann-backend-diskann/third_party/DiskANN/include/disk_utils.h deleted file mode 100644 index 08f046d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/disk_utils.h +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#ifdef __APPLE__ -#else -#include -#endif - -#ifdef _WINDOWS -#include -typedef HANDLE FileHandle; -#else -#include -typedef int FileHandle; -#endif - -#include "cached_io.h" -#include "common_includes.h" - -#include "utils.h" -#include "windows_customizations.h" - -namespace diskann -{ -const size_t MAX_SAMPLE_POINTS_FOR_WARMUP = 100000; -const double PQ_TRAINING_SET_FRACTION = 0.1; -const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; -const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; -const uint32_t NUM_NODES_TO_CACHE = 250000; -const uint32_t WARMUP_L = 20; -const uint32_t NUM_KMEANS_REPS = 12; - -template class PQFlashIndex; - -DISKANN_DLLEXPORT double get_memory_budget(const std::string &mem_budget_str); -DISKANN_DLLEXPORT double get_memory_budget(double search_ram_budget_in_gb); -DISKANN_DLLEXPORT void add_new_file_to_single_index(std::string index_file, std::string new_file); - -DISKANN_DLLEXPORT size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim); - -DISKANN_DLLEXPORT void read_idmap(const std::string &fname, std::vector &ivecs); - -#ifdef EXEC_ENV_OLS -template -DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -#else -template -DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -#endif - -DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, - const std::string &idmaps_prefix, const std::string &idmaps_suffix, - const uint64_t nshards, uint32_t max_degree, const std::string &output_vamana, - const std::string &medoids_file, bool use_filters = false, - const std::string &labels_to_medoids_file = std::string("")); - -DISKANN_DLLEXPORT void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file); - -template -DISKANN_DLLEXPORT std::string preprocess_base_file(const std::string &infile, const std::string &indexPrefix, - diskann::Metric &distMetric); - -template -DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann::Metric _compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_file, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters = false, - const std::string &label_file = std::string(""), - const std::string &labels_to_medoids_file = std::string(""), - const std::string &universal_label = "", const uint32_t Lf = 0); - -template -DISKANN_DLLEXPORT uint32_t optimize_beamwidth(std::unique_ptr> &_pFlashIndex, - T *tuning_sample, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw = 2); - -template -DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, - diskann::Metric _compareMetric, bool use_opq = false, - const std::string &codebook_prefix = "", // default is empty for no codebook pass in - bool use_filters = false, - const std::string &label_file = std::string(""), // default is empty string for no label_file - const std::string &universal_label = "", const uint32_t filter_threshold = 0, - const uint32_t Lf = 0); // default is empty string for no universal label - -template -DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file = std::string("")); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/distance.h b/packages/leann-backend-diskann/third_party/DiskANN/include/distance.h deleted file mode 100644 index 7198308..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/distance.h +++ /dev/null @@ -1,236 +0,0 @@ -#pragma once -#include "windows_customizations.h" -#include -#include - -namespace diskann -{ -enum Metric -{ - L2 = 0, - INNER_PRODUCT = 1, - COSINE = 2, - FAST_L2 = 3 -}; - -template class Distance -{ - public: - DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) - { - } - - // distance comparison function - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const = 0; - - // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB, - uint32_t length) const; - - // For MIPS, normalization adds an extra dimension to the vectors. - // This function lets callers know if the normalization process - // changes the dimension. - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const; - - DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; - - // This is for efficiency. If no normalization is required, the callers - // can simply ignore the normalize_data_for_build() function. - DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - - // Check the preprocessing_required() function before calling this. - // Clients can call the function like this: - // - // if (metric->preprocessing_required()){ - // T* normalized_data_batch; - // Split data into batches of batch_size and for each, call: - // metric->preprocess_base_points(data_batch, batch_size); - // - // TODO: This does not take into account the case for SSD inner product - // where the dimensions change after normalization. - DISKANN_DLLEXPORT virtual void preprocess_base_points(T *original_data, const size_t orig_dim, - const size_t num_points); - - // Invokes normalization for a single vector during search. The scratch space - // has to be created by the caller keeping track of the fact that - // normalization might change the dimension of the query vector. - DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query); - - // If an algorithm has a requirement that some data be aligned to a certain - // boundary it can use this function to indicate that requirement. Currently, - // we are setting it to 8 because that works well for AVX2. If we have AVX512 - // implementations of distance algos, they might have to set this to 16 - // (depending on how they are implemented) - DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; - - // Providing a default implementation for the virtual destructor because we - // don't expect most metric implementations to need it. - DISKANN_DLLEXPORT virtual ~Distance() = default; - - protected: - diskann::Metric _distance_metric; - size_t _alignment_factor = 8; -}; - -class DistanceCosineInt8 : public Distance -{ - public: - DistanceCosineInt8() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; -}; - -class DistanceL2Int8 : public Distance -{ - public: - DistanceL2Int8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; -}; - -// AVX implementations. Borrowed from HNSW code. -class AVXDistanceL2Int8 : public Distance -{ - public: - AVXDistanceL2Int8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; -}; - -class DistanceCosineFloat : public Distance -{ - public: - DistanceCosineFloat() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; -}; - -class DistanceL2Float : public Distance -{ - public: - DistanceL2Float() : Distance(diskann::Metric::L2) - { - } - -#ifdef _WINDOWS - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; -#else - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const __attribute__((hot)); -#endif -}; - -class AVXDistanceL2Float : public Distance -{ - public: - AVXDistanceL2Float() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; -}; - -template class SlowDistanceL2 : public Distance -{ - public: - SlowDistanceL2() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const; -}; - -class SlowDistanceCosineUInt8 : public Distance -{ - public: - SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; -}; - -class DistanceL2UInt8 : public Distance -{ - public: - DistanceL2UInt8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; -}; - -template class DistanceInnerProduct : public Distance -{ - public: - DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) - { - } - - DistanceInnerProduct(diskann::Metric metric) : Distance(metric) - { - } - inline float inner_product(const T *a, const T *b, unsigned size) const; - - inline float compare(const T *a, const T *b, unsigned size) const - { - float result = inner_product(a, b, size); - // if (result < 0) - // return std::numeric_limits::max(); - // else - return -result; - } -}; - -template class DistanceFastL2 : public DistanceInnerProduct -{ - // currently defined only for float. - // templated for future use. - public: - DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) - { - } - float norm(const T *a, unsigned size) const; - float compare(const T *a, const T *b, float norm, unsigned size) const; -}; - -class AVXDistanceInnerProductFloat : public Distance -{ - public: - AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; -}; - -class AVXNormalizedCosineDistanceFloat : public Distance -{ - private: - AVXDistanceInnerProductFloat _innerProduct; - - protected: - void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; - - public: - AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const override - { - // Inner product returns negative values to indicate distance. - // This will ensure that cosine is between -1 and 1. - return 1.0f + _innerProduct.compare(a, b, length); - } - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override; - - DISKANN_DLLEXPORT virtual bool preprocessing_required() const override; - - DISKANN_DLLEXPORT virtual void preprocess_base_points(float *original_data, const size_t orig_dim, - const size_t num_points) override; - - DISKANN_DLLEXPORT virtual void preprocess_query(const float *query_vec, const size_t query_dim, - float *scratch_query_vector) override; -}; - -template Distance *get_distance_function(Metric m); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/embedding.pb.h b/packages/leann-backend-diskann/third_party/DiskANN/include/embedding.pb.h deleted file mode 100644 index 9f5c2b7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/embedding.pb.h +++ /dev/null @@ -1,675 +0,0 @@ -// Generated by the protocol buffer compiler. DO NOT EDIT! -// source: embedding.proto - -#ifndef GOOGLE_PROTOBUF_INCLUDED_embedding_2eproto -#define GOOGLE_PROTOBUF_INCLUDED_embedding_2eproto - -#include -#include - -#include -#if PROTOBUF_VERSION < 3012000 -#error This file was generated by a newer version of protoc which is -#error incompatible with your Protocol Buffer headers. Please update -#error your headers. -#endif -#if 3012004 < PROTOBUF_MIN_PROTOC_VERSION -#error This file was generated by an older version of protoc which is -#error incompatible with your Protocol Buffer headers. Please -#error regenerate this file with a newer version of protoc. -#endif - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include // IWYU pragma: export -#include // IWYU pragma: export -#include -// @@protoc_insertion_point(includes) -#include -#define PROTOBUF_INTERNAL_EXPORT_embedding_2eproto -PROTOBUF_NAMESPACE_OPEN -namespace internal { -class AnyMetadata; -} // namespace internal -PROTOBUF_NAMESPACE_CLOSE - -// Internal implementation detail -- do not use these members. -struct TableStruct_embedding_2eproto { - static const ::PROTOBUF_NAMESPACE_ID::internal::ParseTableField entries[] - PROTOBUF_SECTION_VARIABLE(protodesc_cold); - static const ::PROTOBUF_NAMESPACE_ID::internal::AuxillaryParseTableField aux[] - PROTOBUF_SECTION_VARIABLE(protodesc_cold); - static const ::PROTOBUF_NAMESPACE_ID::internal::ParseTable schema[2] - PROTOBUF_SECTION_VARIABLE(protodesc_cold); - static const ::PROTOBUF_NAMESPACE_ID::internal::FieldMetadata field_metadata[]; - static const ::PROTOBUF_NAMESPACE_ID::internal::SerializationTable serialization_table[]; - static const ::PROTOBUF_NAMESPACE_ID::uint32 offsets[]; -}; -extern const ::PROTOBUF_NAMESPACE_ID::internal::DescriptorTable descriptor_table_embedding_2eproto; -namespace protoembedding { -class NodeEmbeddingRequest; -class NodeEmbeddingRequestDefaultTypeInternal; -extern NodeEmbeddingRequestDefaultTypeInternal _NodeEmbeddingRequest_default_instance_; -class NodeEmbeddingResponse; -class NodeEmbeddingResponseDefaultTypeInternal; -extern NodeEmbeddingResponseDefaultTypeInternal _NodeEmbeddingResponse_default_instance_; -} // namespace protoembedding -PROTOBUF_NAMESPACE_OPEN -template<> ::protoembedding::NodeEmbeddingRequest* Arena::CreateMaybeMessage<::protoembedding::NodeEmbeddingRequest>(Arena*); -template<> ::protoembedding::NodeEmbeddingResponse* Arena::CreateMaybeMessage<::protoembedding::NodeEmbeddingResponse>(Arena*); -PROTOBUF_NAMESPACE_CLOSE -namespace protoembedding { - -// =================================================================== - -class NodeEmbeddingRequest PROTOBUF_FINAL : - public ::PROTOBUF_NAMESPACE_ID::Message /* @@protoc_insertion_point(class_definition:protoembedding.NodeEmbeddingRequest) */ { - public: - inline NodeEmbeddingRequest() : NodeEmbeddingRequest(nullptr) {}; - virtual ~NodeEmbeddingRequest(); - - NodeEmbeddingRequest(const NodeEmbeddingRequest& from); - NodeEmbeddingRequest(NodeEmbeddingRequest&& from) noexcept - : NodeEmbeddingRequest() { - *this = ::std::move(from); - } - - inline NodeEmbeddingRequest& operator=(const NodeEmbeddingRequest& from) { - CopyFrom(from); - return *this; - } - inline NodeEmbeddingRequest& operator=(NodeEmbeddingRequest&& from) noexcept { - if (GetArena() == from.GetArena()) { - if (this != &from) InternalSwap(&from); - } else { - CopyFrom(from); - } - return *this; - } - - static const ::PROTOBUF_NAMESPACE_ID::Descriptor* descriptor() { - return GetDescriptor(); - } - static const ::PROTOBUF_NAMESPACE_ID::Descriptor* GetDescriptor() { - return GetMetadataStatic().descriptor; - } - static const ::PROTOBUF_NAMESPACE_ID::Reflection* GetReflection() { - return GetMetadataStatic().reflection; - } - static const NodeEmbeddingRequest& default_instance(); - - static void InitAsDefaultInstance(); // FOR INTERNAL USE ONLY - static inline const NodeEmbeddingRequest* internal_default_instance() { - return reinterpret_cast( - &_NodeEmbeddingRequest_default_instance_); - } - static constexpr int kIndexInFileMessages = - 0; - - friend void swap(NodeEmbeddingRequest& a, NodeEmbeddingRequest& b) { - a.Swap(&b); - } - inline void Swap(NodeEmbeddingRequest* other) { - if (other == this) return; - if (GetArena() == other->GetArena()) { - InternalSwap(other); - } else { - ::PROTOBUF_NAMESPACE_ID::internal::GenericSwap(this, other); - } - } - void UnsafeArenaSwap(NodeEmbeddingRequest* other) { - if (other == this) return; - GOOGLE_DCHECK(GetArena() == other->GetArena()); - InternalSwap(other); - } - - // implements Message ---------------------------------------------- - - inline NodeEmbeddingRequest* New() const final { - return CreateMaybeMessage(nullptr); - } - - NodeEmbeddingRequest* New(::PROTOBUF_NAMESPACE_ID::Arena* arena) const final { - return CreateMaybeMessage(arena); - } - void CopyFrom(const ::PROTOBUF_NAMESPACE_ID::Message& from) final; - void MergeFrom(const ::PROTOBUF_NAMESPACE_ID::Message& from) final; - void CopyFrom(const NodeEmbeddingRequest& from); - void MergeFrom(const NodeEmbeddingRequest& from); - PROTOBUF_ATTRIBUTE_REINITIALIZES void Clear() final; - bool IsInitialized() const final; - - size_t ByteSizeLong() const final; - const char* _InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) final; - ::PROTOBUF_NAMESPACE_ID::uint8* _InternalSerialize( - ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const final; - int GetCachedSize() const final { return _cached_size_.Get(); } - - private: - inline void SharedCtor(); - inline void SharedDtor(); - void SetCachedSize(int size) const final; - void InternalSwap(NodeEmbeddingRequest* other); - friend class ::PROTOBUF_NAMESPACE_ID::internal::AnyMetadata; - static ::PROTOBUF_NAMESPACE_ID::StringPiece FullMessageName() { - return "protoembedding.NodeEmbeddingRequest"; - } - protected: - explicit NodeEmbeddingRequest(::PROTOBUF_NAMESPACE_ID::Arena* arena); - private: - static void ArenaDtor(void* object); - inline void RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena* arena); - public: - - ::PROTOBUF_NAMESPACE_ID::Metadata GetMetadata() const final; - private: - static ::PROTOBUF_NAMESPACE_ID::Metadata GetMetadataStatic() { - ::PROTOBUF_NAMESPACE_ID::internal::AssignDescriptors(&::descriptor_table_embedding_2eproto); - return ::descriptor_table_embedding_2eproto.file_level_metadata[kIndexInFileMessages]; - } - - public: - - // nested types ---------------------------------------------------- - - // accessors ------------------------------------------------------- - - enum : int { - kNodeIdsFieldNumber = 1, - }; - // repeated uint32 node_ids = 1; - int node_ids_size() const; - private: - int _internal_node_ids_size() const; - public: - void clear_node_ids(); - private: - ::PROTOBUF_NAMESPACE_ID::uint32 _internal_node_ids(int index) const; - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& - _internal_node_ids() const; - void _internal_add_node_ids(::PROTOBUF_NAMESPACE_ID::uint32 value); - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* - _internal_mutable_node_ids(); - public: - ::PROTOBUF_NAMESPACE_ID::uint32 node_ids(int index) const; - void set_node_ids(int index, ::PROTOBUF_NAMESPACE_ID::uint32 value); - void add_node_ids(::PROTOBUF_NAMESPACE_ID::uint32 value); - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& - node_ids() const; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* - mutable_node_ids(); - - // @@protoc_insertion_point(class_scope:protoembedding.NodeEmbeddingRequest) - private: - class _Internal; - - template friend class ::PROTOBUF_NAMESPACE_ID::Arena::InternalHelper; - typedef void InternalArenaConstructable_; - typedef void DestructorSkippable_; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 > node_ids_; - mutable std::atomic _node_ids_cached_byte_size_; - mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_; - friend struct ::TableStruct_embedding_2eproto; -}; -// ------------------------------------------------------------------- - -class NodeEmbeddingResponse PROTOBUF_FINAL : - public ::PROTOBUF_NAMESPACE_ID::Message /* @@protoc_insertion_point(class_definition:protoembedding.NodeEmbeddingResponse) */ { - public: - inline NodeEmbeddingResponse() : NodeEmbeddingResponse(nullptr) {}; - virtual ~NodeEmbeddingResponse(); - - NodeEmbeddingResponse(const NodeEmbeddingResponse& from); - NodeEmbeddingResponse(NodeEmbeddingResponse&& from) noexcept - : NodeEmbeddingResponse() { - *this = ::std::move(from); - } - - inline NodeEmbeddingResponse& operator=(const NodeEmbeddingResponse& from) { - CopyFrom(from); - return *this; - } - inline NodeEmbeddingResponse& operator=(NodeEmbeddingResponse&& from) noexcept { - if (GetArena() == from.GetArena()) { - if (this != &from) InternalSwap(&from); - } else { - CopyFrom(from); - } - return *this; - } - - static const ::PROTOBUF_NAMESPACE_ID::Descriptor* descriptor() { - return GetDescriptor(); - } - static const ::PROTOBUF_NAMESPACE_ID::Descriptor* GetDescriptor() { - return GetMetadataStatic().descriptor; - } - static const ::PROTOBUF_NAMESPACE_ID::Reflection* GetReflection() { - return GetMetadataStatic().reflection; - } - static const NodeEmbeddingResponse& default_instance(); - - static void InitAsDefaultInstance(); // FOR INTERNAL USE ONLY - static inline const NodeEmbeddingResponse* internal_default_instance() { - return reinterpret_cast( - &_NodeEmbeddingResponse_default_instance_); - } - static constexpr int kIndexInFileMessages = - 1; - - friend void swap(NodeEmbeddingResponse& a, NodeEmbeddingResponse& b) { - a.Swap(&b); - } - inline void Swap(NodeEmbeddingResponse* other) { - if (other == this) return; - if (GetArena() == other->GetArena()) { - InternalSwap(other); - } else { - ::PROTOBUF_NAMESPACE_ID::internal::GenericSwap(this, other); - } - } - void UnsafeArenaSwap(NodeEmbeddingResponse* other) { - if (other == this) return; - GOOGLE_DCHECK(GetArena() == other->GetArena()); - InternalSwap(other); - } - - // implements Message ---------------------------------------------- - - inline NodeEmbeddingResponse* New() const final { - return CreateMaybeMessage(nullptr); - } - - NodeEmbeddingResponse* New(::PROTOBUF_NAMESPACE_ID::Arena* arena) const final { - return CreateMaybeMessage(arena); - } - void CopyFrom(const ::PROTOBUF_NAMESPACE_ID::Message& from) final; - void MergeFrom(const ::PROTOBUF_NAMESPACE_ID::Message& from) final; - void CopyFrom(const NodeEmbeddingResponse& from); - void MergeFrom(const NodeEmbeddingResponse& from); - PROTOBUF_ATTRIBUTE_REINITIALIZES void Clear() final; - bool IsInitialized() const final; - - size_t ByteSizeLong() const final; - const char* _InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) final; - ::PROTOBUF_NAMESPACE_ID::uint8* _InternalSerialize( - ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const final; - int GetCachedSize() const final { return _cached_size_.Get(); } - - private: - inline void SharedCtor(); - inline void SharedDtor(); - void SetCachedSize(int size) const final; - void InternalSwap(NodeEmbeddingResponse* other); - friend class ::PROTOBUF_NAMESPACE_ID::internal::AnyMetadata; - static ::PROTOBUF_NAMESPACE_ID::StringPiece FullMessageName() { - return "protoembedding.NodeEmbeddingResponse"; - } - protected: - explicit NodeEmbeddingResponse(::PROTOBUF_NAMESPACE_ID::Arena* arena); - private: - static void ArenaDtor(void* object); - inline void RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena* arena); - public: - - ::PROTOBUF_NAMESPACE_ID::Metadata GetMetadata() const final; - private: - static ::PROTOBUF_NAMESPACE_ID::Metadata GetMetadataStatic() { - ::PROTOBUF_NAMESPACE_ID::internal::AssignDescriptors(&::descriptor_table_embedding_2eproto); - return ::descriptor_table_embedding_2eproto.file_level_metadata[kIndexInFileMessages]; - } - - public: - - // nested types ---------------------------------------------------- - - // accessors ------------------------------------------------------- - - enum : int { - kDimensionsFieldNumber = 2, - kMissingIdsFieldNumber = 3, - kEmbeddingsDataFieldNumber = 1, - }; - // repeated int32 dimensions = 2; - int dimensions_size() const; - private: - int _internal_dimensions_size() const; - public: - void clear_dimensions(); - private: - ::PROTOBUF_NAMESPACE_ID::int32 _internal_dimensions(int index) const; - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >& - _internal_dimensions() const; - void _internal_add_dimensions(::PROTOBUF_NAMESPACE_ID::int32 value); - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >* - _internal_mutable_dimensions(); - public: - ::PROTOBUF_NAMESPACE_ID::int32 dimensions(int index) const; - void set_dimensions(int index, ::PROTOBUF_NAMESPACE_ID::int32 value); - void add_dimensions(::PROTOBUF_NAMESPACE_ID::int32 value); - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >& - dimensions() const; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >* - mutable_dimensions(); - - // repeated uint32 missing_ids = 3; - int missing_ids_size() const; - private: - int _internal_missing_ids_size() const; - public: - void clear_missing_ids(); - private: - ::PROTOBUF_NAMESPACE_ID::uint32 _internal_missing_ids(int index) const; - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& - _internal_missing_ids() const; - void _internal_add_missing_ids(::PROTOBUF_NAMESPACE_ID::uint32 value); - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* - _internal_mutable_missing_ids(); - public: - ::PROTOBUF_NAMESPACE_ID::uint32 missing_ids(int index) const; - void set_missing_ids(int index, ::PROTOBUF_NAMESPACE_ID::uint32 value); - void add_missing_ids(::PROTOBUF_NAMESPACE_ID::uint32 value); - const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& - missing_ids() const; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* - mutable_missing_ids(); - - // bytes embeddings_data = 1; - void clear_embeddings_data(); - const std::string& embeddings_data() const; - void set_embeddings_data(const std::string& value); - void set_embeddings_data(std::string&& value); - void set_embeddings_data(const char* value); - void set_embeddings_data(const void* value, size_t size); - std::string* mutable_embeddings_data(); - std::string* release_embeddings_data(); - void set_allocated_embeddings_data(std::string* embeddings_data); - GOOGLE_PROTOBUF_RUNTIME_DEPRECATED("The unsafe_arena_ accessors for" - " string fields are deprecated and will be removed in a" - " future release.") - std::string* unsafe_arena_release_embeddings_data(); - GOOGLE_PROTOBUF_RUNTIME_DEPRECATED("The unsafe_arena_ accessors for" - " string fields are deprecated and will be removed in a" - " future release.") - void unsafe_arena_set_allocated_embeddings_data( - std::string* embeddings_data); - private: - const std::string& _internal_embeddings_data() const; - void _internal_set_embeddings_data(const std::string& value); - std::string* _internal_mutable_embeddings_data(); - public: - - // @@protoc_insertion_point(class_scope:protoembedding.NodeEmbeddingResponse) - private: - class _Internal; - - template friend class ::PROTOBUF_NAMESPACE_ID::Arena::InternalHelper; - typedef void InternalArenaConstructable_; - typedef void DestructorSkippable_; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 > dimensions_; - mutable std::atomic _dimensions_cached_byte_size_; - ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 > missing_ids_; - mutable std::atomic _missing_ids_cached_byte_size_; - ::PROTOBUF_NAMESPACE_ID::internal::ArenaStringPtr embeddings_data_; - mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_; - friend struct ::TableStruct_embedding_2eproto; -}; -// =================================================================== - - -// =================================================================== - -#ifdef __GNUC__ - #pragma GCC diagnostic push - #pragma GCC diagnostic ignored "-Wstrict-aliasing" -#endif // __GNUC__ -// NodeEmbeddingRequest - -// repeated uint32 node_ids = 1; -inline int NodeEmbeddingRequest::_internal_node_ids_size() const { - return node_ids_.size(); -} -inline int NodeEmbeddingRequest::node_ids_size() const { - return _internal_node_ids_size(); -} -inline void NodeEmbeddingRequest::clear_node_ids() { - node_ids_.Clear(); -} -inline ::PROTOBUF_NAMESPACE_ID::uint32 NodeEmbeddingRequest::_internal_node_ids(int index) const { - return node_ids_.Get(index); -} -inline ::PROTOBUF_NAMESPACE_ID::uint32 NodeEmbeddingRequest::node_ids(int index) const { - // @@protoc_insertion_point(field_get:protoembedding.NodeEmbeddingRequest.node_ids) - return _internal_node_ids(index); -} -inline void NodeEmbeddingRequest::set_node_ids(int index, ::PROTOBUF_NAMESPACE_ID::uint32 value) { - node_ids_.Set(index, value); - // @@protoc_insertion_point(field_set:protoembedding.NodeEmbeddingRequest.node_ids) -} -inline void NodeEmbeddingRequest::_internal_add_node_ids(::PROTOBUF_NAMESPACE_ID::uint32 value) { - node_ids_.Add(value); -} -inline void NodeEmbeddingRequest::add_node_ids(::PROTOBUF_NAMESPACE_ID::uint32 value) { - _internal_add_node_ids(value); - // @@protoc_insertion_point(field_add:protoembedding.NodeEmbeddingRequest.node_ids) -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& -NodeEmbeddingRequest::_internal_node_ids() const { - return node_ids_; -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& -NodeEmbeddingRequest::node_ids() const { - // @@protoc_insertion_point(field_list:protoembedding.NodeEmbeddingRequest.node_ids) - return _internal_node_ids(); -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* -NodeEmbeddingRequest::_internal_mutable_node_ids() { - return &node_ids_; -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* -NodeEmbeddingRequest::mutable_node_ids() { - // @@protoc_insertion_point(field_mutable_list:protoembedding.NodeEmbeddingRequest.node_ids) - return _internal_mutable_node_ids(); -} - -// ------------------------------------------------------------------- - -// NodeEmbeddingResponse - -// bytes embeddings_data = 1; -inline void NodeEmbeddingResponse::clear_embeddings_data() { - embeddings_data_.ClearToEmpty(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), GetArena()); -} -inline const std::string& NodeEmbeddingResponse::embeddings_data() const { - // @@protoc_insertion_point(field_get:protoembedding.NodeEmbeddingResponse.embeddings_data) - return _internal_embeddings_data(); -} -inline void NodeEmbeddingResponse::set_embeddings_data(const std::string& value) { - _internal_set_embeddings_data(value); - // @@protoc_insertion_point(field_set:protoembedding.NodeEmbeddingResponse.embeddings_data) -} -inline std::string* NodeEmbeddingResponse::mutable_embeddings_data() { - // @@protoc_insertion_point(field_mutable:protoembedding.NodeEmbeddingResponse.embeddings_data) - return _internal_mutable_embeddings_data(); -} -inline const std::string& NodeEmbeddingResponse::_internal_embeddings_data() const { - return embeddings_data_.Get(); -} -inline void NodeEmbeddingResponse::_internal_set_embeddings_data(const std::string& value) { - - embeddings_data_.Set(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), value, GetArena()); -} -inline void NodeEmbeddingResponse::set_embeddings_data(std::string&& value) { - - embeddings_data_.Set( - &::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), ::std::move(value), GetArena()); - // @@protoc_insertion_point(field_set_rvalue:protoembedding.NodeEmbeddingResponse.embeddings_data) -} -inline void NodeEmbeddingResponse::set_embeddings_data(const char* value) { - GOOGLE_DCHECK(value != nullptr); - - embeddings_data_.Set(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), ::std::string(value), - GetArena()); - // @@protoc_insertion_point(field_set_char:protoembedding.NodeEmbeddingResponse.embeddings_data) -} -inline void NodeEmbeddingResponse::set_embeddings_data(const void* value, - size_t size) { - - embeddings_data_.Set(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), ::std::string( - reinterpret_cast(value), size), GetArena()); - // @@protoc_insertion_point(field_set_pointer:protoembedding.NodeEmbeddingResponse.embeddings_data) -} -inline std::string* NodeEmbeddingResponse::_internal_mutable_embeddings_data() { - - return embeddings_data_.Mutable(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), GetArena()); -} -inline std::string* NodeEmbeddingResponse::release_embeddings_data() { - // @@protoc_insertion_point(field_release:protoembedding.NodeEmbeddingResponse.embeddings_data) - return embeddings_data_.Release(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), GetArena()); -} -inline void NodeEmbeddingResponse::set_allocated_embeddings_data(std::string* embeddings_data) { - if (embeddings_data != nullptr) { - - } else { - - } - embeddings_data_.SetAllocated(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), embeddings_data, - GetArena()); - // @@protoc_insertion_point(field_set_allocated:protoembedding.NodeEmbeddingResponse.embeddings_data) -} -inline std::string* NodeEmbeddingResponse::unsafe_arena_release_embeddings_data() { - // @@protoc_insertion_point(field_unsafe_arena_release:protoembedding.NodeEmbeddingResponse.embeddings_data) - GOOGLE_DCHECK(GetArena() != nullptr); - - return embeddings_data_.UnsafeArenaRelease(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), - GetArena()); -} -inline void NodeEmbeddingResponse::unsafe_arena_set_allocated_embeddings_data( - std::string* embeddings_data) { - GOOGLE_DCHECK(GetArena() != nullptr); - if (embeddings_data != nullptr) { - - } else { - - } - embeddings_data_.UnsafeArenaSetAllocated(&::PROTOBUF_NAMESPACE_ID::internal::GetEmptyStringAlreadyInited(), - embeddings_data, GetArena()); - // @@protoc_insertion_point(field_unsafe_arena_set_allocated:protoembedding.NodeEmbeddingResponse.embeddings_data) -} - -// repeated int32 dimensions = 2; -inline int NodeEmbeddingResponse::_internal_dimensions_size() const { - return dimensions_.size(); -} -inline int NodeEmbeddingResponse::dimensions_size() const { - return _internal_dimensions_size(); -} -inline void NodeEmbeddingResponse::clear_dimensions() { - dimensions_.Clear(); -} -inline ::PROTOBUF_NAMESPACE_ID::int32 NodeEmbeddingResponse::_internal_dimensions(int index) const { - return dimensions_.Get(index); -} -inline ::PROTOBUF_NAMESPACE_ID::int32 NodeEmbeddingResponse::dimensions(int index) const { - // @@protoc_insertion_point(field_get:protoembedding.NodeEmbeddingResponse.dimensions) - return _internal_dimensions(index); -} -inline void NodeEmbeddingResponse::set_dimensions(int index, ::PROTOBUF_NAMESPACE_ID::int32 value) { - dimensions_.Set(index, value); - // @@protoc_insertion_point(field_set:protoembedding.NodeEmbeddingResponse.dimensions) -} -inline void NodeEmbeddingResponse::_internal_add_dimensions(::PROTOBUF_NAMESPACE_ID::int32 value) { - dimensions_.Add(value); -} -inline void NodeEmbeddingResponse::add_dimensions(::PROTOBUF_NAMESPACE_ID::int32 value) { - _internal_add_dimensions(value); - // @@protoc_insertion_point(field_add:protoembedding.NodeEmbeddingResponse.dimensions) -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >& -NodeEmbeddingResponse::_internal_dimensions() const { - return dimensions_; -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >& -NodeEmbeddingResponse::dimensions() const { - // @@protoc_insertion_point(field_list:protoembedding.NodeEmbeddingResponse.dimensions) - return _internal_dimensions(); -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >* -NodeEmbeddingResponse::_internal_mutable_dimensions() { - return &dimensions_; -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int32 >* -NodeEmbeddingResponse::mutable_dimensions() { - // @@protoc_insertion_point(field_mutable_list:protoembedding.NodeEmbeddingResponse.dimensions) - return _internal_mutable_dimensions(); -} - -// repeated uint32 missing_ids = 3; -inline int NodeEmbeddingResponse::_internal_missing_ids_size() const { - return missing_ids_.size(); -} -inline int NodeEmbeddingResponse::missing_ids_size() const { - return _internal_missing_ids_size(); -} -inline void NodeEmbeddingResponse::clear_missing_ids() { - missing_ids_.Clear(); -} -inline ::PROTOBUF_NAMESPACE_ID::uint32 NodeEmbeddingResponse::_internal_missing_ids(int index) const { - return missing_ids_.Get(index); -} -inline ::PROTOBUF_NAMESPACE_ID::uint32 NodeEmbeddingResponse::missing_ids(int index) const { - // @@protoc_insertion_point(field_get:protoembedding.NodeEmbeddingResponse.missing_ids) - return _internal_missing_ids(index); -} -inline void NodeEmbeddingResponse::set_missing_ids(int index, ::PROTOBUF_NAMESPACE_ID::uint32 value) { - missing_ids_.Set(index, value); - // @@protoc_insertion_point(field_set:protoembedding.NodeEmbeddingResponse.missing_ids) -} -inline void NodeEmbeddingResponse::_internal_add_missing_ids(::PROTOBUF_NAMESPACE_ID::uint32 value) { - missing_ids_.Add(value); -} -inline void NodeEmbeddingResponse::add_missing_ids(::PROTOBUF_NAMESPACE_ID::uint32 value) { - _internal_add_missing_ids(value); - // @@protoc_insertion_point(field_add:protoembedding.NodeEmbeddingResponse.missing_ids) -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& -NodeEmbeddingResponse::_internal_missing_ids() const { - return missing_ids_; -} -inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >& -NodeEmbeddingResponse::missing_ids() const { - // @@protoc_insertion_point(field_list:protoembedding.NodeEmbeddingResponse.missing_ids) - return _internal_missing_ids(); -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* -NodeEmbeddingResponse::_internal_mutable_missing_ids() { - return &missing_ids_; -} -inline ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::uint32 >* -NodeEmbeddingResponse::mutable_missing_ids() { - // @@protoc_insertion_point(field_mutable_list:protoembedding.NodeEmbeddingResponse.missing_ids) - return _internal_mutable_missing_ids(); -} - -#ifdef __GNUC__ - #pragma GCC diagnostic pop -#endif // __GNUC__ -// ------------------------------------------------------------------- - - -// @@protoc_insertion_point(namespace_scope) - -} // namespace protoembedding - -// @@protoc_insertion_point(global_scope) - -#include -#endif // GOOGLE_PROTOBUF_INCLUDED_GOOGLE_PROTOBUF_INCLUDED_embedding_2eproto diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/embedding_compute.h b/packages/leann-backend-diskann/third_party/DiskANN/include/embedding_compute.h deleted file mode 100644 index 354c9c5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/embedding_compute.h +++ /dev/null @@ -1,118 +0,0 @@ -#pragma once - -#include -#include - -#ifdef PYBIND11_EMBEDDED -#include -#else -#include -#endif -#include -#include - -namespace py = pybind11; - -namespace diskann -{ - -class PYBIND11_EXPORT EmbeddingComputer -{ - public: - static EmbeddingComputer &getInstance() - { - static EmbeddingComputer instance; - return instance; - } - - void initialize(const std::string &model_path) - { - try - { - py::module_ sys = py::module_::import("sys"); - py::module_ os = py::module_::import("os"); - - // Add the directory containing embedd_micro.py to Python path - std::string micro_dir = "micro"; - sys.attr("path").attr("append")(micro_dir); - - // Import our module - py::module_ embedd = py::module_::import("embedd_micro"); - - // Create benchmark config - py::object config = embedd.attr("BenchmarkConfig")(model_path, // model_path - py::list(), // empty batch_sizes - 256, // seq_length - 1, // num_runs - true, // use_fp16 - false, // use_cuda_graphs - false // use_flash_attention - ); - - // Create benchmark instance - benchmark = embedd.attr("Benchmark")(config); - } - catch (const std::exception &e) - { - throw std::runtime_error("Failed to initialize Python embedding computer: " + std::string(e.what())); - } - } - - template - std::vector computeEmbeddings(const std::vector &points, size_t dim, size_t batch_size = 32) - { - try - { - // Convert points to numpy array - std::vector flattened_points; - flattened_points.reserve(points.size() * dim); - - for (const auto &point : points) - { - flattened_points.insert(flattened_points.end(), point, point + dim); - } - - py::array_t points_array({static_cast(points.size()), static_cast(dim)}, - flattened_points.data()); - - // Call compute_embeddings - py::object result = benchmark.attr("compute_embeddings")(points_array, batch_size); - - // Convert result back to C++ - py::array_t np_result = result.cast>(); - py::buffer_info buf = np_result.request(); - float *ptr = static_cast(buf.ptr); - - return std::vector(ptr, ptr + buf.size); - } - catch (const std::exception &e) - { - throw std::runtime_error("Failed to compute embeddings: " + std::string(e.what())); - } - } - - private: - EmbeddingComputer() - { -#ifdef PYBIND11_EMBEDDED - if (!Py_IsInitialized()) - { - py::initialize_interpreter(); - } -#endif - } - - ~EmbeddingComputer() - { -#ifdef PYBIND11_EMBEDDED - if (Py_IsInitialized()) - { - py::finalize_interpreter(); - } -#endif - } - - py::object benchmark; -}; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/exceptions.h b/packages/leann-backend-diskann/third_party/DiskANN/include/exceptions.h deleted file mode 100644 index 99e4e73..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/exceptions.h +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include - -namespace diskann -{ - -class NotImplementedException : public std::logic_error -{ - public: - NotImplementedException() : std::logic_error("Function not yet implemented.") - { - } -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/filter_utils.h b/packages/leann-backend-diskann/third_party/DiskANN/include/filter_utils.h deleted file mode 100644 index 55f7aed..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/filter_utils.h +++ /dev/null @@ -1,221 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#ifdef __APPLE__ -#else -#include -#endif - -#ifdef _WINDOWS -#include -typedef HANDLE FileHandle; -#else -#include -typedef int FileHandle; -#endif - -#ifndef _WINDOWS -#include -#endif - -#include "cached_io.h" -#include "common_includes.h" -#include "memory_mapper.h" -#include "utils.h" -#include "windows_customizations.h" - -// custom types (for readability) -typedef tsl::robin_set label_set; -typedef std::string path; - -// structs for returning multiple items from a function -typedef std::tuple, tsl::robin_map, tsl::robin_set> - parse_label_file_return_values; -typedef std::tuple>, uint64_t> load_label_index_return_values; - -namespace diskann -{ -template -DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, - unsigned R, unsigned L, float alpha, unsigned num_threads); - -DISKANN_DLLEXPORT load_label_index_return_values load_label_index(path label_index_path, - uint32_t label_number_of_points); - -template -DISKANN_DLLEXPORT std::tuple>, tsl::robin_set> parse_formatted_label_file( - path label_file); - -DISKANN_DLLEXPORT parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label); - -template -DISKANN_DLLEXPORT tsl::robin_map> generate_label_specific_vector_files_compat( - path input_data_path, tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); - -/* - * For each label, generates a file containing all vectors that have said label. - * Also copies data from original bin file to new dimension-aligned file. - * - * Utilizes POSIX functions mmap and writev in order to minimize memory - * overhead, so we include an STL version as well. - * - * Each data file is saved under the following format: - * input_data_path + "_" + label - */ -#ifndef _WINDOWS -template -inline tsl::robin_map> generate_label_specific_vector_files( - path input_data_path, tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) -{ -#ifndef _WINDOWS - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - diskann::MemoryMapper input_data(input_data_path); - char *input_start = input_data.getBuf(); - - uint32_t number_of_points, dimension; - std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); - std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - const size_t METADATA = 2 * sizeof(uint32_t); - if (number_of_points != point_ids_to_labels.size()) - { - std::cerr << "Error: number of points in labels file and data file differ." << std::endl; - throw; - } - - tsl::robin_map label_to_iovec_map; - tsl::robin_map label_to_curr_iovec; - tsl::robin_map> label_id_to_orig_id; - - // setup iovec list for each label - for (const auto &lbl : all_labels) - { - iovec *label_iovecs = (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec)); - if (label_iovecs == nullptr) - { - throw; - } - label_to_iovec_map[lbl] = label_iovecs; - label_to_curr_iovec[lbl] = 0; - label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]); - } - - // each point added to corresponding per-label iovec list - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) - { - char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); - iovec curr_iovec; - - curr_iovec.iov_base = curr_point; - curr_iovec.iov_len = VECTOR_SIZE; - for (const auto &lbl : point_ids_to_labels[point_id]) - { - *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec; - label_to_curr_iovec[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); - } - } - - // write each label iovec to resp. file - for (const auto &lbl : all_labels) - { - int label_input_data_fd; - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t curr_num_pts = labels_to_number_of_points[lbl]; - - label_input_data_fd = - open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); - if (label_input_data_fd == -1) - throw; - - // write metadata - uint32_t metadata[2] = {curr_num_pts, dimension}; - int return_value = write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); - if (return_value == -1) - { - throw; - } - - // limits on number of iovec structs per writev means we need to perform - // multiple writevs - size_t i = 0; - while (curr_num_pts > IOV_MAX) - { - return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX); - if (return_value == -1) - { - close(label_input_data_fd); - throw; - } - curr_num_pts -= IOV_MAX; - i += 1; - } - return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), curr_num_pts); - if (return_value == -1) - { - close(label_input_data_fd); - throw; - } - - free(label_to_iovec_map[lbl]); - close(label_input_data_fd); - } - - std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; - - return label_id_to_orig_id; -#endif -} -#endif - -inline std::vector loadTags(const std::string &tags_file, const std::string &base_file) -{ - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag; - if (tags_enabled) - { - size_t tag_file_ndims, tag_file_npts; - std::uint32_t *tag_data; - diskann::load_bin(tags_file, tag_data, tag_file_npts, tag_file_ndims); - if (tag_file_ndims != 1) - { - diskann::cerr << "tags file error" << std::endl; - throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // check if the point count match - size_t base_file_npts, base_file_ndims; - diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims); - if (base_file_npts != tag_file_npts) - { - diskann::cerr << "point num in tags file mismatch" << std::endl; - throw diskann::ANNException("point num in tags file mismatch", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - location_to_tag.assign(tag_data, tag_data + tag_file_npts); - delete[] tag_data; - } - return location_to_tag; -} - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_data_store.h b/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_data_store.h deleted file mode 100644 index 0a0a617..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_data_store.h +++ /dev/null @@ -1,89 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. -#pragma once - -#include -#include - -#include "tsl/robin_map.h" -#include "tsl/robin_set.h" -#include "tsl/sparse_map.h" -// #include "boost/dynamic_bitset.hpp" - -#include "abstract_data_store.h" - -#include "distance.h" -#include "natural_number_map.h" -#include "natural_number_set.h" -#include "aligned_file_reader.h" - -namespace diskann -{ -template class InMemDataStore : public AbstractDataStore -{ - public: - InMemDataStore(const location_t capacity, const size_t dim, std::unique_ptr> distance_fn); - virtual ~InMemDataStore(); - - virtual location_t load(const std::string &filename) override; - virtual size_t save(const std::string &filename, const location_t num_points) override; - - virtual size_t get_aligned_dim() const override; - - // Populate internal data from unaligned data while doing alignment and any - // normalization that is required. - virtual void populate_data(const data_t *vectors, const location_t num_pts) override; - virtual void populate_data(const std::string &filename, const size_t offset) override; - - virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override; - - virtual void get_vector(const location_t i, data_t *target) const override; - virtual void set_vector(const location_t i, const data_t *const vector) override; - virtual void prefetch_vector(const location_t loc) override; - - virtual void move_vectors(const location_t old_location_start, const location_t new_location_start, - const location_t num_points) override; - virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override; - - virtual void preprocess_query(const data_t *query, AbstractScratch *query_scratch) const override; - - virtual float get_distance(const data_t *preprocessed_query, const location_t loc) const override; - virtual float get_distance(const location_t loc1, const location_t loc2) const override; - - virtual void get_distance(const data_t *preprocessed_query, const location_t *locations, - const uint32_t location_count, float *distances, - AbstractScratch *scratch) const override; - virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const override; - - virtual location_t calculate_medoid() const override; - - virtual Distance *get_dist_fn() const override; - - virtual size_t get_alignment_factor() const override; - - protected: - virtual location_t expand(const location_t new_size) override; - virtual location_t shrink(const location_t new_size) override; - - virtual location_t load_impl(const std::string &filename); -#ifdef EXEC_ENV_OLS - virtual location_t load_impl(AlignedFileReader &reader); -#endif - - private: - data_t *_data = nullptr; - - size_t _aligned_dim; - - // It may seem weird to put distance metric along with the data store class, - // but this gives us perf benefits as the datastore can do distance - // computations during search and compute norms of vectors internally without - // have to copy data back and forth. - std::unique_ptr> _distance_fn; - - // in case we need to save vector norms for optimization - std::shared_ptr _pre_computed_norms; -}; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_graph_store.h b/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_graph_store.h deleted file mode 100644 index d0206a7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/in_mem_graph_store.h +++ /dev/null @@ -1,51 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include "abstract_graph_store.h" - -namespace diskann -{ - -class InMemGraphStore : public AbstractGraphStore -{ - public: - InMemGraphStore(const size_t total_pts, const size_t reserve_graph_degree); - - // returns tuple of - virtual std::tuple load(const std::string &index_path_prefix, - const size_t num_points) override; - virtual int store(const std::string &index_path_prefix, const size_t num_points, const size_t num_frozen_points, - const uint32_t start) override; - - virtual const std::vector &get_neighbours(const location_t i) const override; - virtual void add_neighbour(const location_t i, location_t neighbour_id) override; - virtual void clear_neighbours(const location_t i) override; - virtual void swap_neighbours(const location_t a, location_t b) override; - - virtual void set_neighbours(const location_t i, std::vector &neighbors) override; - - virtual size_t resize_graph(const size_t new_size) override; - virtual void clear_graph() override; - - virtual size_t get_max_range_of_graph() override; - virtual uint32_t get_max_observed_degree() override; - - protected: - virtual std::tuple load_impl(const std::string &filename, size_t expected_num_points); -#ifdef EXEC_ENV_OLS - virtual std::tuple load_impl(AlignedFileReader &reader, size_t expected_num_points); -#endif - - int save_graph(const std::string &index_path_prefix, const size_t active_points, const size_t num_frozen_points, - const uint32_t start); - - private: - size_t _max_range_of_graph = 0; - uint32_t _max_observed_degree = 0; - - std::vector> _graph; -}; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/index.h b/packages/leann-backend-diskann/third_party/DiskANN/include/index.h deleted file mode 100644 index c4303a1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/index.h +++ /dev/null @@ -1,452 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include "common_includes.h" - -#ifdef EXEC_ENV_OLS -#include "aligned_file_reader.h" -#endif - -#include "distance.h" -#include "locking.h" -#include "natural_number_map.h" -#include "natural_number_set.h" -#include "neighbor.h" -#include "parameters.h" -#include "utils.h" -#include "windows_customizations.h" -#include "scratch.h" -#include "in_mem_data_store.h" -#include "in_mem_graph_store.h" -#include "abstract_index.h" - -#include "quantized_distance.h" -#include "pq_data_store.h" - -#define OVERHEAD_FACTOR 1.1 -#define EXPAND_IF_FULL 0 -#define DEFAULT_MAXC 750 - -namespace diskann -{ - -inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, uint32_t degree) -{ - double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; - double size_of_graph = ((double)size) * degree * sizeof(uint32_t) * defaults::GRAPH_SLACK_FACTOR; - double size_of_locks = ((double)size) * sizeof(non_recursive_mutex); - double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t); - - return OVERHEAD_FACTOR * (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector); -} - -template class Index : public AbstractIndex -{ - /************************************************************************** - * - * Public functions acquire one or more of _update_lock, _consolidate_lock, - * _tag_lock, _delete_lock before calling protected functions which DO NOT - * acquire these locks. They might acquire locks on _locks[i] - * - **************************************************************************/ - - public: - // Constructor for Bulk operations and for creating the index object solely - // for loading a prexisting index. - DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::shared_ptr> data_store, - std::unique_ptr graph_store, - std::shared_ptr> pq_data_store = nullptr); - - // Constructor for incremental index - DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, - const std::shared_ptr index_parameters, - const std::shared_ptr index_search_params, - const size_t num_frozen_pts = 0, const bool dynamic_index = false, - const bool enable_tags = false, const bool concurrent_consolidate = false, - const bool pq_dist_build = false, const size_t num_pq_chunks = 0, - const bool use_opq = false, const bool filtered_index = false); - - DISKANN_DLLEXPORT ~Index(); - - // Saves graph, data, metadata and associated tags. - DISKANN_DLLEXPORT void save(const char *filename, bool compact_before_save = false) override; - - // Load functions -#ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l); -#else - // Reads the number of frozen points from graph's metadata file section. - DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points(const std::string &graph_file); - - DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l) override; -#endif - - // get some private variables - DISKANN_DLLEXPORT size_t get_num_points(); - DISKANN_DLLEXPORT size_t get_max_points(); - - DISKANN_DLLEXPORT bool detect_common_filters(uint32_t point_id, bool search_invocation, - const std::vector &incoming_labels); - - // Batch build from a file. Optionally pass tags vector. - DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, - const std::vector &tags = std::vector()); - - // Batch build from a file. Optionally pass tags file. - DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, const char *tag_filename); - - // Batch build from a data array, which must pad vectors to aligned_dim - DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, const std::vector &tags); - - // Based on filter params builds a filtered or unfiltered index - DISKANN_DLLEXPORT void build(const std::string &data_file, const size_t num_points_to_load, - IndexFilterParams &filter_params) override; - - // Filtered Support - DISKANN_DLLEXPORT void build_filtered_index(const char *filename, const std::string &label_file, - const size_t num_points_to_load, - const std::vector &tags = std::vector()); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - // Get converted integer label from string to int map (_label_map) - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label); - - // Set starting point of an index before inserting any points incrementally. - // The data count should be equal to _num_frozen_pts * _aligned_dim. - DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count); - // Set starting points to random points on a sphere of certain radius. - // A fixed random seed can be specified for scenarios where it's important - // to have higher consistency between index builds. - DISKANN_DLLEXPORT void set_start_points_at_random(T radius, uint32_t random_seed = 0); - - // For FastL2 search on a static index, we interleave the data with graph - DISKANN_DLLEXPORT void optimize_index_layout() override; - - // For FastL2 search on optimized layout - DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices); - - // Added search overload that takes L as parameter, so that we - // can customize L on a per-query basis without tampering with "Parameters" - template - DISKANN_DLLEXPORT std::pair search(const T *query, const size_t K, const uint32_t L, - IDType *indices, float *distances = nullptr); - - // Initialize space for res_vectors before calling. - DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, - float *distances, std::vector &res_vectors, bool use_filters = false, - const std::string filter_label = ""); - - // Filter support search - template - DISKANN_DLLEXPORT std::pair search_with_filters(const T *query, const LabelT &filter_label, - const size_t K, const uint32_t L, - IndexType *indices, float *distances); - - // Will fail if tag already in the index or if tag=0. - DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag); - - // Will fail if tag already in the index or if tag=0. - DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag, const std::vector &label); - - // call this before issuing deletions to sets relevant flags - DISKANN_DLLEXPORT int enable_delete(); - - // Record deleted point now and restructure graph later. Return -1 if tag - // not found, 0 if OK. - DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); - - // Record deleted points now and restructure graph later. Add to failed_tags - // if tag not found. - DISKANN_DLLEXPORT void lazy_delete(const std::vector &tags, std::vector &failed_tags); - - // Call after a series of lazy deletions - // Returns number of live points left after consolidation - // If _conc_consolidates is set in the ctor, then this call can be invoked - // alongside inserts and lazy deletes, else it acquires _update_lock - DISKANN_DLLEXPORT consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters) override; - - DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion, - const float alpha); - - DISKANN_DLLEXPORT bool is_index_saved(); - - // repositions frozen points to the end of _data - if they have been moved - // during deletion - DISKANN_DLLEXPORT void reposition_frozen_point_to_end(); - DISKANN_DLLEXPORT void reposition_points(uint32_t old_location_start, uint32_t new_location_start, - uint32_t num_locations); - - // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); - - DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); - - // memory should be allocated for vec before calling this function - DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); - - DISKANN_DLLEXPORT void print_status(); - - DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); - - // This variable MUST be updated if the number of entries in the metadata - // change. - DISKANN_DLLEXPORT static const int METADATA_ROWS = 5; - - DISKANN_DLLEXPORT void get_degree_stats(size_t &max_deg, size_t &min_deg, size_t &avg_deg, size_t &cnt_deg); - - DISKANN_DLLEXPORT void dump_degree_stats(std::string filename); - - // ******************************** - // - // Internals of the library - // - // ******************************** - - protected: - // overload of abstract index virtual methods - virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) override; - - virtual std::pair _search(const DataType &query, const size_t K, const uint32_t L, - std::any &indices, float *distances = nullptr) override; - virtual std::pair _search_with_filters(const DataType &query, - const std::string &filter_label_raw, const size_t K, - const uint32_t L, std::any &indices, - float *distances) override; - - virtual int _insert_point(const DataType &data_point, const TagType tag) override; - virtual int _insert_point(const DataType &data_point, const TagType tag, Labelvector &labels) override; - - virtual int _lazy_delete(const TagType &tag) override; - - virtual void _lazy_delete(TagVector &tags, TagVector &failed_tags) override; - - virtual void _get_active_tags(TagRobinSet &active_tags) override; - - virtual void _set_start_points_at_random(DataType radius, uint32_t random_seed = 0) override; - - virtual int _get_vector_by_tag(TagType &tag, DataType &vec) override; - - virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) override; - - virtual size_t _search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags, - float *distances, DataVector &res_vectors, bool use_filters = false, - const std::string filter_label = "") override; - - virtual void _set_universal_label(const LabelType universal_label) override; - - // No copy/assign. - Index(const Index &) = delete; - Index &operator=(const Index &) = delete; - - // Use after _data and _nd have been populated - // Acquire exclusive _update_lock before calling - void build_with_data_populated(const std::vector &tags); - - // generates 1 frozen point that will never be deleted from the graph - // This is not visible to the user - void generate_frozen_point(); - - // determines navigating node of the graph by calculating medoid of datafopt - uint32_t calculate_entry_point(); - - void parse_label_file(const std::string &label_file, size_t &num_pts_labels); - - std::unordered_map load_label_map(const std::string &map_file); - - // Returns the locations of start point and frozen points suitable for use - // with iterate_to_fixed_point. - std::vector get_init_ids(); - - // The query to use is placed in scratch->aligned_query - std::pair iterate_to_fixed_point(InMemQueryScratch *scratch, const uint32_t Lindex, - const std::vector &init_ids, bool use_filter, - const std::vector &filters, bool search_invocation); - - void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter = false, - uint32_t filteredLindex = 0); - - void prune_neighbors(const uint32_t location, std::vector &pool, std::vector &pruned_list, - InMemQueryScratch *scratch); - - void prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, - const uint32_t max_candidate_size, const float alpha, std::vector &pruned_list, - InMemQueryScratch *scratch); - - // Prunes candidates in @pool to a shorter list @result - // @pool must be sorted before calling - void occlude_list(const uint32_t location, std::vector &pool, const float alpha, const uint32_t degree, - const uint32_t maxc, std::vector &result, InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr = nullptr); - - // add reverse links from all the visited nodes to node n. - void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, - InMemQueryScratch *scratch); - - void inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch); - - // Acquire exclusive _update_lock before calling - void link(); - - // Acquire exclusive _tag_lock and _delete_lock before calling - int reserve_location(); - - // Acquire exclusive _tag_lock before calling - size_t release_location(int location); - size_t release_locations(const tsl::robin_set &locations); - - // Resize the index when no slots are left for insertion. - // Acquire exclusive _update_lock and _tag_lock before calling. - void resize(size_t new_max_points); - - // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock - // and _delete_lock before calling these functions. - // Renumber nodes, update tag and location maps and compact the - // graph, mode = _consolidated_order in case of lazy deletion and - // _compacted_order in case of eager deletion - DISKANN_DLLEXPORT void compact_data(); - DISKANN_DLLEXPORT void compact_frozen_point(); - - // Remove deleted nodes from adjacency list of node loc - // Replace removed neighbors with second order neighbors. - // Also acquires _locks[i] for i = loc and out-neighbors of loc. - void process_delete(const tsl::robin_set &old_delete_set, size_t loc, const uint32_t range, - const uint32_t maxc, const float alpha, InMemQueryScratch *scratch); - - void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, uint32_t r, - uint32_t maxc, size_t dim); - - // Do not call without acquiring appropriate locks - // call public member functions save and load to invoke these. - DISKANN_DLLEXPORT size_t save_graph(std::string filename); - DISKANN_DLLEXPORT size_t save_data(std::string filename); - DISKANN_DLLEXPORT size_t save_tags(std::string filename); - DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); -#ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_tags(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_delete_set(AlignedFileReader &reader); -#else - DISKANN_DLLEXPORT size_t load_graph(const std::string filename, size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(std::string filename0); - DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name); - DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename); -#endif - - private: - // Distance functions - Metric _dist_metric = diskann::L2; - - // Data - std::shared_ptr> _data_store; - - // Graph related data structures - std::unique_ptr _graph_store; - - char *_opt_graph = nullptr; - - // Dimensions - size_t _dim = 0; - size_t _nd = 0; // number of active points i.e. existing in the graph - size_t _max_points = 0; // total number of points in given data set - - // _num_frozen_pts is the number of points which are used as initial - // candidates when iterating to closest point(s). These are not visible - // externally and won't be returned by search. At least 1 frozen point is - // needed for a dynamic index. The frozen points have consecutive locations. - // See also _start below. - size_t _num_frozen_pts = 0; - size_t _frozen_pts_used = 0; - size_t _node_size; - size_t _data_len; - size_t _neighbor_len; - - // Start point of the search. When _num_frozen_pts is greater than zero, - // this is the location of the first frozen point. Otherwise, this is a - // location of one of the points in index. - uint32_t _start = 0; - - bool _has_built = false; - bool _saturate_graph = false; - bool _save_as_one_file = false; // plan to support in next version - bool _dynamic_index = false; - bool _enable_tags = false; - bool _normalize_vecs = false; // Using normalied L2 for cosine. - bool _deletes_enabled = false; - - // Filter Support - - bool _filtered_index = false; - // Location to label is only updated during insert_point(), all other reads are protected by - // default as a location can only be released at end of consolidate deletes - std::vector> _location_to_labels; - tsl::robin_set _labels; - std::string _labels_file; - std::unordered_map _label_to_start_id; - std::unordered_map _medoid_counts; - - bool _use_universal_label = false; - LabelT _universal_label = 0; - uint32_t _filterIndexingQueueSize; - std::unordered_map _label_map; - - // Indexing parameters - uint32_t _indexingQueueSize; - uint32_t _indexingRange; - uint32_t _indexingMaxC; - float _indexingAlpha; - uint32_t _indexingThreads; - - // Query scratch data structures - ConcurrentQueue *> _query_scratch; - - // Flags for PQ based distance calculation - bool _pq_dist = false; - bool _use_opq = false; - size_t _num_pq_chunks = 0; - // REFACTOR - // uint8_t *_pq_data = nullptr; - std::shared_ptr> _pq_distance_fn = nullptr; - std::shared_ptr> _pq_data_store = nullptr; - bool _pq_generated = false; - FixedChunkPQTable _pq_table; - - // - // Data structures, locks and flags for dynamic indexing and tags - // - - // lazy_delete removes entry from _location_to_tag and _tag_to_location. If - // _location_to_tag does not resolve a location, infer that it was deleted. - tsl::sparse_map _tag_to_location; - natural_number_map _location_to_tag; - - // _empty_slots has unallocated slots and those freed by consolidate_delete. - // _delete_set has locations marked deleted by lazy_delete. Will not be - // immediately available for insert. consolidate_delete will release these - // slots to _empty_slots. - natural_number_set _empty_slots; - std::unique_ptr> _delete_set; - - bool _data_compacted = true; // true if data has been compacted - bool _is_saved = false; // Checking if the index is already saved. - bool _conc_consolidate = false; // use _lock while searching - - // Acquire locks in the order below when acquiring multiple locks - std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and - _update_lock; // search/inserts/deletes/consolidate (shared lock) - std::shared_timed_mutex // Ensure only one consolidate or compact_data is - _consolidate_lock; // ever active - std::shared_timed_mutex // RW lock for _tag_to_location, - _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points, _label_to_start_id - std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted - _delete_lock; // variable - - // Per node lock, cardinality=_max_points + _num_frozen_points - std::vector _locks; - - static const float INDEX_GROWTH_FACTOR; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/index_build_params.h b/packages/leann-backend-diskann/third_party/DiskANN/include/index_build_params.h deleted file mode 100644 index d4f4548..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/index_build_params.h +++ /dev/null @@ -1,73 +0,0 @@ -#pragma once - -#include "common_includes.h" -#include "parameters.h" - -namespace diskann -{ -struct IndexFilterParams -{ - public: - std::string save_path_prefix; - std::string label_file; - std::string tags_file; - std::string universal_label; - uint32_t filter_threshold = 0; - - private: - IndexFilterParams(const std::string &save_path_prefix, const std::string &label_file, - const std::string &universal_label, uint32_t filter_threshold) - : save_path_prefix(save_path_prefix), label_file(label_file), universal_label(universal_label), - filter_threshold(filter_threshold) - { - } - - friend class IndexFilterParamsBuilder; -}; -class IndexFilterParamsBuilder -{ - public: - IndexFilterParamsBuilder() = default; - - IndexFilterParamsBuilder &with_save_path_prefix(const std::string &save_path_prefix) - { - if (save_path_prefix.empty() || save_path_prefix == "") - throw ANNException("Error: save_path_prefix can't be empty", -1); - this->_save_path_prefix = save_path_prefix; - return *this; - } - - IndexFilterParamsBuilder &with_label_file(const std::string &label_file) - { - this->_label_file = label_file; - return *this; - } - - IndexFilterParamsBuilder &with_universal_label(const std::string &univeral_label) - { - this->_universal_label = univeral_label; - return *this; - } - - IndexFilterParamsBuilder &with_filter_threshold(const std::uint32_t &filter_threshold) - { - this->_filter_threshold = filter_threshold; - return *this; - } - - IndexFilterParams build() - { - return IndexFilterParams(_save_path_prefix, _label_file, _universal_label, _filter_threshold); - } - - IndexFilterParamsBuilder(const IndexFilterParamsBuilder &) = delete; - IndexFilterParamsBuilder &operator=(const IndexFilterParamsBuilder &) = delete; - - private: - std::string _save_path_prefix; - std::string _label_file; - std::string _tags_file; - std::string _universal_label; - uint32_t _filter_threshold = 0; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/index_config.h b/packages/leann-backend-diskann/third_party/DiskANN/include/index_config.h deleted file mode 100644 index a8e64d0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/index_config.h +++ /dev/null @@ -1,256 +0,0 @@ -#pragma once - -#include "common_includes.h" -#include "parameters.h" - -namespace diskann -{ -enum class DataStoreStrategy -{ - MEMORY -}; - -enum class GraphStoreStrategy -{ - MEMORY -}; - -struct IndexConfig -{ - DataStoreStrategy data_strategy; - GraphStoreStrategy graph_strategy; - - Metric metric; - size_t dimension; - size_t max_points; - - bool dynamic_index; - bool enable_tags; - bool pq_dist_build; - bool concurrent_consolidate; - bool use_opq; - bool filtered_index; - - size_t num_pq_chunks; - size_t num_frozen_pts; - - std::string label_type; - std::string tag_type; - std::string data_type; - - // Params for building index - std::shared_ptr index_write_params; - // Params for searching index - std::shared_ptr index_search_params; - - private: - IndexConfig(DataStoreStrategy data_strategy, GraphStoreStrategy graph_strategy, Metric metric, size_t dimension, - size_t max_points, size_t num_pq_chunks, size_t num_frozen_points, bool dynamic_index, bool enable_tags, - bool pq_dist_build, bool concurrent_consolidate, bool use_opq, bool filtered_index, - std::string &data_type, const std::string &tag_type, const std::string &label_type, - std::shared_ptr index_write_params, - std::shared_ptr index_search_params) - : data_strategy(data_strategy), graph_strategy(graph_strategy), metric(metric), dimension(dimension), - max_points(max_points), dynamic_index(dynamic_index), enable_tags(enable_tags), pq_dist_build(pq_dist_build), - concurrent_consolidate(concurrent_consolidate), use_opq(use_opq), filtered_index(filtered_index), - num_pq_chunks(num_pq_chunks), num_frozen_pts(num_frozen_points), label_type(label_type), tag_type(tag_type), - data_type(data_type), index_write_params(index_write_params), index_search_params(index_search_params) - { - } - - friend class IndexConfigBuilder; -}; - -class IndexConfigBuilder -{ - public: - IndexConfigBuilder() = default; - - IndexConfigBuilder &with_metric(Metric m) - { - this->_metric = m; - return *this; - } - - IndexConfigBuilder &with_graph_load_store_strategy(GraphStoreStrategy graph_strategy) - { - this->_graph_strategy = graph_strategy; - return *this; - } - - IndexConfigBuilder &with_data_load_store_strategy(DataStoreStrategy data_strategy) - { - this->_data_strategy = data_strategy; - return *this; - } - - IndexConfigBuilder &with_dimension(size_t dimension) - { - this->_dimension = dimension; - return *this; - } - - IndexConfigBuilder &with_max_points(size_t max_points) - { - this->_max_points = max_points; - return *this; - } - - IndexConfigBuilder &is_dynamic_index(bool dynamic_index) - { - this->_dynamic_index = dynamic_index; - return *this; - } - - IndexConfigBuilder &is_enable_tags(bool enable_tags) - { - this->_enable_tags = enable_tags; - return *this; - } - - IndexConfigBuilder &is_pq_dist_build(bool pq_dist_build) - { - this->_pq_dist_build = pq_dist_build; - return *this; - } - - IndexConfigBuilder &is_concurrent_consolidate(bool concurrent_consolidate) - { - this->_concurrent_consolidate = concurrent_consolidate; - return *this; - } - - IndexConfigBuilder &is_use_opq(bool use_opq) - { - this->_use_opq = use_opq; - return *this; - } - - IndexConfigBuilder &is_filtered(bool is_filtered) - { - this->_filtered_index = is_filtered; - return *this; - } - - IndexConfigBuilder &with_num_pq_chunks(size_t num_pq_chunks) - { - this->_num_pq_chunks = num_pq_chunks; - return *this; - } - - IndexConfigBuilder &with_num_frozen_pts(size_t num_frozen_pts) - { - this->_num_frozen_pts = num_frozen_pts; - return *this; - } - - IndexConfigBuilder &with_label_type(const std::string &label_type) - { - this->_label_type = label_type; - return *this; - } - - IndexConfigBuilder &with_tag_type(const std::string &tag_type) - { - this->_tag_type = tag_type; - return *this; - } - - IndexConfigBuilder &with_data_type(const std::string &data_type) - { - this->_data_type = data_type; - return *this; - } - - IndexConfigBuilder &with_index_write_params(IndexWriteParameters &index_write_params) - { - this->_index_write_params = std::make_shared(index_write_params); - return *this; - } - - IndexConfigBuilder &with_index_write_params(std::shared_ptr index_write_params_ptr) - { - if (index_write_params_ptr == nullptr) - { - diskann::cout << "Passed, empty build_params while creating index config" << std::endl; - return *this; - } - this->_index_write_params = index_write_params_ptr; - return *this; - } - - IndexConfigBuilder &with_index_search_params(IndexSearchParams &search_params) - { - this->_index_search_params = std::make_shared(search_params); - return *this; - } - - IndexConfigBuilder &with_index_search_params(std::shared_ptr search_params_ptr) - { - if (search_params_ptr == nullptr) - { - diskann::cout << "Passed, empty search_params while creating index config" << std::endl; - return *this; - } - this->_index_search_params = search_params_ptr; - return *this; - } - - IndexConfig build() - { - if (_data_type == "" || _data_type.empty()) - throw ANNException("Error: data_type can not be empty", -1); - - if (_dynamic_index && _num_frozen_pts == 0) - { - _num_frozen_pts = 1; - } - - if (_dynamic_index) - { - if (_index_search_params != nullptr && _index_search_params->initial_search_list_size == 0) - throw ANNException("Error: please pass initial_search_list_size for building dynamic index.", -1); - } - - // sanity check - if (_dynamic_index && _num_frozen_pts == 0) - { - diskann::cout << "_num_frozen_pts passed as 0 for dynamic_index. Setting it to 1 for safety." << std::endl; - _num_frozen_pts = 1; - } - - return IndexConfig(_data_strategy, _graph_strategy, _metric, _dimension, _max_points, _num_pq_chunks, - _num_frozen_pts, _dynamic_index, _enable_tags, _pq_dist_build, _concurrent_consolidate, - _use_opq, _filtered_index, _data_type, _tag_type, _label_type, _index_write_params, - _index_search_params); - } - - IndexConfigBuilder(const IndexConfigBuilder &) = delete; - IndexConfigBuilder &operator=(const IndexConfigBuilder &) = delete; - - private: - DataStoreStrategy _data_strategy; - GraphStoreStrategy _graph_strategy; - - Metric _metric; - size_t _dimension; - size_t _max_points; - - bool _dynamic_index = false; - bool _enable_tags = false; - bool _pq_dist_build = false; - bool _concurrent_consolidate = false; - bool _use_opq = false; - bool _filtered_index{defaults::HAS_LABELS}; - - size_t _num_pq_chunks = 0; - size_t _num_frozen_pts{defaults::NUM_FROZEN_POINTS_STATIC}; - - std::string _label_type{"uint32"}; - std::string _tag_type{"uint32"}; - std::string _data_type; - - std::shared_ptr _index_write_params; - std::shared_ptr _index_search_params; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/index_factory.h b/packages/leann-backend-diskann/third_party/DiskANN/include/index_factory.h deleted file mode 100644 index 76fb0b9..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/index_factory.h +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include "index.h" -#include "abstract_graph_store.h" -#include "in_mem_graph_store.h" -#include "pq_data_store.h" - -namespace diskann -{ -class IndexFactory -{ - public: - DISKANN_DLLEXPORT explicit IndexFactory(const IndexConfig &config); - DISKANN_DLLEXPORT std::unique_ptr create_instance(); - - DISKANN_DLLEXPORT static std::unique_ptr construct_graphstore( - const GraphStoreStrategy stratagy, const size_t size, const size_t reserve_graph_degree); - - template - DISKANN_DLLEXPORT static std::shared_ptr> construct_datastore(DataStoreStrategy stratagy, - size_t num_points, - size_t dimension, Metric m); - // For now PQDataStore incorporates within itself all variants of quantization that we support. In the - // future it may be necessary to introduce an AbstractPQDataStore class to spearate various quantization - // flavours. - template - DISKANN_DLLEXPORT static std::shared_ptr> construct_pq_datastore(DataStoreStrategy strategy, - size_t num_points, size_t dimension, - Metric m, size_t num_pq_chunks, - bool use_opq); - template static Distance *construct_inmem_distance_fn(Metric m); - - private: - void check_config(); - - template - std::unique_ptr create_instance(); - - std::unique_ptr create_instance(const std::string &data_type, const std::string &tag_type, - const std::string &label_type); - - template - std::unique_ptr create_instance(const std::string &tag_type, const std::string &label_type); - - template - std::unique_ptr create_instance(const std::string &label_type); - - std::unique_ptr _config; -}; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/linux_aligned_file_reader.h b/packages/leann-backend-diskann/third_party/DiskANN/include/linux_aligned_file_reader.h deleted file mode 100644 index d1d1e74..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/linux_aligned_file_reader.h +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#ifndef _WINDOWS -#ifndef __APPLE__ - -#include "aligned_file_reader.h" - -class LinuxAlignedFileReader : public AlignedFileReader -{ - private: - uint64_t file_sz; - FileHandle file_desc; - io_context_t bad_ctx = (io_context_t)-1; - - public: - LinuxAlignedFileReader(); - ~LinuxAlignedFileReader(); - - IOContext &get_ctx(); - - // register thread-id for a context - void register_thread(); - - // de-register thread-id for a context - void deregister_thread(); - void deregister_all_threads(); - - // Open & close ops - // Blocking calls - void open(const std::string &fname); - void close(); - - // process batch of aligned requests in parallel - // NOTE :: blocking call - void read(std::vector &read_reqs, IOContext &ctx, bool async = false); -}; - -#endif -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/locking.h b/packages/leann-backend-diskann/third_party/DiskANN/include/locking.h deleted file mode 100644 index 890c24a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/locking.h +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. -#pragma once - -#include - -#ifdef _WINDOWS -#include "windows_slim_lock.h" -#endif - -namespace diskann -{ -#ifdef _WINDOWS -using non_recursive_mutex = windows_exclusive_slim_lock; -using LockGuard = windows_exclusive_slim_lock_guard; -#else -using non_recursive_mutex = std::mutex; -using LockGuard = std::lock_guard; -#endif -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/logger.h b/packages/leann-backend-diskann/third_party/DiskANN/include/logger.h deleted file mode 100644 index 0b17807..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/logger.h +++ /dev/null @@ -1,35 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. -#pragma once - -#include -#include -#include "windows_customizations.h" - -#ifdef EXEC_ENV_OLS -#ifndef ENABLE_CUSTOM_LOGGER -#define ENABLE_CUSTOM_LOGGER -#endif // !ENABLE_CUSTOM_LOGGER -#endif // EXEC_ENV_OLS - -namespace diskann -{ -#ifdef ENABLE_CUSTOM_LOGGER -DISKANN_DLLEXPORT extern std::basic_ostream cout; -DISKANN_DLLEXPORT extern std::basic_ostream cerr; -#else -using std::cerr; -using std::cout; -#endif - -enum class DISKANN_DLLEXPORT LogLevel -{ - LL_Info = 0, - LL_Error, - LL_Count -}; - -#ifdef ENABLE_CUSTOM_LOGGER -DISKANN_DLLEXPORT void SetCustomLogger(std::function logger); -#endif -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/logger_impl.h b/packages/leann-backend-diskann/third_party/DiskANN/include/logger_impl.h deleted file mode 100644 index 03c65e0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/logger_impl.h +++ /dev/null @@ -1,61 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include "ann_exception.h" -#include "logger.h" - -namespace diskann -{ -#ifdef ENABLE_CUSTOM_LOGGER -class ANNStreamBuf : public std::basic_streambuf -{ - public: - DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp); - DISKANN_DLLEXPORT ~ANNStreamBuf(); - - DISKANN_DLLEXPORT bool is_open() const - { - return true; // because stdout and stderr are always open. - } - DISKANN_DLLEXPORT void close(); - DISKANN_DLLEXPORT virtual int underflow(); - DISKANN_DLLEXPORT virtual int overflow(int c); - DISKANN_DLLEXPORT virtual int sync(); - - private: - FILE *_fp; - char *_buf; - int _bufIndex; - std::mutex _mutex; - LogLevel _logLevel; - - int flush(); - void logImpl(char *str, int numchars); - - // Why the two buffer-sizes? If we are running normally, we are basically - // interacting with a character output system, so we short-circuit the - // output process by keeping an empty buffer and writing each character - // to stdout/stderr. But if we are running in OLS, we have to take all - // the text that is written to diskann::cout/diskann:cerr, consolidate it - // and push it out in one-shot, because the OLS infra does not give us - // character based output. Therefore, we use a larger buffer that is large - // enough to store the longest message, and continuously add characters - // to it. When the calling code outputs a std::endl or std::flush, sync() - // will be called and will output a log level, component name, and the text - // that has been collected. (sync() is also called if the buffer is full, so - // overflows/missing text are not a concern). - // This implies calling code _must_ either print std::endl or std::flush - // to ensure that the message is written immediately. - - static const int BUFFER_SIZE = 1024; - - ANNStreamBuf(const ANNStreamBuf &); - ANNStreamBuf &operator=(const ANNStreamBuf &); -}; -#endif -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/math_utils.h b/packages/leann-backend-diskann/third_party/DiskANN/include/math_utils.h deleted file mode 100644 index 83d189f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/math_utils.h +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include "common_includes.h" -#include "utils.h" - -namespace math_utils -{ - -float calc_distance(float *vec_1, float *vec_2, size_t dim); - -// compute l2-squared norms of data stored in row major num_points * dim, -// needs -// to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim); - -void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, - bool transpose_rot = false); - -// calculate closest center to data of num_points * dim (row major) -// centers is num_centers * dim (row major) -// data_l2sq has pre-computed squared norms of data -// centers_l2sq has pre-computed squared norms of centers -// pre-allocated center_index will contain id of k nearest centers -// pre-allocated dist_matrix shound be num_points * num_centers and contain -// squared distances - -// Ideally used only by compute_closest_centers -void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k = 1); - -// Given data in num_points * new_dim row major -// Pivots stored in full_pivot_data as k * new_dim row major -// Calculate the closest pivot for each point and store it in vector -// closest_centers_ivf (which needs to be allocated outside) -// Additionally, if inverted index is not null (and pre-allocated), it will -// return inverted index for each center Additionally, if pts_norms_squared is -// not null, then it will assume that point norms are pre-computed and use -// those -// values - -void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, - size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index = NULL, - float *pts_norms_squared = NULL); - -// if to_subtract is 1, will subtract nearest center from each row. Else will -// add. Output will be in data_load iself. -// Nearest centers need to be provided in closst_centers. - -void process_residuals(float *data_load, size_t num_points, size_t dim, float *cur_pivot_data, size_t num_centers, - uint32_t *closest_centers, bool to_subtract); - -} // namespace math_utils - -namespace kmeans -{ - -// run Lloyds one iteration -// Given data in row major num_points * dim, and centers in row major -// num_centers * dim -// And squared lengths of data points, output the closest center to each data -// point, update centers, and also return inverted index. -// If closest_centers == NULL, will allocate memory and return. -// Similarly, if closest_docs == NULL, will allocate memory and return. - -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, - std::vector *closest_docs, uint32_t *&closest_center); - -// Run Lloyds until max_reps or stopping criterion -// If you pass NULL for closest_docs and closest_center, it will NOT return -// the results, else it will assume appriate allocation as closest_docs = new -// vector [num_centers], and closest_center = new size_t[num_points] -// Final centers are output in centers as row major num_centers * dim -// -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, - const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center); - -// assumes already memory allocated for pivot_data as new -// float[num_centers*dim] and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); - -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); -} // namespace kmeans diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/memory_mapper.h b/packages/leann-backend-diskann/third_party/DiskANN/include/memory_mapper.h deleted file mode 100644 index 75faca1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/memory_mapper.h +++ /dev/null @@ -1,43 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#ifndef _WINDOWS -#include -#include -#include -#include -#include - -#else -#include -#endif -#include - -namespace diskann -{ -class MemoryMapper -{ - private: -#ifndef _WINDOWS - int _fd; -#else - HANDLE _bareFile; - HANDLE _fd; - -#endif - char *_buf; - size_t _fileSize; - const char *_fileName; - - public: - MemoryMapper(const char *filename); - MemoryMapper(const std::string &filename); - - char *getBuf(); - size_t getFileSize(); - - ~MemoryMapper(); -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_map.h b/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_map.h deleted file mode 100644 index e846882..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_map.h +++ /dev/null @@ -1,86 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include - -#include - -namespace diskann -{ -// A map whose key is a natural number (from 0 onwards) and maps to a value. -// Made as both memory and performance efficient map for scenario such as -// DiskANN location-to-tag map. There, the pool of numbers is consecutive from -// zero to some max value, and it's expected that most if not all keys from 0 -// up to some current maximum will be present in the map. The memory usage of -// the map is determined by the largest inserted key since it uses vector as a -// backing store and bitset for presence indication. -// -// Thread-safety: this class is not thread-safe in general. -// Exception: multiple read-only operations are safe on the object only if -// there are no writers to it in parallel. -template class natural_number_map -{ - public: - static_assert(std::is_trivial::value, "Key must be a trivial type"); - - // Represents a reference to a element in the map. Used while iterating - // over map entries. - struct position - { - size_t _key; - // The number of keys that were enumerated when iterating through the - // map so far. Used to early-terminate enumeration when ithere are no - // more entries in the map. - size_t _keys_already_enumerated; - - // Returns whether it's valid to access the element at this position in - // the map. - bool is_valid() const; - }; - - natural_number_map(); - - void reserve(size_t count); - size_t size() const; - - void set(Key key, Value value); - void erase(Key key); - - bool contains(Key key) const; - bool try_get(Key key, Value &value) const; - - // Returns the value at the specified position. Prerequisite: position is - // valid. - Value get(const position &pos) const; - - // Finds the first element in the map, if any. Invalidated by changes in the - // map. - position find_first() const; - - // Finds the next element in the map after the specified position. - // Invalidated by changes in the map. - position find_next(const position &after_position) const; - - void clear(); - - private: - // Number of entries in the map. Not the same as size() of the - // _values_vector below. - size_t _size; - - // Array of values. The key is the index of the value. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_set.h b/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_set.h deleted file mode 100644 index ec5b827..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/natural_number_set.h +++ /dev/null @@ -1,50 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include "boost_dynamic_bitset_fwd.h" - -namespace diskann -{ -// A set of natural numbers (from 0 onwards). Made for scenario where the -// pool of numbers is consecutive from zero to some max value and very -// efficient methods for "add to set", "get any value from set", "is in set" -// are needed. The memory usage of the set is determined by the largest -// number of inserted entries (uses a vector as a backing store) as well as -// the largest value to be placed in it (uses bitset as well). -// -// Thread-safety: this class is not thread-safe in general. -// Exception: multiple read-only operations (e.g. is_in_set, empty, size) are -// safe on the object only if there are no writers to it in parallel. -template class natural_number_set -{ - public: - static_assert(std::is_trivial::value, "Identifier must be a trivial type"); - - natural_number_set(); - - bool is_empty() const; - void reserve(size_t count); - void insert(T id); - T pop_any(); - void clear(); - size_t size() const; - bool is_in_set(T id) const; - - private: - // Values that are currently in set. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/neighbor.h b/packages/leann-backend-diskann/third_party/DiskANN/include/neighbor.h deleted file mode 100644 index d7c0c25..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/neighbor.h +++ /dev/null @@ -1,152 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include "utils.h" - -namespace diskann -{ - -struct Neighbor -{ - unsigned id; - float distance; - bool expanded; - - Neighbor() = default; - - Neighbor(unsigned id, float distance) : id{id}, distance{distance}, expanded(false) - { - } - - inline bool operator<(const Neighbor &other) const - { - return distance < other.distance || (distance == other.distance && id < other.id); - } - - inline bool operator==(const Neighbor &other) const - { - return (id == other.id); - } -}; - -// Invariant: after every `insert` and `closest_unexpanded()`, `_cur` points to -// the first Neighbor which is unexpanded. -class NeighborPriorityQueue -{ - public: - NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0) - { - } - - explicit NeighborPriorityQueue(size_t capacity) : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1) - { - } - - // Inserts the item ordered into the set up to the sets capacity. - // The item will be dropped if it is the same id as an exiting - // set item or it has a greated distance than the final - // item in the set. The set cursor that is used to pop() the - // next item will be set to the lowest index of an uncheck item - void insert(const Neighbor &nbr) - { - if (_size == _capacity && _data[_size - 1] < nbr) - { - return; - } - - size_t lo = 0, hi = _size; - while (lo < hi) - { - size_t mid = (lo + hi) >> 1; - if (nbr < _data[mid]) - { - hi = mid; - // Make sure the same id isn't inserted into the set - } - else if (_data[mid].id == nbr.id) - { - return; - } - else - { - lo = mid + 1; - } - } - - if (lo < _capacity) - { - std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor)); - } - _data[lo] = {nbr.id, nbr.distance}; - if (_size < _capacity) - { - _size++; - } - if (lo < _cur) - { - _cur = lo; - } - } - - Neighbor closest_unexpanded() - { - _data[_cur].expanded = true; - size_t pre = _cur; - while (_cur < _size && _data[_cur].expanded) - { - _cur++; - } - return _data[pre]; - } - - bool has_unexpanded_node() const - { - return _cur < _size; - } - - size_t size() const - { - return _size; - } - - size_t capacity() const - { - return _capacity; - } - - void reserve(size_t capacity) - { - if (capacity + 1 > _data.size()) - { - _data.resize(capacity + 1); - } - _capacity = capacity; - } - - Neighbor &operator[](size_t i) - { - return _data[i]; - } - - Neighbor operator[](size_t i) const - { - return _data[i]; - } - - void clear() - { - _size = 0; - _cur = 0; - } - - private: - size_t _size, _capacity, _cur; - std::vector _data; -}; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/parameters.h b/packages/leann-backend-diskann/third_party/DiskANN/include/parameters.h deleted file mode 100644 index 0206814..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/parameters.h +++ /dev/null @@ -1,119 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include - -#include "omp.h" -#include "defaults.h" - -namespace diskann -{ - -class IndexWriteParameters - -{ - public: - const uint32_t search_list_size; // L - const uint32_t max_degree; // R - const bool saturate_graph; - const uint32_t max_occlusion_size; // C - const float alpha; - const uint32_t num_threads; - const uint32_t filter_list_size; // Lf - - IndexWriteParameters(const uint32_t search_list_size, const uint32_t max_degree, const bool saturate_graph, - const uint32_t max_occlusion_size, const float alpha, const uint32_t num_threads, - const uint32_t filter_list_size) - : search_list_size(search_list_size), max_degree(max_degree), saturate_graph(saturate_graph), - max_occlusion_size(max_occlusion_size), alpha(alpha), num_threads(num_threads), - filter_list_size(filter_list_size) - { - } - - friend class IndexWriteParametersBuilder; -}; - -class IndexSearchParams -{ - public: - IndexSearchParams(const uint32_t initial_search_list_size, const uint32_t num_search_threads) - : initial_search_list_size(initial_search_list_size), num_search_threads(num_search_threads) - { - } - const uint32_t initial_search_list_size; // search L - const uint32_t num_search_threads; // search threads -}; - -class IndexWriteParametersBuilder -{ - /** - * Fluent builder pattern to keep track of the 7 non-default properties - * and their order. The basic ctor was getting unwieldy. - */ - public: - IndexWriteParametersBuilder(const uint32_t search_list_size, // L - const uint32_t max_degree // R - ) - : _search_list_size(search_list_size), _max_degree(max_degree) - { - } - - IndexWriteParametersBuilder &with_max_occlusion_size(const uint32_t max_occlusion_size) - { - _max_occlusion_size = max_occlusion_size; - return *this; - } - - IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph) - { - _saturate_graph = saturate_graph; - return *this; - } - - IndexWriteParametersBuilder &with_alpha(const float alpha) - { - _alpha = alpha; - return *this; - } - - IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads) - { - _num_threads = num_threads == 0 ? omp_get_num_procs() : num_threads; - return *this; - } - - IndexWriteParametersBuilder &with_filter_list_size(const uint32_t filter_list_size) - { - _filter_list_size = filter_list_size == 0 ? _search_list_size : filter_list_size; - return *this; - } - - IndexWriteParameters build() const - { - return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, _max_occlusion_size, _alpha, - _num_threads, _filter_list_size); - } - - IndexWriteParametersBuilder(const IndexWriteParameters &wp) - : _search_list_size(wp.search_list_size), _max_degree(wp.max_degree), - _max_occlusion_size(wp.max_occlusion_size), _saturate_graph(wp.saturate_graph), _alpha(wp.alpha), - _filter_list_size(wp.filter_list_size) - { - } - IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete; - IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = delete; - - private: - uint32_t _search_list_size{}; - uint32_t _max_degree{}; - uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE}; - bool _saturate_graph{defaults::SATURATE_GRAPH}; - float _alpha{defaults::ALPHA}; - uint32_t _num_threads{defaults::NUM_THREADS}; - uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE}; -}; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/partition.h b/packages/leann-backend-diskann/third_party/DiskANN/include/partition.h deleted file mode 100644 index c2c4c76..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/partition.h +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include -#include -#include -#include -#include - -#include "neighbor.h" -#include "parameters.h" -#include "tsl/robin_set.h" -#include "utils.h" - -#include "windows_customizations.h" - -template -void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate); - -template -void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, - size_t &ndims); - -template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, - size_t &slice_size); - -int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::vector &cluster_sizes); - -template -int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, - const size_t k_base, std::string prefix_path); - -template -int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::string prefix_path); - -template -int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename); - -template -int partition(const std::string data_file, const float sampling_rate, size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); - -template -int partition_with_ram_budget(const std::string data_file, const double sampling_rate, double ram_budget, - size_t graph_degree, const std::string prefix_path, size_t k_base); diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/percentile_stats.h b/packages/leann-backend-diskann/third_party/DiskANN/include/percentile_stats.h deleted file mode 100644 index 7932575..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/percentile_stats.h +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include -#ifdef _WINDOWS -#include -#endif -#include -#include - -#include "distance.h" -#include "parameters.h" - -namespace diskann -{ -struct QueryStats -{ - float total_us = 0; // total time to process query in micros - float io_us = 0; // total time spent in IO - float cpu_us = 0; // total time spent in CPU - - unsigned n_4k = 0; // # of 4kB reads - unsigned n_8k = 0; // # of 8kB reads - unsigned n_12k = 0; // # of 12kB reads - unsigned n_ios = 0; // total # of IOs issued - unsigned read_size = 0; // total # of bytes read - unsigned n_cmps_saved = 0; // # cmps saved - unsigned n_cmps = 0; // # cmps - unsigned n_cache_hits = 0; // # cache_hits - unsigned n_hops = 0; // # search hops -}; - -template -inline T get_percentile_stats(QueryStats *stats, uint64_t len, float percentile, - const std::function &member_fn) -{ - std::vector vals(len); - for (uint64_t i = 0; i < len; i++) - { - vals[i] = member_fn(stats[i]); - } - - std::sort(vals.begin(), vals.end(), [](const T &left, const T &right) { return left < right; }); - - auto retval = vals[(uint64_t)(percentile * len)]; - vals.clear(); - return retval; -} - -template -inline double get_mean_stats(QueryStats *stats, uint64_t len, const std::function &member_fn) -{ - double avg = 0; - for (uint64_t i = 0; i < len; i++) - { - avg += (double)member_fn(stats[i]); - } - return avg / len; -} -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq.h deleted file mode 100644 index 3e6119f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq.h +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include "utils.h" -#include "pq_common.h" - -namespace diskann -{ -class FixedChunkPQTable -{ - float *tables = nullptr; // pq_tables = float array of size [256 * ndims] - uint64_t ndims = 0; // ndims = true dimension of vectors - uint64_t n_chunks = 0; - bool use_rotation = false; - uint32_t *chunk_offsets = nullptr; - float *centroid = nullptr; - float *tables_tr = nullptr; // same as pq_tables, but col-major - float *rotmat_tr = nullptr; - - public: - FixedChunkPQTable(); - - virtual ~FixedChunkPQTable(); - -#ifdef EXEC_ENV_OLS - void load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks); -#else - void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); -#endif - - uint32_t get_num_chunks(); - - void preprocess_query(float *query_vec); - - // assumes pre-processed query - void populate_chunk_distances(const float *query_vec, float *dist_vec); - - float l2_distance(const float *query_vec, uint8_t *base_vec); - - float inner_product(const float *query_vec, uint8_t *base_vec); - - // assumes no rotation is involved - void inflate_vector(uint8_t *base_vec, float *out_vec); - - void populate_chunk_inner_products(const float *query_vec, float *dist_vec); -}; - -void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); - -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - std::vector &dists_out); - -// Need to replace calls to these with calls to vector& based functions above -void aggregate_coords(const unsigned *ids, const uint64_t n_ids, const uint8_t *all_coords, const uint64_t ndims, - uint8_t *out); - -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - float *dists_out); - -DISKANN_DLLEXPORT int generate_pq_pivots(const float *const train_data, size_t num_train, unsigned dim, - unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, - std::string pq_pivots_path, bool make_zero_mean = false); - -DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers, - unsigned num_pq_chunks, std::string opq_pivots_path, - bool make_zero_mean = false); - -DISKANN_DLLEXPORT int generate_pq_pivots_simplified(const float *train_data, size_t num_train, size_t dim, - size_t num_pq_chunks, std::vector &pivot_data_vector); - -template -int generate_pq_data_from_pivots(const std::string &data_file, unsigned num_centers, unsigned num_pq_chunks, - const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, - bool use_opq = false); - -DISKANN_DLLEXPORT int generate_pq_data_from_pivots_simplified(const float *data, const size_t num, - const float *pivot_data, const size_t pivots_num, - const size_t dim, const size_t num_pq_chunks, - std::vector &pq); - -template -void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - const diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); - -template -void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, const diskann::Metric compareMetric, - const double p_val, const uint64_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix = ""); -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_common.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq_common.h deleted file mode 100644 index c6a3a57..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_common.h +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include -#include - -#define NUM_PQ_BITS 8 -#define NUM_PQ_CENTROIDS (1 << NUM_PQ_BITS) -#define MAX_OPQ_ITERS 20 -#define NUM_KMEANS_REPS_PQ 12 -#define MAX_PQ_TRAINING_SET_SIZE 256000 -#define MAX_PQ_CHUNKS 512 - -namespace diskann -{ -inline std::string get_quantized_vectors_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) -{ - return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_compressed.bin"; -} - -inline std::string get_pivot_data_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) -{ - return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_pivots.bin"; -} - -inline std::string get_rotation_matrix_suffix(const std::string &pivot_data_filename) -{ - return pivot_data_filename + "_rotation_matrix.bin"; -} - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_data_store.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq_data_store.h deleted file mode 100644 index 7c0cb5f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_data_store.h +++ /dev/null @@ -1,97 +0,0 @@ -#pragma once -#include -#include "distance.h" -#include "quantized_distance.h" -#include "pq.h" -#include "abstract_data_store.h" - -namespace diskann -{ -// REFACTOR TODO: By default, the PQDataStore is an in-memory datastore because both Vamana and -// DiskANN treat it the same way. But with DiskPQ, that may need to change. -template class PQDataStore : public AbstractDataStore -{ - - public: - PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, - std::unique_ptr> pq_distance_fn); - PQDataStore(const PQDataStore &) = delete; - PQDataStore &operator=(const PQDataStore &) = delete; - ~PQDataStore(); - - // Load quantized vectors from a set of files. Here filename is treated - // as a prefix and the files are assumed to be named with DiskANN - // conventions. - virtual location_t load(const std::string &file_prefix) override; - - // Save quantized vectors to a set of files whose names start with - // file_prefix. - // Currently, the plan is to save the quantized vectors to the quantized - // vectors file. - virtual size_t save(const std::string &file_prefix, const location_t num_points) override; - - // Since base class function is pure virtual, we need to declare it here, even though alignent concept is not needed - // for Quantized data stores. - virtual size_t get_aligned_dim() const override; - - // Populate quantized data from unaligned data using PQ functionality - virtual void populate_data(const data_t *vectors, const location_t num_pts) override; - virtual void populate_data(const std::string &filename, const size_t offset) override; - - virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override; - - virtual void get_vector(const location_t i, data_t *target) const override; - virtual void set_vector(const location_t i, const data_t *const vector) override; - virtual void prefetch_vector(const location_t loc) override; - - virtual void move_vectors(const location_t old_location_start, const location_t new_location_start, - const location_t num_points) override; - virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override; - - virtual void preprocess_query(const data_t *query, AbstractScratch *scratch) const override; - - virtual float get_distance(const data_t *query, const location_t loc) const override; - virtual float get_distance(const location_t loc1, const location_t loc2) const override; - - // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling - // this function. - virtual void get_distance(const data_t *preprocessed_query, const location_t *locations, - const uint32_t location_count, float *distances, - AbstractScratch *scratch_space) const override; - - // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling - // this function. - virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const override; - - // We are returning the distance function that is used for full precision - // vectors here, not the PQ distance function. This is because the callers - // all are expecting a Distance not QuantizedDistance. - virtual Distance *get_dist_fn() const override; - - virtual location_t calculate_medoid() const override; - - virtual size_t get_alignment_factor() const override; - - protected: - virtual location_t expand(const location_t new_size) override; - virtual location_t shrink(const location_t new_size) override; - - virtual location_t load_impl(const std::string &filename); -#ifdef EXEC_ENV_OLS - virtual location_t load_impl(AlignedFileReader &reader); -#endif - - private: - uint8_t *_quantized_data = nullptr; - size_t _num_chunks = 0; - - // REFACTOR TODO: Doing this temporarily before refactoring OPQ into - // its own class. Remove later. - bool _use_opq = false; - - Metric _distance_metric; - std::unique_ptr> _distance_fn = nullptr; - std::unique_ptr> _pq_distance_fn = nullptr; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_flash_index.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq_flash_index.h deleted file mode 100644 index 174df5c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_flash_index.h +++ /dev/null @@ -1,286 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#include "common_includes.h" - -#include "aligned_file_reader.h" -#include "concurrent_queue.h" -#include "neighbor.h" -#include "parameters.h" -#include "percentile_stats.h" -#include "pq.h" -#include "utils.h" -#include "windows_customizations.h" -#include "scratch.h" -#include "tsl/robin_map.h" -#include "tsl/robin_set.h" - -#define FULL_PRECISION_REORDER_MULTIPLIER 3 - -namespace diskann -{ - -template class PQFlashIndex -{ - public: - DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr &fileReader, - std::shared_ptr &graphReader, - diskann::Metric metric = diskann::Metric::L2); - DISKANN_DLLEXPORT ~PQFlashIndex(); - -#ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix, - const char *pq_prefix = nullptr); -#else - // load compressed data, and obtains the handle to the disk-resident index - DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix, const char *pq_prefix = nullptr, - const char *partition_prefix = nullptr); -#endif - -#ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath, const char *graph_file); -#else - DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath, - const char *graph_file, const char *partition_file); -#endif - - DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); - -#ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, - uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list); -#else - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, - uint32_t num_threads, - std::vector &node_list); -#endif - - DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, - const bool shuffle = false); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_reorder_data = false, QueryStats *stats = nullptr, - const bool USE_DEFERRED_FETCH = false, - const bool skip_search_reorder = false, - const bool recompute_beighbor_embeddings = false, - const bool dedup_node_dis = false, float prune_ratio = 0, - const bool batch_recompute = false, bool global_pruning = false); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data = false, QueryStats *stats = nullptr, - const bool USE_DEFERRED_FETCH = false, - const bool skip_search_reorder = false, - const bool recompute_beighbor_embeddings = false, - const bool dedup_node_dis = false, float prune_ratio = 0, - const bool batch_recompute = false, bool global_pruning = false); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr, const bool USE_DEFERRED_FETCH = false, - const bool skip_search_reorder = false, - const bool recompute_beighbor_embeddings = false, - const bool dedup_node_dis = false, float prune_ratio = 0, - const bool batch_recompute = false, bool global_pruning = false); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr, const bool USE_DEFERRED_FETCH = false, - const bool skip_search_reorder = false, - const bool recompute_beighbor_embeddings = false, - const bool dedup_node_dis = false, float prune_ratio = 0, - const bool batch_recompute = false, bool global_pruning = false); - - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); - - DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, const uint64_t min_l_search, - const uint64_t max_l_search, std::vector &indices, - std::vector &distances, const uint64_t min_beam_width, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT uint64_t get_data_dim(); - - std::shared_ptr &reader; - - DISKANN_DLLEXPORT diskann::Metric get_metric(); - - // - // node_ids: input list of node_ids to be read - // coord_buffers: pointers to pre-allocated buffers that coords need to copied to. If null, dont copy. - // nbr_buffers: pre-allocated buffers to copy neighbors into - // - // returns a vector of bool one for each node_id: true if read is success, else false - // - DISKANN_DLLEXPORT std::vector read_nodes(const std::vector &node_ids, - std::vector &coord_buffers, - std::vector> &nbr_buffers); - - DISKANN_DLLEXPORT std::vector get_pq_vector(std::uint64_t vid); - DISKANN_DLLEXPORT uint64_t get_num_points(); - - protected: - DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); - DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - private: - DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, LabelT label_id); - std::unordered_map load_label_map(std::basic_istream &infile); - DISKANN_DLLEXPORT void parse_label_file(std::basic_istream &infile, size_t &num_pts_labels); - DISKANN_DLLEXPORT void get_label_file_metadata(const std::string &fileContent, uint32_t &num_pts, - uint32_t &num_total_labels); - DISKANN_DLLEXPORT void generate_random_labels(std::vector &labels, const uint32_t num_labels, - const uint32_t nthreads); - void reset_stream_for_reading(std::basic_istream &infile); - - // sector # on disk where node_id is present with in the graph part - DISKANN_DLLEXPORT uint64_t get_node_sector(uint64_t node_id); - - // ptr to start of the node - DISKANN_DLLEXPORT char *offset_to_node(char *sector_buf, uint64_t node_id); - - // returns region of `node_buf` containing [NNBRS][NBR_ID(uint32_t)] - DISKANN_DLLEXPORT uint32_t *offset_to_node_nhood(char *node_buf); - - // returns region of `node_buf` containing [COORD(T)] - DISKANN_DLLEXPORT T *offset_to_node_coords(char *node_buf); - - DISKANN_DLLEXPORT int load_graph_index(const std::string &graph_index_file); - - DISKANN_DLLEXPORT int read_partition_info(const std::string &partition_bin); - - DISKANN_DLLEXPORT int read_neighbors(const std::string &graph_index_file, uint64_t target_node_id); - - // index info for multi-node sectors - // nhood of node `i` is in sector: [i / nnodes_per_sector] - // offset in sector: [(i % nnodes_per_sector) * max_node_len] - // - // index info for multi-sector nodes - // nhood of node `i` is in sector: [i * DIV_ROUND_UP(_max_node_len, SECTOR_LEN)] - // offset in sector: [0] - // - // Common info - // coords start at ofsset - // #nbrs of node `i`: *(unsigned*) (offset + disk_bytes_per_point) - // nbrs of node `i` : (unsigned*) (offset + disk_bytes_per_point + 1) - - uint64_t _max_node_len = 0; - uint64_t _nnodes_per_sector = 0; // 0 for multi-sector nodes, >0 for multi-node sectors - uint64_t _max_degree = 0; - uint64_t _C = 0; - // Data used for searching with re-order vectors - uint64_t _ndims_reorder_vecs = 0; - uint64_t _reorder_data_start_sector = 0; - uint64_t _nvecs_per_sector = 0; - - diskann::Metric metric = diskann::Metric::L2; - - // used only for inner product search to re-scale the result value - // (due to the pre-processing of base during index build) - float _max_base_norm = 0.0f; - - // data info - uint64_t _num_points = 0; - uint64_t _num_frozen_points = 0; - uint64_t _frozen_location = 0; - uint64_t _data_dim = 0; - uint64_t _aligned_dim = 0; - uint64_t _disk_bytes_per_point = 0; // Number of bytes - - std::string _disk_index_file; - std::vector> _node_visit_counter; - - // PQ data - // _n_chunks = # of chunks ndims is split into - // data: char * _n_chunks - // chunk_size = chunk size of each dimension chunk - // pq_tables = float* [[2^8 * [chunk_size]] * _n_chunks] - uint8_t *data = nullptr; - uint64_t _n_chunks; - FixedChunkPQTable _pq_table; - - // distance comparator - std::shared_ptr> _dist_cmp; - std::shared_ptr> _dist_cmp_float; - - // for very large datasets: we use PQ even for the disk resident index - bool _use_disk_index_pq = false; - uint64_t _disk_pq_n_chunks = 0; - FixedChunkPQTable _disk_pq_table; - - // medoid/start info - - // graph has one entry point by default, - // we can optionally have multiple starting points - uint32_t *_medoids = nullptr; - // defaults to 1 - size_t _num_medoids; - // by default, it is empty. If there are multiple - // centroids, we pick the medoid corresponding to the - // closest centroid as the starting point of search - float *_centroid_data = nullptr; - - // nhood_cache; the uint32_t in nhood_Cache are offsets into nhood_cache_buf - unsigned *_nhood_cache_buf = nullptr; - tsl::robin_map> _nhood_cache; - - // coord_cache; The T* in coord_cache are offsets into coord_cache_buf - T *_coord_cache_buf = nullptr; - tsl::robin_map _coord_cache; - - // thread-specific scratch - ConcurrentQueue *> _thread_data; - uint64_t _max_nthreads; - bool _load_flag = false; - bool _count_visited_nodes = false; - bool _reorder_data_exists = false; - uint64_t _reoreder_data_offset = 0; - - // filter support - uint32_t *_pts_to_label_offsets = nullptr; - uint32_t *_pts_to_label_counts = nullptr; - LabelT *_pts_to_labels = nullptr; - std::unordered_map> _filter_to_medoid_ids; - bool _use_universal_label = false; - LabelT _universal_filter_label; - tsl::robin_set _dummy_pts; - tsl::robin_set _has_dummy_pts; - tsl::robin_map _dummy_to_real_map; - tsl::robin_map> _real_to_dummy_map; - std::unordered_map _label_map; - - private: - bool _use_partition = false; - - std::shared_ptr graph_reader; // Graph file reader - std::string _graph_index_file; // Graph file path - uint64_t _graph_node_len; // Graph node length - uint64_t _emb_node_len; // Embedding node length - - // Partition related data structures - uint64_t _num_partitions; // Number of partitions - std::vector> _graph_partitions; // Partition information - std::vector _id2partition; // ID to partition mapping - -#ifdef EXEC_ENV_OLS - // Set to a larger value than the actual header to accommodate - // any additions we make to the header. This is an outer limit - // on how big the header can be. - static const int HEADER_SIZE = defaults::SECTOR_LEN; - char *getHeaderBytes(); -#endif -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_l2_distance.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq_l2_distance.h deleted file mode 100644 index e6fc6e4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_l2_distance.h +++ /dev/null @@ -1,87 +0,0 @@ -#pragma once -#include "quantized_distance.h" - -namespace diskann -{ -template class PQL2Distance : public QuantizedDistance -{ - public: - // REFACTOR TODO: We could take a file prefix here and load the - // PQ pivots file, so that the distance object is initialized - // immediately after construction. But this would not work well - // with our data store concept where the store is created first - // and data populated after. - // REFACTOR TODO: Ideally, we should only read the num_chunks from - // the pivots file. However, we read the pivots file only later, but - // clients can call functions like get__filename without calling - // load_pivot_data. Hence this. The TODO is whether we should check - // that the num_chunks from the file is the same as this one. - - PQL2Distance(uint32_t num_chunks, bool use_opq = false); - - virtual ~PQL2Distance() override; - - virtual bool is_opq() const override; - - virtual std::string get_quantized_vectors_filename(const std::string &prefix) const override; - virtual std::string get_pivot_data_filename(const std::string &prefix) const override; - virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const override; - -#ifdef EXEC_ENV_OLS - virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, - size_t num_chunks) override; -#else - virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) override; -#endif - - // Number of chunks in the PQ table. Depends on the compression level used. - // Has to be < ndim - virtual uint32_t get_num_chunks() const override; - - // Preprocess the query by computing chunk distances from the query vector to - // various centroids. Since we don't want this class to do scratch management, - // we will take a PQScratch object which can come either from Index class or - // PQFlashIndex class. - virtual void preprocess_query(const data_t *aligned_query, uint32_t original_dim, - PQScratch &pq_scratch) override; - - // Distance function used for graph traversal. This function must be called - // after - // preprocess_query. The reason we do not call preprocess ourselves is because - // that function has to be called once per query, while this function is - // called at each iteration of the graph walk. NOTE: This function expects - // 1. the query to be preprocessed using preprocess_query() - // 2. the scratch object to contain the quantized vectors corresponding to ids - // in aligned_pq_coord_scratch. Done by calling aggregate_coords() - // - virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t id_count, - float *dists_out) override; - - // Same as above, but returns the distances in a vector instead of an array. - // Convenience function for index.cpp. - virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, - std::vector &dists_out) override; - - // Currently this function is required for DiskPQ. However, it too can be - // subsumed under preprocessed_distance if we add the appropriate scratch - // variables to PQScratch and initialize them in - // pq_flash_index.cpp::disk_iterate_to_fixed_point() - virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) override; - - protected: - // assumes pre-processed query - virtual void prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec); - - // assumes no rotation is involved - // virtual void inflate_vector(uint8_t *base_vec, float *out_vec); - - float *_tables = nullptr; // pq_tables = float array of size [256 * ndims] - uint64_t _ndims = 0; // ndims = true dimension of vectors - uint64_t _num_chunks = 0; - bool _is_opq = false; - uint32_t *_chunk_offsets = nullptr; - float *_centroid = nullptr; - float *_tables_tr = nullptr; // same as pq_tables, but col-major - float *_rotmat_tr = nullptr; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_scratch.h b/packages/leann-backend-diskann/third_party/DiskANN/include/pq_scratch.h deleted file mode 100644 index 95f1b13..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/pq_scratch.h +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once -#include -#include "pq_common.h" -#include "utils.h" - -namespace diskann -{ - -template class PQScratch -{ - public: - float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] - float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE - uint8_t *aligned_pq_coord_scratch = nullptr; // AT LEAST [N_CHUNKS * MAX_DEGREE] - float *rotated_query = nullptr; - float *aligned_query_float = nullptr; - - PQScratch(size_t graph_degree, size_t aligned_dim); - void initialize(size_t dim, const T *query, const float norm = 1.0f); - virtual ~PQScratch(); -}; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/program_options_utils.hpp b/packages/leann-backend-diskann/third_party/DiskANN/include/program_options_utils.hpp deleted file mode 100644 index 2be6059..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/program_options_utils.hpp +++ /dev/null @@ -1,81 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include - -namespace program_options_utils -{ -const std::string make_program_description(const char *executable_name, const char *description) -{ - return std::string("\n") - .append(description) - .append("\n\n") - .append("Usage: ") - .append(executable_name) - .append(" [OPTIONS]"); -} - -// Required parameters -const char *DATA_TYPE_DESCRIPTION = "data type, one of {int8, uint8, float} - float is single precision (32 bit)"; -const char *DISTANCE_FUNCTION_DESCRIPTION = - "distance function {l2, mips, fast_l2, cosine}. 'fast l2' and 'mips' only support data_type float"; -const char *INDEX_PATH_PREFIX_DESCRIPTION = "Path prefix to the index, e.g. '/mnt/data/my_ann_index'"; -const char *RESULT_PATH_DESCRIPTION = - "Path prefix for saving results of the queries, e.g. '/mnt/data/query_file_X.bin'"; -const char *QUERY_FILE_DESCRIPTION = "Query file in binary format, e.g. '/mnt/data/query_file_X.bin'"; -const char *NUMBER_OF_RESULTS_DESCRIPTION = "Number of neighbors to be returned (K in the DiskANN white paper)"; -const char *SEARCH_LIST_DESCRIPTION = - "Size of search list to use. This value is the number of neighbor/distance pairs to keep in memory at the same " - "time while performing a query. This can also be described as the size of the working set at query time. This " - "must be greater than or equal to the number of results/neighbors to return (K in the white paper). Corresponds " - "to L in the DiskANN white paper."; -const char *INPUT_DATA_PATH = "Input data file in bin format. This is the file you want to build the index over. " - "File format: Shape of the vector followed by the vector of embeddings as binary data."; - -// Optional parameters -const char *FILTER_LABEL_DESCRIPTION = - "Filter to use when running a query. 'filter_label' and 'query_filters_file' are mutually exclusive."; -const char *FILTERS_FILE_DESCRIPTION = - "Filter file for Queries for Filtered Search. File format is text with one filter per line. File must " - "have exactly one filter OR the same number of filters as there are queries in the 'query_file'."; -const char *LABEL_TYPE_DESCRIPTION = - "Storage type of Labels {uint/uint32, ushort/uint16}, default value is uint which will consume memory 4 bytes per " - "filter. 'uint' is an alias for 'uint32' and 'ushort' is an alias for 'uint16'."; -const char *GROUND_TRUTH_FILE_DESCRIPTION = - "ground truth file for the queryset"; // what's the format, what's the requirements? does it need to include an - // entry for every item or just a small subset? I have so many questions about - // this file -const char *NUMBER_THREADS_DESCRIPTION = "Number of threads used for building index. Defaults to number of logical " - "processor cores on your this machine returned by omp_get_num_procs()"; -const char *FAIL_IF_RECALL_BELOW = - "Value between 0 (inclusive) and 100 (exclusive) indicating the recall tolerance percentage threshold before " - "program fails with a non-zero exit code. The default value of 0 means that the program will complete " - "successfully with any recall value. A non-zero value indicates the floor for acceptable recall values. If the " - "calculated recall value is below this threshold then the program will write out the results but return a non-zero " - "exit code as a signal that the recall was not acceptable."; // does it continue running or die immediately? Will I - // still get my results even if the return code is -1? - -const char *NUMBER_OF_NODES_TO_CACHE = "Number of BFS nodes around medoid(s) to cache. Default value: 0"; -const char *BEAMWIDTH = "Beamwidth for search. Set 0 to optimize internally. Default value: 2"; -const char *MAX_BUILD_DEGREE = "Maximum graph degree"; -const char *GRAPH_BUILD_COMPLEXITY = - "Size of the search working set during build time. This is the numer of neighbor/distance pairs to keep in memory " - "while building the index. Higher value results in a higher quality graph but it will take more time to build the " - "graph."; -const char *GRAPH_BUILD_ALPHA = "Alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for " - "denser graphs with lower diameter"; -const char *BUIlD_GRAPH_PQ_BYTES = "Number of PQ bytes to build the index; 0 for full precision build"; -const char *USE_OPQ = "Use Optimized Product Quantization (OPQ)."; -const char *LABEL_FILE = "Input label file in txt format for Filtered Index build. The file should contain comma " - "separated filters for each node with each line corresponding to a graph node"; -const char *UNIVERSAL_LABEL = - "Universal label, Use only in conjunction with label file for filtered index build. If a " - "graph node has all the labels against it, we can assign a special universal filter to the " - "point instead of comma separated filters for that point. The universal label should be assigned to nodes " - "in the labels file instead of listing all labels for a node. DiskANN will not automatically assign a " - "universal label to a node."; -const char *FILTERED_LBUILD = "Build complexity for filtered points, higher value results in better graphs"; - -} // namespace program_options_utils diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/proto_embedding.h b/packages/leann-backend-diskann/third_party/DiskANN/include/proto_embedding.h deleted file mode 100644 index f17e225..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/proto_embedding.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "embedding.pb.h" - -// This header ensures that the protobuf files are included correctly -// and provides a namespace alias for convenience -namespace diskann { - namespace proto = protoembedding; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/quantized_distance.h b/packages/leann-backend-diskann/third_party/DiskANN/include/quantized_distance.h deleted file mode 100644 index cc4aea9..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/quantized_distance.h +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once -#include -#include -#include -#include "abstract_scratch.h" - -namespace diskann -{ -template class PQScratch; - -template class QuantizedDistance -{ - public: - QuantizedDistance() = default; - QuantizedDistance(const QuantizedDistance &) = delete; - QuantizedDistance &operator=(const QuantizedDistance &) = delete; - virtual ~QuantizedDistance() = default; - - virtual bool is_opq() const = 0; - virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0; - virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0; - virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const = 0; - - // Loading the PQ centroid table need not be part of the abstract class. - // However, we want to indicate that this function will change once we have a - // file reader hierarchy, so leave it here as-is. -#ifdef EXEC_ENV_OLS - virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file, size_t num_chunks) = 0; -#else - virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) = 0; -#endif - - // Number of chunks in the PQ table. Depends on the compression level used. - // Has to be < ndim - virtual uint32_t get_num_chunks() const = 0; - - // Preprocess the query by computing chunk distances from the query vector to - // various centroids. Since we don't want this class to do scratch management, - // we will take a PQScratch object which can come either from Index class or - // PQFlashIndex class. - virtual void preprocess_query(const data_t *query_vec, uint32_t query_dim, PQScratch &pq_scratch) = 0; - - // Workhorse - // This function must be called after preprocess_query - virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t id_count, float *dists_out) = 0; - - // Same as above, but convenience function for index.cpp. - virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, - std::vector &dists_out) = 0; - - // Currently this function is required for DiskPQ. However, it too can be subsumed - // under preprocessed_distance if we add the appropriate scratch variables to - // PQScratch and initialize them in pq_flash_index.cpp::disk_iterate_to_fixed_point() - virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) = 0; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/common.h b/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/common.h deleted file mode 100644 index b833963..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/common.h +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -namespace diskann -{ -// Constants -static const std::string VECTOR_KEY = "query", K_KEY = "k", INDICES_KEY = "indices", DISTANCES_KEY = "distances", - TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", ERROR_MESSAGE_KEY = "error", L_KEY = "Ls", - TIME_TAKEN_KEY = "time_taken_in_us", PARTITION_KEY = "partition", - UNKNOWN_ERROR = "unknown_error"; -const unsigned int DEFAULT_L = 100; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/search_wrapper.h b/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/search_wrapper.h deleted file mode 100644 index ebd067d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/search_wrapper.h +++ /dev/null @@ -1,140 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include - -#include -#include - -namespace diskann -{ -class SearchResult -{ - public: - SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, - const float *const distances, const std::string *const tags = nullptr, - const unsigned *const partitions = nullptr); - - const std::vector &get_indices() const - { - return _indices; - } - const std::vector &get_distances() const - { - return _distances; - } - bool tags_enabled() const - { - return _tags_enabled; - } - const std::vector &get_tags() const - { - return _tags; - } - bool partitions_enabled() const - { - return _partitions_enabled; - } - const std::vector &get_partitions() const - { - return _partitions; - } - unsigned get_time() const - { - return _search_time_in_ms; - } - - private: - unsigned int _K; - unsigned int _search_time_in_ms; - std::vector _indices; - std::vector _distances; - - bool _tags_enabled; - std::vector _tags; - - bool _partitions_enabled; - std::vector _partitions; -}; - -class SearchNotImplementedException : public std::logic_error -{ - private: - std::string _errormsg; - - public: - SearchNotImplementedException(const char *type) : std::logic_error("Not Implemented") - { - _errormsg = "Search with data type "; - _errormsg += std::string(type); - _errormsg += " not implemented : "; - _errormsg += __FUNCTION__; - } - - virtual const char *what() const throw() - { - return _errormsg.c_str(); - } -}; - -class BaseSearch -{ - public: - BaseSearch(const std::string &tagsFile = nullptr); - virtual SearchResult search(const float *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("float"); - } - virtual SearchResult search(const int8_t *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("int8_t"); - } - - virtual SearchResult search(const uint8_t *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("uint8_t"); - } - - void lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags); - - protected: - bool _tags_enabled; - std::vector _tags_str; -}; - -template class InMemorySearch : public BaseSearch -{ - public: - InMemorySearch(const std::string &baseFile, const std::string &indexFile, const std::string &tagsFile, Metric m, - uint32_t num_threads, uint32_t search_l); - virtual ~InMemorySearch(); - - SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); - - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; -}; - -template class PQFlashSearch : public BaseSearch -{ - public: - PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, const unsigned num_threads, - const std::string &tagsFile, Metric m); - virtual ~PQFlashSearch(); - - SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); - - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; - std::shared_ptr reader; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/server.h b/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/server.h deleted file mode 100644 index 1d75847..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/restapi/server.h +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -namespace diskann -{ -class Server -{ - public: - Server(web::uri &url, std::vector> &multi_searcher, - const std::string &typestring); - virtual ~Server(); - - pplx::task open(); - pplx::task close(); - - protected: - template void handle_post(web::http::http_request message); - - template - web::json::value toJsonArray(const std::vector &v, std::function valConverter); - web::json::value prepareResponse(const int64_t &queryId, const int k); - - template - void parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, - unsigned int &dimensions, unsigned &Ls); - - web::json::value idsToJsonArray(const diskann::SearchResult &result); - web::json::value distancesToJsonArray(const diskann::SearchResult &result); - web::json::value tagsToJsonArray(const diskann::SearchResult &result); - web::json::value partitionsToJsonArray(const diskann::SearchResult &result); - - SearchResult aggregate_results(const unsigned K, const std::vector &results); - - private: - bool _isDebug; - std::unique_ptr _listener; - const bool _multi_search; - std::vector> _multi_searcher; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/scratch.h b/packages/leann-backend-diskann/third_party/DiskANN/include/scratch.h deleted file mode 100644 index 2f43e33..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/scratch.h +++ /dev/null @@ -1,216 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include - -#include "boost_dynamic_bitset_fwd.h" -// #include "boost/dynamic_bitset.hpp" -#include "tsl/robin_set.h" -#include "tsl/robin_map.h" -#include "tsl/sparse_map.h" - -#include "aligned_file_reader.h" -#include "abstract_scratch.h" -#include "neighbor.h" -#include "defaults.h" -#include "concurrent_queue.h" - -namespace diskann -{ -template class PQScratch; - -// -// AbstractScratch space for in-memory index based search -// -template class InMemQueryScratch : public AbstractScratch -{ - public: - ~InMemQueryScratch(); - InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim, - size_t alignment_factor, bool init_pq_scratch = false); - void resize_for_new_L(uint32_t new_search_l); - void clear(); - - inline uint32_t get_L() - { - return _L; - } - inline uint32_t get_R() - { - return _R; - } - inline uint32_t get_maxc() - { - return _maxc; - } - inline T *aligned_query() - { - return this->_aligned_query_T; - } - inline PQScratch *pq_scratch() - { - return this->_pq_scratch; - } - inline std::vector &pool() - { - return _pool; - } - inline NeighborPriorityQueue &best_l_nodes() - { - return _best_l_nodes; - } - inline std::vector &occlude_factor() - { - return _occlude_factor; - } - inline tsl::robin_set &inserted_into_pool_rs() - { - return _inserted_into_pool_rs; - } - inline boost::dynamic_bitset<> &inserted_into_pool_bs() - { - return *_inserted_into_pool_bs; - } - inline std::vector &id_scratch() - { - return _id_scratch; - } - inline std::vector &dist_scratch() - { - return _dist_scratch; - } - inline tsl::robin_set &expanded_nodes_set() - { - return _expanded_nodes_set; - } - inline std::vector &expanded_nodes_vec() - { - return _expanded_nghrs_vec; - } - inline std::vector &occlude_list_output() - { - return _occlude_list_output; - } - - private: - uint32_t _L; - uint32_t _R; - uint32_t _maxc; - - // _pool stores all neighbors explored from best_L_nodes. - // Usually around L+R, but could be higher. - // Initialized to 3L+R for some slack, expands as needed. - std::vector _pool; - - // _best_l_nodes is reserved for storing best L entries - // Underlying storage is L+1 to support inserts - NeighborPriorityQueue _best_l_nodes; - - // _occlude_factor.size() >= pool.size() in occlude_list function - // _pool is clipped to maxc in occlude_list before affecting _occlude_factor - // _occlude_factor is initialized to maxc size - std::vector _occlude_factor; - - // Capacity initialized to 20L - tsl::robin_set _inserted_into_pool_rs; - - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - boost::dynamic_bitset<> *_inserted_into_pool_bs; - - // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - std::vector _id_scratch; - - // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - // _dist_scratch should be at least the size of id_scratch - std::vector _dist_scratch; - - // Buffers used in process delete, capacity increases as needed - tsl::robin_set _expanded_nodes_set; - std::vector _expanded_nghrs_vec; - std::vector _occlude_list_output; -}; - -// -// AbstractScratch space for SSD index based search -// - -template class SSDQueryScratch : public AbstractScratch -{ - public: - T *coord_scratch = nullptr; // MUST BE AT LEAST [sizeof(T) * data_dim] - - char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] - size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use - - tsl::robin_set visited; - NeighborPriorityQueue retset; - std::vector full_retset; - - SSDQueryScratch(size_t aligned_dim, size_t visited_reserve); - ~SSDQueryScratch(); - - void reset(); -}; - -template class SSDThreadData -{ - public: - SSDQueryScratch scratch; - IOContext ctx; - - SSDThreadData(size_t aligned_dim, size_t visited_reserve); - void clear(); -}; - -// -// Class to avoid the hassle of pushing and popping the query scratch. -// -template class ScratchStoreManager -{ - public: - ScratchStoreManager(ConcurrentQueue &query_scratch) : _scratch_pool(query_scratch) - { - _scratch = query_scratch.pop(); - while (_scratch == nullptr) - { - query_scratch.wait_for_push_notify(); - _scratch = query_scratch.pop(); - } - } - T *scratch_space() - { - return _scratch; - } - - ~ScratchStoreManager() - { - _scratch->clear(); - _scratch_pool.push(_scratch); - _scratch_pool.push_notify_all(); - } - - void destroy() - { - while (!_scratch_pool.empty()) - { - auto scratch = _scratch_pool.pop(); - while (scratch == nullptr) - { - _scratch_pool.wait_for_push_notify(); - scratch = _scratch_pool.pop(); - } - delete scratch; - } - } - - private: - T *_scratch; - ConcurrentQueue &_scratch_pool; - ScratchStoreManager(const ScratchStoreManager &); - ScratchStoreManager &operator=(const ScratchStoreManager &); -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/simd_utils.h b/packages/leann-backend-diskann/third_party/DiskANN/include/simd_utils.h deleted file mode 100644 index 4b07369..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/simd_utils.h +++ /dev/null @@ -1,106 +0,0 @@ -#pragma once - -#ifdef _WINDOWS -#include -#include -#include -#include -#else -#include -#endif - -namespace diskann -{ -static inline __m256 _mm256_mul_epi8(__m256i X) -{ - __m256i zero = _mm256_setzero_si256(); - - __m256i sign_x = _mm256_cmpgt_epi8(zero, X); - - __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); - __m256i xhi = _mm256_unpackhi_epi8(X, sign_x); - - return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, xlo), _mm256_madd_epi16(xhi, xhi))); -} - -static inline __m128 _mm_mulhi_epi8(__m128i X) -{ - __m128i zero = _mm_setzero_si128(); - __m128i sign_x = _mm_cmplt_epi8(X, zero); - __m128i xhi = _mm_unpackhi_epi8(X, sign_x); - - return _mm_cvtepi32_ps(_mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi))); -} - -static inline __m128 _mm_mulhi_epi8_shift32(__m128i X) -{ - __m128i zero = _mm_setzero_si128(); - X = _mm_srli_epi64(X, 32); - __m128i sign_x = _mm_cmplt_epi8(X, zero); - __m128i xhi = _mm_unpackhi_epi8(X, sign_x); - - return _mm_cvtepi32_ps(_mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi))); -} -static inline __m128 _mm_mul_epi8(__m128i X, __m128i Y) -{ - __m128i zero = _mm_setzero_si128(); - - __m128i sign_x = _mm_cmplt_epi8(X, zero); - __m128i sign_y = _mm_cmplt_epi8(Y, zero); - - __m128i xlo = _mm_unpacklo_epi8(X, sign_x); - __m128i xhi = _mm_unpackhi_epi8(X, sign_x); - __m128i ylo = _mm_unpacklo_epi8(Y, sign_y); - __m128i yhi = _mm_unpackhi_epi8(Y, sign_y); - - return _mm_cvtepi32_ps(_mm_add_epi32(_mm_madd_epi16(xlo, ylo), _mm_madd_epi16(xhi, yhi))); -} -static inline __m128 _mm_mul_epi8(__m128i X) -{ - __m128i zero = _mm_setzero_si128(); - __m128i sign_x = _mm_cmplt_epi8(X, zero); - __m128i xlo = _mm_unpacklo_epi8(X, sign_x); - __m128i xhi = _mm_unpackhi_epi8(X, sign_x); - - return _mm_cvtepi32_ps(_mm_add_epi32(_mm_madd_epi16(xlo, xlo), _mm_madd_epi16(xhi, xhi))); -} - -static inline __m128 _mm_mul32_pi8(__m128i X, __m128i Y) -{ - __m128i xlo = _mm_cvtepi8_epi16(X), ylo = _mm_cvtepi8_epi16(Y); - return _mm_cvtepi32_ps(_mm_unpacklo_epi32(_mm_madd_epi16(xlo, ylo), _mm_setzero_si128())); -} - -static inline __m256 _mm256_mul_epi8(__m256i X, __m256i Y) -{ - __m256i zero = _mm256_setzero_si256(); - - __m256i sign_x = _mm256_cmpgt_epi8(zero, X); - __m256i sign_y = _mm256_cmpgt_epi8(zero, Y); - - __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); - __m256i xhi = _mm256_unpackhi_epi8(X, sign_x); - __m256i ylo = _mm256_unpacklo_epi8(Y, sign_y); - __m256i yhi = _mm256_unpackhi_epi8(Y, sign_y); - - return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, ylo), _mm256_madd_epi16(xhi, yhi))); -} - -static inline __m256 _mm256_mul32_pi8(__m128i X, __m128i Y) -{ - __m256i xlo = _mm256_cvtepi8_epi16(X), ylo = _mm256_cvtepi8_epi16(Y); - return _mm256_blend_ps(_mm256_cvtepi32_ps(_mm256_madd_epi16(xlo, ylo)), _mm256_setzero_ps(), 252); -} - -static inline float _mm256_reduce_add_ps(__m256 x) -{ - /* ( x3+x7, x2+x6, x1+x5, x0+x4 ) */ - const __m128 x128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); - /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */ - const __m128 x64 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); - /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */ - const __m128 x32 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); - /* Conversion to float is a no-op on x86-64 */ - return _mm_cvtss_f32(x32); -} -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tag_uint128.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tag_uint128.h deleted file mode 100644 index 642de31..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tag_uint128.h +++ /dev/null @@ -1,68 +0,0 @@ -#pragma once -#include -#include - -namespace diskann -{ -#pragma pack(push, 1) - -struct tag_uint128 -{ - std::uint64_t _data1 = 0; - std::uint64_t _data2 = 0; - - bool operator==(const tag_uint128 &other) const - { - return _data1 == other._data1 && _data2 == other._data2; - } - - bool operator==(std::uint64_t other) const - { - return _data1 == other && _data2 == 0; - } - - tag_uint128 &operator=(const tag_uint128 &other) - { - _data1 = other._data1; - _data2 = other._data2; - - return *this; - } - - tag_uint128 &operator=(std::uint64_t other) - { - _data1 = other; - _data2 = 0; - - return *this; - } -}; - -#pragma pack(pop) -} // namespace diskann - -namespace std -{ -// Hash 128 input bits down to 64 bits of output. -// This is intended to be a reasonably good hash function. -inline std::uint64_t Hash128to64(const std::uint64_t &low, const std::uint64_t &high) -{ - // Murmur-inspired hashing. - const std::uint64_t kMul = 0x9ddfea08eb382d69ULL; - std::uint64_t a = (low ^ high) * kMul; - a ^= (a >> 47); - std::uint64_t b = (high ^ a) * kMul; - b ^= (b >> 47); - b *= kMul; - return b; -} - -template <> struct hash -{ - size_t operator()(const diskann::tag_uint128 &key) const noexcept - { - return Hash128to64(key._data1, key._data2); // map -0 to 0 - } -}; - -} // namespace std \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/timer.h b/packages/leann-backend-diskann/third_party/DiskANN/include/timer.h deleted file mode 100644 index 325edf3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/timer.h +++ /dev/null @@ -1,40 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. -#pragma once - -#include - -namespace diskann -{ -class Timer -{ - typedef std::chrono::high_resolution_clock _clock; - std::chrono::time_point<_clock> check_point; - - public: - Timer() : check_point(_clock::now()) - { - } - - void reset() - { - check_point = _clock::now(); - } - - long long elapsed() const - { - return std::chrono::duration_cast(_clock::now() - check_point).count(); - } - - float elapsed_seconds() const - { - return (float)elapsed() / 1000000.0f; - } - - std::string elapsed_seconds_for_step(const std::string &step) const - { - return std::string("Time for ") + step + std::string(": ") + std::to_string(elapsed_seconds()) + - std::string(" seconds"); - } -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/.clang-format b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/.clang-format deleted file mode 100644 index 9d15924..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/.clang-format +++ /dev/null @@ -1,2 +0,0 @@ -DisableFormat: true -SortIncludes: false diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_growth_policy.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_growth_policy.h deleted file mode 100644 index 6bfa9e5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_growth_policy.h +++ /dev/null @@ -1,330 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Tessil - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_ROBIN_GROWTH_POLICY_H -#define TSL_ROBIN_GROWTH_POLICY_H - - -#include -#include -#include -#include -#include -#include -#include -#include -#include - - -#ifndef tsl_assert -# ifdef TSL_DEBUG -# define tsl_assert(expr) assert(expr) -# else -# define tsl_assert(expr) (static_cast(0)) -# endif -#endif - - -/** - * If exceptions are enabled, throw the exception passed in parameter, otherwise call std::terminate. - */ -#ifndef TSL_THROW_OR_TERMINATE -# if (defined(__cpp_exceptions) || defined(__EXCEPTIONS) || (defined (_MSC_VER) && defined (_CPPUNWIND))) && !defined(TSL_NO_EXCEPTIONS) -# define TSL_THROW_OR_TERMINATE(ex, msg) throw ex(msg) -# else -# ifdef NDEBUG -# define TSL_THROW_OR_TERMINATE(ex, msg) std::terminate() -# else -# include -# define TSL_THROW_OR_TERMINATE(ex, msg) do { std::fprintf(stderr, msg); std::terminate(); } while(0) -# endif -# endif -#endif - - -#ifndef TSL_LIKELY -# if defined(__GNUC__) || defined(__clang__) -# define TSL_LIKELY(exp) (__builtin_expect(!!(exp), true)) -# else -# define TSL_LIKELY(exp) (exp) -# endif -#endif - - -namespace tsl { -namespace rh { - -/** - * Grow the hash table by a factor of GrowthFactor keeping the bucket count to a power of two. It allows - * the table to use a mask operation instead of a modulo operation to map a hash to a bucket. - * - * GrowthFactor must be a power of two >= 2. - */ -template -class power_of_two_growth_policy { -public: - /** - * Called on the hash table creation and on rehash. The number of buckets for the table is passed in parameter. - * This number is a minimum, the policy may update this value with a higher value if needed (but not lower). - * - * If 0 is given, min_bucket_count_in_out must still be 0 after the policy creation and - * bucket_for_hash must always return 0 in this case. - */ - explicit power_of_two_growth_policy(std::size_t& min_bucket_count_in_out) { - if(min_bucket_count_in_out > max_bucket_count()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - if(min_bucket_count_in_out > 0) { - min_bucket_count_in_out = round_up_to_power_of_two(min_bucket_count_in_out); - m_mask = min_bucket_count_in_out - 1; - } - else { - m_mask = 0; - } - } - - /** - * Return the bucket [0, bucket_count()) to which the hash belongs. - * If bucket_count() is 0, it must always return 0. - */ - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return hash & m_mask; - } - - /** - * Return the number of buckets that should be used on next growth. - */ - std::size_t next_bucket_count() const { - if((m_mask + 1) > max_bucket_count() / GrowthFactor) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - return (m_mask + 1) * GrowthFactor; - } - - /** - * Return the maximum number of buckets supported by the policy. - */ - std::size_t max_bucket_count() const { - // Largest power of two. - return ((std::numeric_limits::max)() / 2) + 1; - } - - /** - * Reset the growth policy as if it was created with a bucket count of 0. - * After a clear, the policy must always return 0 when bucket_for_hash is called. - */ - void clear() noexcept { - m_mask = 0; - } - -private: - static std::size_t round_up_to_power_of_two(std::size_t value) { - if(is_power_of_two(value)) { - return value; - } - - if(value == 0) { - return 1; - } - - --value; - for(std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) { - value |= value >> i; - } - - return value + 1; - } - - static constexpr bool is_power_of_two(std::size_t value) { - return value != 0 && (value & (value - 1)) == 0; - } - -protected: - static_assert(is_power_of_two(GrowthFactor) && GrowthFactor >= 2, "GrowthFactor must be a power of two >= 2."); - - std::size_t m_mask; -}; - - -/** - * Grow the hash table by GrowthFactor::num / GrowthFactor::den and use a modulo to map a hash - * to a bucket. Slower but it can be useful if you want a slower growth. - */ -template> -class mod_growth_policy { -public: - explicit mod_growth_policy(std::size_t& min_bucket_count_in_out) { - if(min_bucket_count_in_out > max_bucket_count()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - if(min_bucket_count_in_out > 0) { - m_mod = min_bucket_count_in_out; - } - else { - m_mod = 1; - } - } - - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return hash % m_mod; - } - - std::size_t next_bucket_count() const { - if(m_mod == max_bucket_count()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - const double next_bucket_count = std::ceil(double(m_mod) * REHASH_SIZE_MULTIPLICATION_FACTOR); - if(!std::isnormal(next_bucket_count)) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - if(next_bucket_count > double(max_bucket_count())) { - return max_bucket_count(); - } - else { - return std::size_t(next_bucket_count); - } - } - - std::size_t max_bucket_count() const { - return MAX_BUCKET_COUNT; - } - - void clear() noexcept { - m_mod = 1; - } - -private: - static constexpr double REHASH_SIZE_MULTIPLICATION_FACTOR = 1.0 * GrowthFactor::num / GrowthFactor::den; - static const std::size_t MAX_BUCKET_COUNT = - std::size_t(double( - (std::numeric_limits::max)() / REHASH_SIZE_MULTIPLICATION_FACTOR - )); - - static_assert(REHASH_SIZE_MULTIPLICATION_FACTOR >= 1.1, "Growth factor should be >= 1.1."); - - std::size_t m_mod; -}; - - - -namespace detail { - -static constexpr const std::array PRIMES = {{ - 1ul, 5ul, 17ul, 29ul, 37ul, 53ul, 67ul, 79ul, 97ul, 131ul, 193ul, 257ul, 389ul, 521ul, 769ul, 1031ul, - 1543ul, 2053ul, 3079ul, 6151ul, 12289ul, 24593ul, 49157ul, 98317ul, 196613ul, 393241ul, 786433ul, - 1572869ul, 3145739ul, 6291469ul, 12582917ul, 25165843ul, 50331653ul, 100663319ul, 201326611ul, - 402653189ul, 805306457ul, 1610612741ul, 3221225473ul, 4294967291ul -}}; - -template -static constexpr std::size_t mod(std::size_t hash) { return hash % PRIMES[IPrime]; } - -// MOD_PRIME[iprime](hash) returns hash % PRIMES[iprime]. This table allows for faster modulo as the -// compiler can optimize the modulo code better with a constant known at the compilation. -static constexpr const std::array MOD_PRIME = {{ - &mod<0>, &mod<1>, &mod<2>, &mod<3>, &mod<4>, &mod<5>, &mod<6>, &mod<7>, &mod<8>, &mod<9>, &mod<10>, - &mod<11>, &mod<12>, &mod<13>, &mod<14>, &mod<15>, &mod<16>, &mod<17>, &mod<18>, &mod<19>, &mod<20>, - &mod<21>, &mod<22>, &mod<23>, &mod<24>, &mod<25>, &mod<26>, &mod<27>, &mod<28>, &mod<29>, &mod<30>, - &mod<31>, &mod<32>, &mod<33>, &mod<34>, &mod<35>, &mod<36>, &mod<37> , &mod<38>, &mod<39> -}}; - -} - -/** - * Grow the hash table by using prime numbers as bucket count. Slower than tsl::rh::power_of_two_growth_policy in - * general but will probably distribute the values around better in the buckets with a poor hash function. - * - * To allow the compiler to optimize the modulo operation, a lookup table is used with constant primes numbers. - * - * With a switch the code would look like: - * \code - * switch(iprime) { // iprime is the current prime of the hash table - * case 0: hash % 5ul; - * break; - * case 1: hash % 17ul; - * break; - * case 2: hash % 29ul; - * break; - * ... - * } - * \endcode - * - * Due to the constant variable in the modulo the compiler is able to optimize the operation - * by a series of multiplications, substractions and shifts. - * - * The 'hash % 5' could become something like 'hash - (hash * 0xCCCCCCCD) >> 34) * 5' in a 64 bits environement. - */ -class prime_growth_policy { -public: - explicit prime_growth_policy(std::size_t& min_bucket_count_in_out) { - auto it_prime = std::lower_bound(detail::PRIMES.begin(), - detail::PRIMES.end(), min_bucket_count_in_out); - if(it_prime == detail::PRIMES.end()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - m_iprime = static_cast(std::distance(detail::PRIMES.begin(), it_prime)); - if(min_bucket_count_in_out > 0) { - min_bucket_count_in_out = *it_prime; - } - else { - min_bucket_count_in_out = 0; - } - } - - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return detail::MOD_PRIME[m_iprime](hash); - } - - std::size_t next_bucket_count() const { - if(m_iprime + 1 >= detail::PRIMES.size()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size."); - } - - return detail::PRIMES[m_iprime + 1]; - } - - std::size_t max_bucket_count() const { - return detail::PRIMES.back(); - } - - void clear() noexcept { - m_iprime = 0; - } - -private: - unsigned int m_iprime; - - static_assert((std::numeric_limits::max)() >= detail::PRIMES.size(), - "The type of m_iprime is not big enough."); -}; - -} -} - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_hash.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_hash.h deleted file mode 100644 index 5ecc962..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_hash.h +++ /dev/null @@ -1,1285 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Tessil - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_ROBIN_HASH_H -#define TSL_ROBIN_HASH_H - - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "robin_growth_policy.h" - - -namespace tsl { - -namespace detail_robin_hash { - -template -struct make_void { - using type = void; -}; - -template -struct has_is_transparent: std::false_type { -}; - -template -struct has_is_transparent::type>: std::true_type { -}; - -template -struct is_power_of_two_policy: std::false_type { -}; - -template -struct is_power_of_two_policy>: std::true_type { -}; - - - -using truncated_hash_type = std::uint_least32_t; - -/** - * Helper class that store a truncated hash if StoreHash is true and nothing otherwise. - */ -template -class bucket_entry_hash { -public: - bool bucket_hash_equal(std::size_t /*hash*/) const noexcept { - return true; - } - - truncated_hash_type truncated_hash() const noexcept { - return 0; - } - -protected: - void set_hash(truncated_hash_type /*hash*/) noexcept { - } -}; - -template<> -class bucket_entry_hash { -public: - bool bucket_hash_equal(std::size_t hash) const noexcept { - return m_hash == truncated_hash_type(hash); - } - - truncated_hash_type truncated_hash() const noexcept { - return m_hash; - } - -protected: - void set_hash(truncated_hash_type hash) noexcept { - m_hash = truncated_hash_type(hash); - } - -private: - truncated_hash_type m_hash; -}; - - -/** - * Each bucket entry has: - * - A value of type `ValueType`. - * - An integer to store how far the value of the bucket, if any, is from its ideal bucket - * (ex: if the current bucket 5 has the value 'foo' and `hash('foo') % nb_buckets` == 3, - * `dist_from_ideal_bucket()` will return 2 as the current value of the bucket is two - * buckets away from its ideal bucket) - * If there is no value in the bucket (i.e. `empty()` is true) `dist_from_ideal_bucket()` will be < 0. - * - A marker which tells us if the bucket is the last bucket of the bucket array (useful for the - * iterator of the hash table). - * - If `StoreHash` is true, 32 bits of the hash of the value, if any, are also stored in the bucket. - * If the size of the hash is more than 32 bits, it is truncated. We don't store the full hash - * as storing the hash is a potential opportunity to use the unused space due to the alignement - * of the bucket_entry structure. We can thus potentially store the hash without any extra space - * (which would not be possible with 64 bits of the hash). - */ -template -class bucket_entry: public bucket_entry_hash { - using bucket_hash = bucket_entry_hash; - -public: - using value_type = ValueType; - using distance_type = std::int_least16_t; - - - bucket_entry() noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), - m_last_bucket(false) - { - tsl_assert(empty()); - } - - bucket_entry(bool last_bucket) noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), - m_last_bucket(last_bucket) - { - tsl_assert(empty()); - } - - bucket_entry(const bucket_entry& other) noexcept(std::is_nothrow_copy_constructible::value): - bucket_hash(other), - m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), - m_last_bucket(other.m_last_bucket) - { - if(!other.empty()) { - ::new (static_cast(std::addressof(m_value))) value_type(other.value()); - m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; - } - } - - /** - * Never really used, but still necessary as we must call resize on an empty `std::vector`. - * and we need to support move-only types. See robin_hash constructor for details. - */ - bucket_entry(bucket_entry&& other) noexcept(std::is_nothrow_move_constructible::value): - bucket_hash(std::move(other)), - m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), - m_last_bucket(other.m_last_bucket) - { - if(!other.empty()) { - ::new (static_cast(std::addressof(m_value))) value_type(std::move(other.value())); - m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; - } - } - - bucket_entry& operator=(const bucket_entry& other) - noexcept(std::is_nothrow_copy_constructible::value) - { - if(this != &other) { - clear(); - - bucket_hash::operator=(other); - if(!other.empty()) { - ::new (static_cast(std::addressof(m_value))) value_type(other.value()); - } - - m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; - m_last_bucket = other.m_last_bucket; - } - - return *this; - } - - bucket_entry& operator=(bucket_entry&& ) = delete; - - ~bucket_entry() noexcept { - clear(); - } - - void clear() noexcept { - if(!empty()) { - destroy_value(); - m_dist_from_ideal_bucket = EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET; - } - } - - bool empty() const noexcept { - return m_dist_from_ideal_bucket == EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET; - } - - value_type& value() noexcept { - tsl_assert(!empty()); - return *reinterpret_cast(std::addressof(m_value)); - } - - const value_type& value() const noexcept { - tsl_assert(!empty()); - return *reinterpret_cast(std::addressof(m_value)); - } - - distance_type dist_from_ideal_bucket() const noexcept { - return m_dist_from_ideal_bucket; - } - - bool last_bucket() const noexcept { - return m_last_bucket; - } - - void set_as_last_bucket() noexcept { - m_last_bucket = true; - } - - template - void set_value_of_empty_bucket(distance_type dist_from_ideal_bucket, - truncated_hash_type hash, Args&&... value_type_args) - { - tsl_assert(dist_from_ideal_bucket >= 0); - tsl_assert(empty()); - - ::new (static_cast(std::addressof(m_value))) value_type(std::forward(value_type_args)...); - this->set_hash(hash); - m_dist_from_ideal_bucket = dist_from_ideal_bucket; - - tsl_assert(!empty()); - } - - void swap_with_value_in_bucket(distance_type& dist_from_ideal_bucket, - truncated_hash_type& hash, value_type& value) - { - tsl_assert(!empty()); - - using std::swap; - swap(value, this->value()); - swap(dist_from_ideal_bucket, m_dist_from_ideal_bucket); - - // Avoid warning of unused variable if StoreHash is false - (void) hash; - if(StoreHash) { - const truncated_hash_type tmp_hash = this->truncated_hash(); - this->set_hash(hash); - hash = tmp_hash; - } - } - - static truncated_hash_type truncate_hash(std::size_t hash) noexcept { - return truncated_hash_type(hash); - } - -private: - void destroy_value() noexcept { - tsl_assert(!empty()); - value().~value_type(); - } - -private: - using storage = typename std::aligned_storage::type; - - static const distance_type EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET = -1; - - distance_type m_dist_from_ideal_bucket; - bool m_last_bucket; - storage m_value; -}; - - - -/** - * Internal common class used by `robin_map` and `robin_set`. - * - * ValueType is what will be stored by `robin_hash` (usually `std::pair` for map and `Key` for set). - * - * `KeySelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a - * reference to the key. - * - * `ValueSelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a - * reference to the value. `ValueSelect` should be void if there is no value (in a set for example). - * - * The strong exception guarantee only holds if the expression - * `std::is_nothrow_swappable::value && std::is_nothrow_move_constructible::value` is true. - * - * Behaviour is undefined if the destructor of `ValueType` throws. - */ -template -class robin_hash: private Hash, private KeyEqual, private GrowthPolicy { -private: - template - using has_mapped_type = typename std::integral_constant::value>; - - static_assert(noexcept(std::declval().bucket_for_hash(std::size_t(0))), "GrowthPolicy::bucket_for_hash must be noexcept."); - static_assert(noexcept(std::declval().clear()), "GrowthPolicy::clear must be noexcept."); - -public: - template - class robin_iterator; - - using key_type = typename KeySelect::key_type; - using value_type = ValueType; - using size_type = std::size_t; - using difference_type = std::ptrdiff_t; - using hasher = Hash; - using key_equal = KeyEqual; - using allocator_type = Allocator; - using reference = value_type&; - using const_reference = const value_type&; - using pointer = value_type*; - using const_pointer = const value_type*; - using iterator = robin_iterator; - using const_iterator = robin_iterator; - - -private: - /** - * Either store the hash because we are asked by the `StoreHash` template parameter - * or store the hash because it doesn't cost us anything in size and can be used to speed up rehash. - */ - static constexpr bool STORE_HASH = StoreHash || - ( - (sizeof(tsl::detail_robin_hash::bucket_entry) == - sizeof(tsl::detail_robin_hash::bucket_entry)) - && - (sizeof(std::size_t) == sizeof(truncated_hash_type) || - is_power_of_two_policy::value) - && - // Don't store the hash for primitive types with default hash. - (!std::is_arithmetic::value || - !std::is_same>::value) - ); - - /** - * Only use the stored hash on lookup if we are explictly asked. We are not sure how slow - * the KeyEqual operation is. An extra comparison may slow things down with a fast KeyEqual. - */ - static constexpr bool USE_STORED_HASH_ON_LOOKUP = StoreHash; - - /** - * We can only use the hash on rehash if the size of the hash type is the same as the stored one or - * if we use a power of two modulo. In the case of the power of two modulo, we just mask - * the least significant bytes, we just have to check that the truncated_hash_type didn't truncated - * more bytes. - */ - static bool USE_STORED_HASH_ON_REHASH(size_type bucket_count) { - (void) bucket_count; - if(STORE_HASH && sizeof(std::size_t) == sizeof(truncated_hash_type)) { - return true; - } - else if(STORE_HASH && is_power_of_two_policy::value) { - tsl_assert(bucket_count > 0); - return (bucket_count - 1) <= (std::numeric_limits::max)(); - } - else { - return false; - } - } - - using bucket_entry = tsl::detail_robin_hash::bucket_entry; - using distance_type = typename bucket_entry::distance_type; - - using buckets_allocator = typename std::allocator_traits::template rebind_alloc; - using buckets_container_type = std::vector; - - -public: - /** - * The 'operator*()' and 'operator->()' methods return a const reference and const pointer respectively to the - * stored value type. - * - * In case of a map, to get a mutable reference to the value associated to a key (the '.second' in the - * stored pair), you have to call 'value()'. - * - * The main reason for this is that if we returned a `std::pair&` instead - * of a `const std::pair&`, the user may modify the key which will put the map in a undefined state. - */ - template - class robin_iterator { - friend class robin_hash; - - private: - using iterator_bucket = typename std::conditional::type; - - - robin_iterator(iterator_bucket it) noexcept: m_iterator(it) { - } - - public: - using iterator_category = std::forward_iterator_tag; - using value_type = const typename robin_hash::value_type; - using difference_type = std::ptrdiff_t; - using reference = value_type&; - using pointer = value_type*; - - - robin_iterator() noexcept { - } - - robin_iterator(const robin_iterator& other) noexcept: m_iterator(other.m_iterator) { - } - - const typename robin_hash::key_type& key() const { - return KeySelect()(m_iterator->value()); - } - - template::value && IsConst>::type* = nullptr> - const typename U::value_type& value() const { - return U()(m_iterator->value()); - } - - template::value && !IsConst>::type* = nullptr> - typename U::value_type& value() { - return U()(m_iterator->value()); - } - - reference operator*() const { - return m_iterator->value(); - } - - pointer operator->() const { - return std::addressof(m_iterator->value()); - } - - robin_iterator& operator++() { - while(true) { - if(m_iterator->last_bucket()) { - ++m_iterator; - return *this; - } - - ++m_iterator; - if(!m_iterator->empty()) { - return *this; - } - } - } - - robin_iterator operator++(int) { - robin_iterator tmp(*this); - ++*this; - - return tmp; - } - - friend bool operator==(const robin_iterator& lhs, const robin_iterator& rhs) { - return lhs.m_iterator == rhs.m_iterator; - } - - friend bool operator!=(const robin_iterator& lhs, const robin_iterator& rhs) { - return !(lhs == rhs); - } - - private: - iterator_bucket m_iterator; - }; - - -public: - robin_hash(size_type bucket_count, - const Hash& hash, - const KeyEqual& equal, - const Allocator& alloc, - float max_load_factor): Hash(hash), - KeyEqual(equal), - GrowthPolicy(bucket_count), - m_buckets(alloc), - m_first_or_empty_bucket(static_empty_bucket_ptr()), - m_bucket_count(bucket_count), - m_nb_elements(0), - m_grow_on_next_insert(false) - { - if(bucket_count > max_bucket_count()) { - TSL_THROW_OR_TERMINATE(std::length_error, "The map exceeds its maxmimum size."); - } - - if(m_bucket_count > 0) { - /* - * We can't use the `vector(size_type count, const Allocator& alloc)` constructor - * as it's only available in C++14 and we need to support C++11. We thus must resize after using - * the `vector(const Allocator& alloc)` constructor. - * - * We can't use `vector(size_type count, const T& value, const Allocator& alloc)` as it requires the - * value T to be copyable. - */ - m_buckets.resize(m_bucket_count); - m_first_or_empty_bucket = m_buckets.data(); - - tsl_assert(!m_buckets.empty()); - m_buckets.back().set_as_last_bucket(); - } - - - this->max_load_factor(max_load_factor); - } - - robin_hash(const robin_hash& other): Hash(other), - KeyEqual(other), - GrowthPolicy(other), - m_buckets(other.m_buckets), - m_first_or_empty_bucket(m_buckets.empty()?static_empty_bucket_ptr():m_buckets.data()), - m_bucket_count(other.m_bucket_count), - m_nb_elements(other.m_nb_elements), - m_load_threshold(other.m_load_threshold), - m_max_load_factor(other.m_max_load_factor), - m_grow_on_next_insert(other.m_grow_on_next_insert) - { - } - - robin_hash(robin_hash&& other) noexcept(std::is_nothrow_move_constructible::value && - std::is_nothrow_move_constructible::value && - std::is_nothrow_move_constructible::value && - std::is_nothrow_move_constructible::value) - : Hash(std::move(static_cast(other))), - KeyEqual(std::move(static_cast(other))), - GrowthPolicy(std::move(static_cast(other))), - m_buckets(std::move(other.m_buckets)), - m_first_or_empty_bucket(m_buckets.empty()?static_empty_bucket_ptr():m_buckets.data()), - m_bucket_count(other.m_bucket_count), - m_nb_elements(other.m_nb_elements), - m_load_threshold(other.m_load_threshold), - m_max_load_factor(other.m_max_load_factor), - m_grow_on_next_insert(other.m_grow_on_next_insert) - { - other.GrowthPolicy::clear(); - other.m_buckets.clear(); - other.m_first_or_empty_bucket = static_empty_bucket_ptr(); - other.m_bucket_count = 0; - other.m_nb_elements = 0; - other.m_load_threshold = 0; - other.m_grow_on_next_insert = false; - } - - robin_hash& operator=(const robin_hash& other) { - if(&other != this) { - Hash::operator=(other); - KeyEqual::operator=(other); - GrowthPolicy::operator=(other); - - m_buckets = other.m_buckets; - m_first_or_empty_bucket = m_buckets.empty()?static_empty_bucket_ptr(): - m_buckets.data(); - m_bucket_count = other.m_bucket_count; - m_nb_elements = other.m_nb_elements; - m_load_threshold = other.m_load_threshold; - m_max_load_factor = other.m_max_load_factor; - m_grow_on_next_insert = other.m_grow_on_next_insert; - } - - return *this; - } - - robin_hash& operator=(robin_hash&& other) { - other.swap(*this); - other.clear(); - - return *this; - } - - allocator_type get_allocator() const { - return m_buckets.get_allocator(); - } - - - /* - * Iterators - */ - iterator begin() noexcept { - auto begin = m_buckets.begin(); - while(begin != m_buckets.end() && begin->empty()) { - ++begin; - } - - return iterator(begin); - } - - const_iterator begin() const noexcept { - return cbegin(); - } - - const_iterator cbegin() const noexcept { - auto begin = m_buckets.cbegin(); - while(begin != m_buckets.cend() && begin->empty()) { - ++begin; - } - - return const_iterator(begin); - } - - iterator end() noexcept { - return iterator(m_buckets.end()); - } - - const_iterator end() const noexcept { - return cend(); - } - - const_iterator cend() const noexcept { - return const_iterator(m_buckets.cend()); - } - - - /* - * Capacity - */ - bool empty() const noexcept { - return m_nb_elements == 0; - } - - size_type size() const noexcept { - return m_nb_elements; - } - - size_type max_size() const noexcept { - return m_buckets.max_size(); - } - - /* - * Modifiers - */ - void clear() noexcept { - for(auto& bucket: m_buckets) { - bucket.clear(); - } - - m_nb_elements = 0; - m_grow_on_next_insert = false; - } - - - - template - std::pair insert(P&& value) { - return insert_impl(KeySelect()(value), std::forward

(value)); - } - - template - iterator insert(const_iterator hint, P&& value) { - if(hint != cend() && compare_keys(KeySelect()(*hint), KeySelect()(value))) { - return mutable_iterator(hint); - } - - return insert(std::forward

(value)).first; - } - - template - void insert(InputIt first, InputIt last) { - if(std::is_base_of::iterator_category>::value) - { - const auto nb_elements_insert = std::distance(first, last); - const size_type nb_free_buckets = m_load_threshold - size(); - tsl_assert(m_load_threshold >= size()); - - if(nb_elements_insert > 0 && nb_free_buckets < size_type(nb_elements_insert)) { - reserve(size() + size_type(nb_elements_insert)); - } - } - - for(; first != last; ++first) { - insert(*first); - } - } - - - - template - std::pair insert_or_assign(K&& key, M&& obj) { - auto it = try_emplace(std::forward(key), std::forward(obj)); - if(!it.second) { - it.first.value() = std::forward(obj); - } - - return it; - } - - template - iterator insert_or_assign(const_iterator hint, K&& key, M&& obj) { - if(hint != cend() && compare_keys(KeySelect()(*hint), key)) { - auto it = mutable_iterator(hint); - it.value() = std::forward(obj); - - return it; - } - - return insert_or_assign(std::forward(key), std::forward(obj)).first; - } - - - template - std::pair emplace(Args&&... args) { - return insert(value_type(std::forward(args)...)); - } - - template - iterator emplace_hint(const_iterator hint, Args&&... args) { - return insert(hint, value_type(std::forward(args)...)); - } - - - - template - std::pair try_emplace(K&& key, Args&&... args) { - return insert_impl(key, std::piecewise_construct, - std::forward_as_tuple(std::forward(key)), - std::forward_as_tuple(std::forward(args)...)); - } - - template - iterator try_emplace(const_iterator hint, K&& key, Args&&... args) { - if(hint != cend() && compare_keys(KeySelect()(*hint), key)) { - return mutable_iterator(hint); - } - - return try_emplace(std::forward(key), std::forward(args)...).first; - } - - /** - * Here to avoid `template size_type erase(const K& key)` being used when - * we use an `iterator` instead of a `const_iterator`. - */ - iterator erase(iterator pos) { - erase_from_bucket(pos); - - /** - * Erase bucket used a backward shift after clearing the bucket. - * Check if there is a new value in the bucket, if not get the next non-empty. - */ - if(pos.m_iterator->empty()) { - ++pos; - } - - return pos; - } - - iterator erase(const_iterator pos) { - return erase(mutable_iterator(pos)); - } - - iterator erase(const_iterator first, const_iterator last) { - if(first == last) { - return mutable_iterator(first); - } - - auto first_mutable = mutable_iterator(first); - auto last_mutable = mutable_iterator(last); - for(auto it = first_mutable.m_iterator; it != last_mutable.m_iterator; ++it) { - if(!it->empty()) { - it->clear(); - m_nb_elements--; - } - } - - if(last_mutable == end()) { - return end(); - } - - - /* - * Backward shift on the values which come after the deleted values. - * We try to move the values closer to their ideal bucket. - */ - std::size_t icloser_bucket = std::size_t(std::distance(m_buckets.begin(), first_mutable.m_iterator)); - std::size_t ito_move_closer_value = std::size_t(std::distance(m_buckets.begin(), last_mutable.m_iterator)); - tsl_assert(ito_move_closer_value > icloser_bucket); - - const std::size_t ireturn_bucket = ito_move_closer_value - - (std::min)(ito_move_closer_value - icloser_bucket, - std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket())); - - while(ito_move_closer_value < m_buckets.size() && m_buckets[ito_move_closer_value].dist_from_ideal_bucket() > 0) { - icloser_bucket = ito_move_closer_value - - (std::min)(ito_move_closer_value - icloser_bucket, - std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket())); - - - tsl_assert(m_buckets[icloser_bucket].empty()); - const distance_type new_distance = distance_type(m_buckets[ito_move_closer_value].dist_from_ideal_bucket() - - (ito_move_closer_value - icloser_bucket)); - m_buckets[icloser_bucket].set_value_of_empty_bucket(new_distance, - m_buckets[ito_move_closer_value].truncated_hash(), - std::move(m_buckets[ito_move_closer_value].value())); - m_buckets[ito_move_closer_value].clear(); - - - ++icloser_bucket; - ++ito_move_closer_value; - } - - - return iterator(m_buckets.begin() + ireturn_bucket); - } - - - template - size_type erase(const K& key) { - return erase(key, hash_key(key)); - } - - template - size_type erase(const K& key, std::size_t hash) { - auto it = find(key, hash); - if(it != end()) { - erase_from_bucket(it); - - return 1; - } - else { - return 0; - } - } - - - - - - void swap(robin_hash& other) { - using std::swap; - - swap(static_cast(*this), static_cast(other)); - swap(static_cast(*this), static_cast(other)); - swap(static_cast(*this), static_cast(other)); - swap(m_buckets, other.m_buckets); - swap(m_first_or_empty_bucket, other.m_first_or_empty_bucket); - swap(m_bucket_count, other.m_bucket_count); - swap(m_nb_elements, other.m_nb_elements); - swap(m_load_threshold, other.m_load_threshold); - swap(m_max_load_factor, other.m_max_load_factor); - swap(m_grow_on_next_insert, other.m_grow_on_next_insert); - } - - - /* - * Lookup - */ - template::value>::type* = nullptr> - typename U::value_type& at(const K& key) { - return at(key, hash_key(key)); - } - - template::value>::type* = nullptr> - typename U::value_type& at(const K& key, std::size_t hash) { - return const_cast(static_cast(this)->at(key, hash)); - } - - - template::value>::type* = nullptr> - const typename U::value_type& at(const K& key) const { - return at(key, hash_key(key)); - } - - template::value>::type* = nullptr> - const typename U::value_type& at(const K& key, std::size_t hash) const { - auto it = find(key, hash); - if(it != cend()) { - return it.value(); - } - else { - TSL_THROW_OR_TERMINATE(std::out_of_range, "Couldn't find key."); - } - } - - template::value>::type* = nullptr> - typename U::value_type& operator[](K&& key) { - return try_emplace(std::forward(key)).first.value(); - } - - - template - size_type count(const K& key) const { - return count(key, hash_key(key)); - } - - template - size_type count(const K& key, std::size_t hash) const { - if(find(key, hash) != cend()) { - return 1; - } - else { - return 0; - } - } - - - template - iterator find(const K& key) { - return find_impl(key, hash_key(key)); - } - - template - iterator find(const K& key, std::size_t hash) { - return find_impl(key, hash); - } - - - template - const_iterator find(const K& key) const { - return find_impl(key, hash_key(key)); - } - - template - const_iterator find(const K& key, std::size_t hash) const { - return find_impl(key, hash); - } - - - template - std::pair equal_range(const K& key) { - return equal_range(key, hash_key(key)); - } - - template - std::pair equal_range(const K& key, std::size_t hash) { - iterator it = find(key, hash); - return std::make_pair(it, (it == end())?it:std::next(it)); - } - - - template - std::pair equal_range(const K& key) const { - return equal_range(key, hash_key(key)); - } - - template - std::pair equal_range(const K& key, std::size_t hash) const { - const_iterator it = find(key, hash); - return std::make_pair(it, (it == cend())?it:std::next(it)); - } - - /* - * Bucket interface - */ - size_type bucket_count() const { - return m_bucket_count; - } - - size_type max_bucket_count() const { - return (std::min)(GrowthPolicy::max_bucket_count(), m_buckets.max_size()); - } - - /* - * Hash policy - */ - float load_factor() const { - if(bucket_count() == 0) { - return 0; - } - - return float(m_nb_elements)/float(bucket_count()); - } - - float max_load_factor() const { - return m_max_load_factor; - } - - void max_load_factor(float ml) { - m_max_load_factor = (std::max)(0.1f, (std::min)(ml, 0.95f)); - m_load_threshold = size_type(float(bucket_count())*m_max_load_factor); - } - - void rehash(size_type count) { - count = (std::max)(count, size_type(std::ceil(float(size())/max_load_factor()))); - rehash_impl(count); - } - - void reserve(size_type count) { - rehash(size_type(std::ceil(float(count)/max_load_factor()))); - } - - /* - * Observers - */ - hasher hash_function() const { - return static_cast(*this); - } - - key_equal key_eq() const { - return static_cast(*this); - } - - - /* - * Other - */ - iterator mutable_iterator(const_iterator pos) { - return iterator(m_buckets.begin() + std::distance(m_buckets.cbegin(), pos.m_iterator)); - } - -private: - template - std::size_t hash_key(const K& key) const { - return Hash::operator()(key); - } - - template - bool compare_keys(const K1& key1, const K2& key2) const { - return KeyEqual::operator()(key1, key2); - } - - std::size_t bucket_for_hash(std::size_t hash) const { - const std::size_t bucket = GrowthPolicy::bucket_for_hash(hash); - tsl_assert(bucket < m_buckets.size() || (bucket == 0 && m_buckets.empty())); - - return bucket; - } - - template::value>::type* = nullptr> - std::size_t next_bucket(std::size_t index) const noexcept { - tsl_assert(index < bucket_count()); - - return (index + 1) & this->m_mask; - } - - template::value>::type* = nullptr> - std::size_t next_bucket(std::size_t index) const noexcept { - tsl_assert(index < bucket_count()); - - index++; - return (index != bucket_count())?index:0; - } - - - - template - iterator find_impl(const K& key, std::size_t hash) { - return mutable_iterator(static_cast(this)->find(key, hash)); - } - - template - const_iterator find_impl(const K& key, std::size_t hash) const { - std::size_t ibucket = bucket_for_hash(hash); - distance_type dist_from_ideal_bucket = 0; - - while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) { - if(TSL_LIKELY((!USE_STORED_HASH_ON_LOOKUP || (m_first_or_empty_bucket + ibucket)->bucket_hash_equal(hash)) && - compare_keys(KeySelect()((m_first_or_empty_bucket + ibucket)->value()), key))) - { - return const_iterator(m_buckets.begin() + ibucket); - } - - ibucket = next_bucket(ibucket); - dist_from_ideal_bucket++; - } - - return cend(); - } - - void erase_from_bucket(iterator pos) { - pos.m_iterator->clear(); - m_nb_elements--; - - /** - * Backward shift, swap the empty bucket, previous_ibucket, with the values on its right, ibucket, - * until we cross another empty bucket or if the other bucket has a distance_from_ideal_bucket == 0. - * - * We try to move the values closer to their ideal bucket. - */ - std::size_t previous_ibucket = std::size_t(std::distance(m_buckets.begin(), pos.m_iterator)); - std::size_t ibucket = next_bucket(previous_ibucket); - - while(m_buckets[ibucket].dist_from_ideal_bucket() > 0) { - tsl_assert(m_buckets[previous_ibucket].empty()); - - const distance_type new_distance = distance_type(m_buckets[ibucket].dist_from_ideal_bucket() - 1); - m_buckets[previous_ibucket].set_value_of_empty_bucket(new_distance, m_buckets[ibucket].truncated_hash(), - std::move(m_buckets[ibucket].value())); - m_buckets[ibucket].clear(); - - previous_ibucket = ibucket; - ibucket = next_bucket(ibucket); - } - } - - template - std::pair insert_impl(const K& key, Args&&... value_type_args) { - const std::size_t hash = hash_key(key); - - std::size_t ibucket = bucket_for_hash(hash); - distance_type dist_from_ideal_bucket = 0; - - while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) { - if((!USE_STORED_HASH_ON_LOOKUP || (m_first_or_empty_bucket + ibucket)->bucket_hash_equal(hash)) && - compare_keys(KeySelect()((m_first_or_empty_bucket + ibucket)->value()), key)) - { - return std::make_pair(iterator(m_buckets.begin() + ibucket), false); - } - - ibucket = next_bucket(ibucket); - dist_from_ideal_bucket++; - } - - if(grow_on_high_load()) { - ibucket = bucket_for_hash(hash); - dist_from_ideal_bucket = 0; - - while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) { - ibucket = next_bucket(ibucket); - dist_from_ideal_bucket++; - } - } - - - if((m_first_or_empty_bucket + ibucket)->empty()) { - (m_first_or_empty_bucket + ibucket)->set_value_of_empty_bucket(dist_from_ideal_bucket, bucket_entry::truncate_hash(hash), - std::forward(value_type_args)...); - } - else { - insert_value(ibucket, dist_from_ideal_bucket, bucket_entry::truncate_hash(hash), - std::forward(value_type_args)...); - } - - - m_nb_elements++; - /* - * The value will be inserted in ibucket in any case, either because it was - * empty or by stealing the bucket (robin hood). - */ - return std::make_pair(iterator(m_buckets.begin() + ibucket), true); - } - - - template - void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket, - truncated_hash_type hash, Args&&... value_type_args) - { - insert_value(ibucket, dist_from_ideal_bucket, hash, value_type(std::forward(value_type_args)...)); - } - - void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket, - truncated_hash_type hash, value_type&& value) - { - m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); - ibucket = next_bucket(ibucket); - dist_from_ideal_bucket++; - - while(!m_buckets[ibucket].empty()) { - if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) { - if(dist_from_ideal_bucket >= REHASH_ON_HIGH_NB_PROBES__NPROBES && - load_factor() >= REHASH_ON_HIGH_NB_PROBES__MIN_LOAD_FACTOR) - { - /** - * The number of probes is really high, rehash the map on the next insert. - * Difficult to do now as rehash may throw an exception. - */ - m_grow_on_next_insert = true; - } - - m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); - } - - ibucket = next_bucket(ibucket); - dist_from_ideal_bucket++; - } - - m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value)); - } - - - void rehash_impl(size_type count) { - robin_hash new_table(count, static_cast(*this), static_cast(*this), - get_allocator(), m_max_load_factor); - - const bool use_stored_hash = USE_STORED_HASH_ON_REHASH(new_table.bucket_count()); - for(auto& bucket: m_buckets) { - if(bucket.empty()) { - continue; - } - - const std::size_t hash = use_stored_hash?bucket.truncated_hash(): - new_table.hash_key(KeySelect()(bucket.value())); - - new_table.insert_value_on_rehash(new_table.bucket_for_hash(hash), 0, - bucket_entry::truncate_hash(hash), std::move(bucket.value())); - } - - new_table.m_nb_elements = m_nb_elements; - new_table.swap(*this); - } - - void insert_value_on_rehash(std::size_t ibucket, distance_type dist_from_ideal_bucket, - truncated_hash_type hash, value_type&& value) - { - while(true) { - if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) { - if(m_buckets[ibucket].empty()) { - m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value)); - return; - } - else { - m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); - } - } - - dist_from_ideal_bucket++; - ibucket = next_bucket(ibucket); - } - } - - - - /** - * Return true if the map has been rehashed. - */ - bool grow_on_high_load() { - if(m_grow_on_next_insert || size() >= m_load_threshold) { - rehash_impl(GrowthPolicy::next_bucket_count()); - m_grow_on_next_insert = false; - - return true; - } - - return false; - } - - -public: - static const size_type DEFAULT_INIT_BUCKETS_SIZE = 16; - static constexpr float DEFAULT_MAX_LOAD_FACTOR = 0.5f; - -private: - static const distance_type REHASH_ON_HIGH_NB_PROBES__NPROBES = 128; - static constexpr float REHASH_ON_HIGH_NB_PROBES__MIN_LOAD_FACTOR = 0.15f; - - - /** - * Return an always valid pointer to an static empty bucket_entry with last_bucket() == true. - */ - bucket_entry* static_empty_bucket_ptr() { - static bucket_entry empty_bucket(true); - return &empty_bucket; - } - -private: - buckets_container_type m_buckets; - - /** - * Points to m_buckets.data() if !m_buckets.empty() otherwise points to static_empty_bucket_ptr. - * This variable is useful to avoid the cost of checking if m_buckets is empty when trying - * to find an element. - */ - bucket_entry* m_first_or_empty_bucket; - - /** - * Used a lot in find, avoid the call to m_buckets.size() which is a bit slower. - */ - size_type m_bucket_count; - - size_type m_nb_elements; - - size_type m_load_threshold; - float m_max_load_factor; - - bool m_grow_on_next_insert; -}; - -} - -} - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_map.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_map.h deleted file mode 100644 index 5958e70..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_map.h +++ /dev/null @@ -1,668 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Tessil - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_ROBIN_MAP_H -#define TSL_ROBIN_MAP_H - - -#include -#include -#include -#include -#include -#include -#include "robin_hash.h" - - -namespace tsl { - - -/** - * Implementation of a hash map using open-adressing and the robin hood hashing algorithm with backward shift deletion. - * - * For operations modifying the hash map (insert, erase, rehash, ...), the strong exception guarantee - * is only guaranteed when the expression `std::is_nothrow_swappable>::value && - * std::is_nothrow_move_constructible>::value` is true, otherwise if an exception - * is thrown during the swap or the move, the hash map may end up in a undefined state. Per the standard - * a `Key` or `T` with a noexcept copy constructor and no move constructor also satisfies the - * `std::is_nothrow_move_constructible>::value` criterion (and will thus guarantee the - * strong exception for the map). - * - * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve - * the performance during lookups if the `KeyEqual` function takes time (if it engenders a cache-miss for example) - * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used - * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash. - * When it is detected that storing the hash will not incur any memory penality due to alignement (i.e. - * `sizeof(tsl::detail_robin_hash::bucket_entry) == - * sizeof(tsl::detail_robin_hash::bucket_entry)`) and `tsl::rh::power_of_two_growth_policy` is - * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will - * not be used on lookups unless `StoreHash` is true). - * - * `GrowthPolicy` defines how the map grows and consequently how a hash value is mapped to a bucket. - * By default the map uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets - * to a power of two and uses a mask to map the hash to a bucket instead of the slow modulo. - * Other growth policies are available and you may define your own growth policy, - * check `tsl::rh::power_of_two_growth_policy` for the interface. - * - * If the destructor of `Key` or `T` throws an exception, the behaviour of the class is undefined. - * - * Iterators invalidation: - * - clear, operator=, reserve, rehash: always invalidate the iterators. - * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators. - * - erase: always invalidate the iterators. - */ -template, - class KeyEqual = std::equal_to, - class Allocator = std::allocator>, - bool StoreHash = false, - class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>> -class robin_map { -private: - template - using has_is_transparent = tsl::detail_robin_hash::has_is_transparent; - - class KeySelect { - public: - using key_type = Key; - - const key_type& operator()(const std::pair& key_value) const noexcept { - return key_value.first; - } - - key_type& operator()(std::pair& key_value) noexcept { - return key_value.first; - } - }; - - class ValueSelect { - public: - using value_type = T; - - const value_type& operator()(const std::pair& key_value) const noexcept { - return key_value.second; - } - - value_type& operator()(std::pair& key_value) noexcept { - return key_value.second; - } - }; - - using ht = detail_robin_hash::robin_hash, KeySelect, ValueSelect, - Hash, KeyEqual, Allocator, StoreHash, GrowthPolicy>; - -public: - using key_type = typename ht::key_type; - using mapped_type = T; - using value_type = typename ht::value_type; - using size_type = typename ht::size_type; - using difference_type = typename ht::difference_type; - using hasher = typename ht::hasher; - using key_equal = typename ht::key_equal; - using allocator_type = typename ht::allocator_type; - using reference = typename ht::reference; - using const_reference = typename ht::const_reference; - using pointer = typename ht::pointer; - using const_pointer = typename ht::const_pointer; - using iterator = typename ht::iterator; - using const_iterator = typename ht::const_iterator; - - -public: - /* - * Constructors - */ - robin_map(): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE) { - } - - explicit robin_map(size_type bucket_count, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): - m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR) - { - } - - robin_map(size_type bucket_count, - const Allocator& alloc): robin_map(bucket_count, Hash(), KeyEqual(), alloc) - { - } - - robin_map(size_type bucket_count, - const Hash& hash, - const Allocator& alloc): robin_map(bucket_count, hash, KeyEqual(), alloc) - { - } - - explicit robin_map(const Allocator& alloc): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) { - } - - template - robin_map(InputIt first, InputIt last, - size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): robin_map(bucket_count, hash, equal, alloc) - { - insert(first, last); - } - - template - robin_map(InputIt first, InputIt last, - size_type bucket_count, - const Allocator& alloc): robin_map(first, last, bucket_count, Hash(), KeyEqual(), alloc) - { - } - - template - robin_map(InputIt first, InputIt last, - size_type bucket_count, - const Hash& hash, - const Allocator& alloc): robin_map(first, last, bucket_count, hash, KeyEqual(), alloc) - { - } - - robin_map(std::initializer_list init, - size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): - robin_map(init.begin(), init.end(), bucket_count, hash, equal, alloc) - { - } - - robin_map(std::initializer_list init, - size_type bucket_count, - const Allocator& alloc): - robin_map(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc) - { - } - - robin_map(std::initializer_list init, - size_type bucket_count, - const Hash& hash, - const Allocator& alloc): - robin_map(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc) - { - } - - robin_map& operator=(std::initializer_list ilist) { - m_ht.clear(); - - m_ht.reserve(ilist.size()); - m_ht.insert(ilist.begin(), ilist.end()); - - return *this; - } - - allocator_type get_allocator() const { return m_ht.get_allocator(); } - - - /* - * Iterators - */ - iterator begin() noexcept { return m_ht.begin(); } - const_iterator begin() const noexcept { return m_ht.begin(); } - const_iterator cbegin() const noexcept { return m_ht.cbegin(); } - - iterator end() noexcept { return m_ht.end(); } - const_iterator end() const noexcept { return m_ht.end(); } - const_iterator cend() const noexcept { return m_ht.cend(); } - - - /* - * Capacity - */ - bool empty() const noexcept { return m_ht.empty(); } - size_type size() const noexcept { return m_ht.size(); } - size_type max_size() const noexcept { return m_ht.max_size(); } - - /* - * Modifiers - */ - void clear() noexcept { m_ht.clear(); } - - - - std::pair insert(const value_type& value) { - return m_ht.insert(value); - } - - template::value>::type* = nullptr> - std::pair insert(P&& value) { - return m_ht.emplace(std::forward

(value)); - } - - std::pair insert(value_type&& value) { - return m_ht.insert(std::move(value)); - } - - - iterator insert(const_iterator hint, const value_type& value) { - return m_ht.insert(hint, value); - } - - template::value>::type* = nullptr> - iterator insert(const_iterator hint, P&& value) { - return m_ht.emplace_hint(hint, std::forward

(value)); - } - - iterator insert(const_iterator hint, value_type&& value) { - return m_ht.insert(hint, std::move(value)); - } - - - template - void insert(InputIt first, InputIt last) { - m_ht.insert(first, last); - } - - void insert(std::initializer_list ilist) { - m_ht.insert(ilist.begin(), ilist.end()); - } - - - - - template - std::pair insert_or_assign(const key_type& k, M&& obj) { - return m_ht.insert_or_assign(k, std::forward(obj)); - } - - template - std::pair insert_or_assign(key_type&& k, M&& obj) { - return m_ht.insert_or_assign(std::move(k), std::forward(obj)); - } - - template - iterator insert_or_assign(const_iterator hint, const key_type& k, M&& obj) { - return m_ht.insert_or_assign(hint, k, std::forward(obj)); - } - - template - iterator insert_or_assign(const_iterator hint, key_type&& k, M&& obj) { - return m_ht.insert_or_assign(hint, std::move(k), std::forward(obj)); - } - - - - /** - * Due to the way elements are stored, emplace will need to move or copy the key-value once. - * The method is equivalent to insert(value_type(std::forward(args)...)); - * - * Mainly here for compatibility with the std::unordered_map interface. - */ - template - std::pair emplace(Args&&... args) { - return m_ht.emplace(std::forward(args)...); - } - - - - /** - * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once. - * The method is equivalent to insert(hint, value_type(std::forward(args)...)); - * - * Mainly here for compatibility with the std::unordered_map interface. - */ - template - iterator emplace_hint(const_iterator hint, Args&&... args) { - return m_ht.emplace_hint(hint, std::forward(args)...); - } - - - - - template - std::pair try_emplace(const key_type& k, Args&&... args) { - return m_ht.try_emplace(k, std::forward(args)...); - } - - template - std::pair try_emplace(key_type&& k, Args&&... args) { - return m_ht.try_emplace(std::move(k), std::forward(args)...); - } - - template - iterator try_emplace(const_iterator hint, const key_type& k, Args&&... args) { - return m_ht.try_emplace(hint, k, std::forward(args)...); - } - - template - iterator try_emplace(const_iterator hint, key_type&& k, Args&&... args) { - return m_ht.try_emplace(hint, std::move(k), std::forward(args)...); - } - - - - - iterator erase(iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); } - size_type erase(const key_type& key) { return m_ht.erase(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash. - */ - size_type erase(const key_type& key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - size_type erase(const K& key) { return m_ht.erase(key); } - - /** - * @copydoc erase(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash. - */ - template::value>::type* = nullptr> - size_type erase(const K& key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - - - void swap(robin_map& other) { other.m_ht.swap(m_ht); } - - - - /* - * Lookup - */ - T& at(const Key& key) { return m_ht.at(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - T& at(const Key& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); } - - - const T& at(const Key& key) const { return m_ht.at(key); } - - /** - * @copydoc at(const Key& key, std::size_t precalculated_hash) - */ - const T& at(const Key& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); } - - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - T& at(const K& key) { return m_ht.at(key); } - - /** - * @copydoc at(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - T& at(const K& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); } - - - /** - * @copydoc at(const K& key) - */ - template::value>::type* = nullptr> - const T& at(const K& key) const { return m_ht.at(key); } - - /** - * @copydoc at(const K& key, std::size_t precalculated_hash) - */ - template::value>::type* = nullptr> - const T& at(const K& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); } - - - - - T& operator[](const Key& key) { return m_ht[key]; } - T& operator[](Key&& key) { return m_ht[std::move(key)]; } - - - - - size_type count(const Key& key) const { return m_ht.count(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - size_type count(const Key& key, std::size_t precalculated_hash) const { - return m_ht.count(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - size_type count(const K& key) const { return m_ht.count(key); } - - /** - * @copydoc count(const K& key) const - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } - - - - - iterator find(const Key& key) { return m_ht.find(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } - - const_iterator find(const Key& key) const { return m_ht.find(key); } - - /** - * @copydoc find(const Key& key, std::size_t precalculated_hash) - */ - const_iterator find(const Key& key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - iterator find(const K& key) { return m_ht.find(key); } - - /** - * @copydoc find(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } - - /** - * @copydoc find(const K& key) - */ - template::value>::type* = nullptr> - const_iterator find(const K& key) const { return m_ht.find(key); } - - /** - * @copydoc find(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - const_iterator find(const K& key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - - - - std::pair equal_range(const Key& key) { return m_ht.equal_range(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - std::pair equal_range(const Key& key, std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - std::pair equal_range(const Key& key) const { return m_ht.equal_range(key); } - - /** - * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) - */ - std::pair equal_range(const Key& key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key) { return m_ht.equal_range(key); } - - - /** - * @copydoc equal_range(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key, std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * @copydoc equal_range(const K& key) - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key) const { return m_ht.equal_range(key); } - - /** - * @copydoc equal_range(const K& key, std::size_t precalculated_hash) - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - - - - /* - * Bucket interface - */ - size_type bucket_count() const { return m_ht.bucket_count(); } - size_type max_bucket_count() const { return m_ht.max_bucket_count(); } - - - /* - * Hash policy - */ - float load_factor() const { return m_ht.load_factor(); } - float max_load_factor() const { return m_ht.max_load_factor(); } - void max_load_factor(float ml) { m_ht.max_load_factor(ml); } - - void rehash(size_type count) { m_ht.rehash(count); } - void reserve(size_type count) { m_ht.reserve(count); } - - - /* - * Observers - */ - hasher hash_function() const { return m_ht.hash_function(); } - key_equal key_eq() const { return m_ht.key_eq(); } - - /* - * Other - */ - - /** - * Convert a const_iterator to an iterator. - */ - iterator mutable_iterator(const_iterator pos) { - return m_ht.mutable_iterator(pos); - } - - friend bool operator==(const robin_map& lhs, const robin_map& rhs) { - if(lhs.size() != rhs.size()) { - return false; - } - - for(const auto& element_lhs: lhs) { - const auto it_element_rhs = rhs.find(element_lhs.first); - if(it_element_rhs == rhs.cend() || element_lhs.second != it_element_rhs->second) { - return false; - } - } - - return true; - } - - friend bool operator!=(const robin_map& lhs, const robin_map& rhs) { - return !operator==(lhs, rhs); - } - - friend void swap(robin_map& lhs, robin_map& rhs) { - lhs.swap(rhs); - } - -private: - ht m_ht; -}; - - -/** - * Same as `tsl::robin_map`. - */ -template, - class KeyEqual = std::equal_to, - class Allocator = std::allocator>, - bool StoreHash = false> -using robin_pg_map = robin_map; - -} // end namespace tsl - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_set.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_set.h deleted file mode 100644 index 4e4667e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/robin_set.h +++ /dev/null @@ -1,535 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Tessil - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_ROBIN_SET_H -#define TSL_ROBIN_SET_H - - -#include -#include -#include -#include -#include -#include -#include "robin_hash.h" - - -namespace tsl { - - -/** - * Implementation of a hash set using open-adressing and the robin hood hashing algorithm with backward shift deletion. - * - * For operations modifying the hash set (insert, erase, rehash, ...), the strong exception guarantee - * is only guaranteed when the expression `std::is_nothrow_swappable::value && - * std::is_nothrow_move_constructible::value` is true, otherwise if an exception - * is thrown during the swap or the move, the hash set may end up in a undefined state. Per the standard - * a `Key` with a noexcept copy constructor and no move constructor also satisfies the - * `std::is_nothrow_move_constructible::value` criterion (and will thus guarantee the - * strong exception for the set). - * - * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve - * the performance during lookups if the `KeyEqual` function takes time (or engenders a cache-miss for example) - * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used - * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash. - * When it is detected that storing the hash will not incur any memory penality due to alignement (i.e. - * `sizeof(tsl::detail_robin_hash::bucket_entry) == - * sizeof(tsl::detail_robin_hash::bucket_entry)`) and `tsl::rh::power_of_two_growth_policy` is - * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will - * not be used on lookups unless `StoreHash` is true). - * - * `GrowthPolicy` defines how the set grows and consequently how a hash value is mapped to a bucket. - * By default the set uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets - * to a power of two and uses a mask to set the hash to a bucket instead of the slow modulo. - * Other growth policies are available and you may define your own growth policy, - * check `tsl::rh::power_of_two_growth_policy` for the interface. - * - * If the destructor of `Key` throws an exception, the behaviour of the class is undefined. - * - * Iterators invalidation: - * - clear, operator=, reserve, rehash: always invalidate the iterators. - * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators. - * - erase: always invalidate the iterators. - */ -template, - class KeyEqual = std::equal_to, - class Allocator = std::allocator, - bool StoreHash = false, - class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>> -class robin_set { -private: - template - using has_is_transparent = tsl::detail_robin_hash::has_is_transparent; - - class KeySelect { - public: - using key_type = Key; - - const key_type& operator()(const Key& key) const noexcept { - return key; - } - - key_type& operator()(Key& key) noexcept { - return key; - } - }; - - using ht = detail_robin_hash::robin_hash; - -public: - using key_type = typename ht::key_type; - using value_type = typename ht::value_type; - using size_type = typename ht::size_type; - using difference_type = typename ht::difference_type; - using hasher = typename ht::hasher; - using key_equal = typename ht::key_equal; - using allocator_type = typename ht::allocator_type; - using reference = typename ht::reference; - using const_reference = typename ht::const_reference; - using pointer = typename ht::pointer; - using const_pointer = typename ht::const_pointer; - using iterator = typename ht::iterator; - using const_iterator = typename ht::const_iterator; - - - /* - * Constructors - */ - robin_set(): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE) { - } - - explicit robin_set(size_type bucket_count, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): - m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR) - { - } - - robin_set(size_type bucket_count, - const Allocator& alloc): robin_set(bucket_count, Hash(), KeyEqual(), alloc) - { - } - - robin_set(size_type bucket_count, - const Hash& hash, - const Allocator& alloc): robin_set(bucket_count, hash, KeyEqual(), alloc) - { - } - - explicit robin_set(const Allocator& alloc): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) { - } - - template - robin_set(InputIt first, InputIt last, - size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): robin_set(bucket_count, hash, equal, alloc) - { - insert(first, last); - } - - template - robin_set(InputIt first, InputIt last, - size_type bucket_count, - const Allocator& alloc): robin_set(first, last, bucket_count, Hash(), KeyEqual(), alloc) - { - } - - template - robin_set(InputIt first, InputIt last, - size_type bucket_count, - const Hash& hash, - const Allocator& alloc): robin_set(first, last, bucket_count, hash, KeyEqual(), alloc) - { - } - - robin_set(std::initializer_list init, - size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, - const Hash& hash = Hash(), - const KeyEqual& equal = KeyEqual(), - const Allocator& alloc = Allocator()): - robin_set(init.begin(), init.end(), bucket_count, hash, equal, alloc) - { - } - - robin_set(std::initializer_list init, - size_type bucket_count, - const Allocator& alloc): - robin_set(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc) - { - } - - robin_set(std::initializer_list init, - size_type bucket_count, - const Hash& hash, - const Allocator& alloc): - robin_set(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc) - { - } - - - robin_set& operator=(std::initializer_list ilist) { - m_ht.clear(); - - m_ht.reserve(ilist.size()); - m_ht.insert(ilist.begin(), ilist.end()); - - return *this; - } - - allocator_type get_allocator() const { return m_ht.get_allocator(); } - - - /* - * Iterators - */ - iterator begin() noexcept { return m_ht.begin(); } - const_iterator begin() const noexcept { return m_ht.begin(); } - const_iterator cbegin() const noexcept { return m_ht.cbegin(); } - - iterator end() noexcept { return m_ht.end(); } - const_iterator end() const noexcept { return m_ht.end(); } - const_iterator cend() const noexcept { return m_ht.cend(); } - - - /* - * Capacity - */ - bool empty() const noexcept { return m_ht.empty(); } - size_type size() const noexcept { return m_ht.size(); } - size_type max_size() const noexcept { return m_ht.max_size(); } - - /* - * Modifiers - */ - void clear() noexcept { m_ht.clear(); } - - - - - std::pair insert(const value_type& value) { - return m_ht.insert(value); - } - - std::pair insert(value_type&& value) { - return m_ht.insert(std::move(value)); - } - - iterator insert(const_iterator hint, const value_type& value) { - return m_ht.insert(hint, value); - } - - iterator insert(const_iterator hint, value_type&& value) { - return m_ht.insert(hint, std::move(value)); - } - - template - void insert(InputIt first, InputIt last) { - m_ht.insert(first, last); - } - - void insert(std::initializer_list ilist) { - m_ht.insert(ilist.begin(), ilist.end()); - } - - - - - /** - * Due to the way elements are stored, emplace will need to move or copy the key-value once. - * The method is equivalent to insert(value_type(std::forward(args)...)); - * - * Mainly here for compatibility with the std::unordered_map interface. - */ - template - std::pair emplace(Args&&... args) { - return m_ht.emplace(std::forward(args)...); - } - - - - /** - * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once. - * The method is equivalent to insert(hint, value_type(std::forward(args)...)); - * - * Mainly here for compatibility with the std::unordered_map interface. - */ - template - iterator emplace_hint(const_iterator hint, Args&&... args) { - return m_ht.emplace_hint(hint, std::forward(args)...); - } - - - - iterator erase(iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); } - size_type erase(const key_type& key) { return m_ht.erase(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash. - */ - size_type erase(const key_type& key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - size_type erase(const K& key) { return m_ht.erase(key); } - - /** - * @copydoc erase(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash. - */ - template::value>::type* = nullptr> - size_type erase(const K& key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - - - void swap(robin_set& other) { other.m_ht.swap(m_ht); } - - - - /* - * Lookup - */ - size_type count(const Key& key) const { return m_ht.count(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - size_type count(const Key& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - size_type count(const K& key) const { return m_ht.count(key); } - - /** - * @copydoc count(const K& key) const - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } - - - - - iterator find(const Key& key) { return m_ht.find(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } - - const_iterator find(const Key& key) const { return m_ht.find(key); } - - /** - * @copydoc find(const Key& key, std::size_t precalculated_hash) - */ - const_iterator find(const Key& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - iterator find(const K& key) { return m_ht.find(key); } - - /** - * @copydoc find(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } - - /** - * @copydoc find(const K& key) - */ - template::value>::type* = nullptr> - const_iterator find(const K& key) const { return m_ht.find(key); } - - /** - * @copydoc find(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - const_iterator find(const K& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); } - - - - - std::pair equal_range(const Key& key) { return m_ht.equal_range(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - std::pair equal_range(const Key& key, std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - std::pair equal_range(const Key& key) const { return m_ht.equal_range(key); } - - /** - * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) - */ - std::pair equal_range(const Key& key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. - * If so, K must be hashable and comparable to Key. - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key) { return m_ht.equal_range(key); } - - /** - * @copydoc equal_range(const K& key) - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same - * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash. - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key, std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * @copydoc equal_range(const K& key) - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key) const { return m_ht.equal_range(key); } - - /** - * @copydoc equal_range(const K& key, std::size_t precalculated_hash) - */ - template::value>::type* = nullptr> - std::pair equal_range(const K& key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - - - - /* - * Bucket interface - */ - size_type bucket_count() const { return m_ht.bucket_count(); } - size_type max_bucket_count() const { return m_ht.max_bucket_count(); } - - - /* - * Hash policy - */ - float load_factor() const { return m_ht.load_factor(); } - float max_load_factor() const { return m_ht.max_load_factor(); } - void max_load_factor(float ml) { m_ht.max_load_factor(ml); } - - void rehash(size_type count) { m_ht.rehash(count); } - void reserve(size_type count) { m_ht.reserve(count); } - - - /* - * Observers - */ - hasher hash_function() const { return m_ht.hash_function(); } - key_equal key_eq() const { return m_ht.key_eq(); } - - - /* - * Other - */ - - /** - * Convert a const_iterator to an iterator. - */ - iterator mutable_iterator(const_iterator pos) { - return m_ht.mutable_iterator(pos); - } - - friend bool operator==(const robin_set& lhs, const robin_set& rhs) { - if(lhs.size() != rhs.size()) { - return false; - } - - for(const auto& element_lhs: lhs) { - const auto it_element_rhs = rhs.find(element_lhs); - if(it_element_rhs == rhs.cend()) { - return false; - } - } - - return true; - } - - friend bool operator!=(const robin_set& lhs, const robin_set& rhs) { - return !operator==(lhs, rhs); - } - - friend void swap(robin_set& lhs, robin_set& rhs) { - lhs.swap(rhs); - } - -private: - ht m_ht; -}; - - -/** - * Same as `tsl::robin_set`. - */ -template, - class KeyEqual = std::equal_to, - class Allocator = std::allocator, - bool StoreHash = false> -using robin_pg_set = robin_set; - -} // end namespace tsl - -#endif - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_growth_policy.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_growth_policy.h deleted file mode 100644 index d73aaaf..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_growth_policy.h +++ /dev/null @@ -1,301 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Thibaut Goetghebuer-Planchon - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_SPARSE_GROWTH_POLICY_H -#define TSL_SPARSE_GROWTH_POLICY_H - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tsl { -namespace sh { - -/** - * Grow the hash table by a factor of GrowthFactor keeping the bucket count to a - * power of two. It allows the table to use a mask operation instead of a modulo - * operation to map a hash to a bucket. - * - * GrowthFactor must be a power of two >= 2. - */ -template -class power_of_two_growth_policy { - public: - /** - * Called on the hash table creation and on rehash. The number of buckets for - * the table is passed in parameter. This number is a minimum, the policy may - * update this value with a higher value if needed (but not lower). - * - * If 0 is given, min_bucket_count_in_out must still be 0 after the policy - * creation and bucket_for_hash must always return 0 in this case. - */ - explicit power_of_two_growth_policy(std::size_t &min_bucket_count_in_out) { - if (min_bucket_count_in_out > max_bucket_count()) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - if (min_bucket_count_in_out > 0) { - min_bucket_count_in_out = - round_up_to_power_of_two(min_bucket_count_in_out); - m_mask = min_bucket_count_in_out - 1; - } else { - m_mask = 0; - } - } - - /** - * Return the bucket [0, bucket_count()) to which the hash belongs. - * If bucket_count() is 0, it must always return 0. - */ - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return hash & m_mask; - } - - /** - * Return the number of buckets that should be used on next growth. - */ - std::size_t next_bucket_count() const { - if ((m_mask + 1) > max_bucket_count() / GrowthFactor) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - return (m_mask + 1) * GrowthFactor; - } - - /** - * Return the maximum number of buckets supported by the policy. - */ - std::size_t max_bucket_count() const { - // Largest power of two. - return (std::numeric_limits::max() / 2) + 1; - } - - /** - * Reset the growth policy as if it was created with a bucket count of 0. - * After a clear, the policy must always return 0 when bucket_for_hash is - * called. - */ - void clear() noexcept { m_mask = 0; } - - private: - static std::size_t round_up_to_power_of_two(std::size_t value) { - if (is_power_of_two(value)) { - return value; - } - - if (value == 0) { - return 1; - } - - --value; - for (std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) { - value |= value >> i; - } - - return value + 1; - } - - static constexpr bool is_power_of_two(std::size_t value) { - return value != 0 && (value & (value - 1)) == 0; - } - - protected: - static_assert(is_power_of_two(GrowthFactor) && GrowthFactor >= 2, - "GrowthFactor must be a power of two >= 2."); - - std::size_t m_mask; -}; - -/** - * Grow the hash table by GrowthFactor::num / GrowthFactor::den and use a modulo - * to map a hash to a bucket. Slower but it can be useful if you want a slower - * growth. - */ -template > -class mod_growth_policy { - public: - explicit mod_growth_policy(std::size_t &min_bucket_count_in_out) { - if (min_bucket_count_in_out > max_bucket_count()) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - if (min_bucket_count_in_out > 0) { - m_mod = min_bucket_count_in_out; - } else { - m_mod = 1; - } - } - - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return hash % m_mod; - } - - std::size_t next_bucket_count() const { - if (m_mod == max_bucket_count()) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - const double next_bucket_count = - std::ceil(double(m_mod) * REHASH_SIZE_MULTIPLICATION_FACTOR); - if (!std::isnormal(next_bucket_count)) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - if (next_bucket_count > double(max_bucket_count())) { - return max_bucket_count(); - } else { - return std::size_t(next_bucket_count); - } - } - - std::size_t max_bucket_count() const { return MAX_BUCKET_COUNT; } - - void clear() noexcept { m_mod = 1; } - - private: - static constexpr double REHASH_SIZE_MULTIPLICATION_FACTOR = - 1.0 * GrowthFactor::num / GrowthFactor::den; - static const std::size_t MAX_BUCKET_COUNT = - std::size_t(double(std::numeric_limits::max() / - REHASH_SIZE_MULTIPLICATION_FACTOR)); - - static_assert(REHASH_SIZE_MULTIPLICATION_FACTOR >= 1.1, - "Growth factor should be >= 1.1."); - - std::size_t m_mod; -}; - -/** - * Grow the hash table by using prime numbers as bucket count. Slower than - * tsl::sh::power_of_two_growth_policy in general but will probably distribute - * the values around better in the buckets with a poor hash function. - * - * To allow the compiler to optimize the modulo operation, a lookup table is - * used with constant primes numbers. - * - * With a switch the code would look like: - * \code - * switch(iprime) { // iprime is the current prime of the hash table - * case 0: hash % 5ul; - * break; - * case 1: hash % 17ul; - * break; - * case 2: hash % 29ul; - * break; - * ... - * } - * \endcode - * - * Due to the constant variable in the modulo the compiler is able to optimize - * the operation by a series of multiplications, substractions and shifts. - * - * The 'hash % 5' could become something like 'hash - (hash * 0xCCCCCCCD) >> 34) - * * 5' in a 64 bits environment. - */ -class prime_growth_policy { - public: - explicit prime_growth_policy(std::size_t &min_bucket_count_in_out) { - auto it_prime = std::lower_bound(primes().begin(), primes().end(), - min_bucket_count_in_out); - if (it_prime == primes().end()) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - m_iprime = - static_cast(std::distance(primes().begin(), it_prime)); - if (min_bucket_count_in_out > 0) { - min_bucket_count_in_out = *it_prime; - } else { - min_bucket_count_in_out = 0; - } - } - - std::size_t bucket_for_hash(std::size_t hash) const noexcept { - return mod_prime()[m_iprime](hash); - } - - std::size_t next_bucket_count() const { - if (m_iprime + 1 >= primes().size()) { - throw std::length_error("The hash table exceeds its maximum size."); - } - - return primes()[m_iprime + 1]; - } - - std::size_t max_bucket_count() const { return primes().back(); } - - void clear() noexcept { m_iprime = 0; } - - private: - static const std::array &primes() { - static const std::array PRIMES = { - {1ul, 5ul, 17ul, 29ul, 37ul, - 53ul, 67ul, 79ul, 97ul, 131ul, - 193ul, 257ul, 389ul, 521ul, 769ul, - 1031ul, 1543ul, 2053ul, 3079ul, 6151ul, - 12289ul, 24593ul, 49157ul, 98317ul, 196613ul, - 393241ul, 786433ul, 1572869ul, 3145739ul, 6291469ul, - 12582917ul, 25165843ul, 50331653ul, 100663319ul, 201326611ul, - 402653189ul, 805306457ul, 1610612741ul, 3221225473ul, 4294967291ul}}; - - static_assert( - std::numeric_limits::max() >= PRIMES.size(), - "The type of m_iprime is not big enough."); - - return PRIMES; - } - - static const std::array &mod_prime() { - // MOD_PRIME[iprime](hash) returns hash % PRIMES[iprime]. This table allows - // for faster modulo as the compiler can optimize the modulo code better - // with a constant known at the compilation. - static const std::array MOD_PRIME = { - {&mod<0>, &mod<1>, &mod<2>, &mod<3>, &mod<4>, &mod<5>, &mod<6>, - &mod<7>, &mod<8>, &mod<9>, &mod<10>, &mod<11>, &mod<12>, &mod<13>, - &mod<14>, &mod<15>, &mod<16>, &mod<17>, &mod<18>, &mod<19>, &mod<20>, - &mod<21>, &mod<22>, &mod<23>, &mod<24>, &mod<25>, &mod<26>, &mod<27>, - &mod<28>, &mod<29>, &mod<30>, &mod<31>, &mod<32>, &mod<33>, &mod<34>, - &mod<35>, &mod<36>, &mod<37>, &mod<38>, &mod<39>}}; - - return MOD_PRIME; - } - - template - static std::size_t mod(std::size_t hash) { - return hash % primes()[IPrime]; - } - - private: - unsigned int m_iprime; -}; - -} // namespace sh -} // namespace tsl - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_hash.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_hash.h deleted file mode 100644 index e2115b4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_hash.h +++ /dev/null @@ -1,2215 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Thibaut Goetghebuer-Planchon - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_SPARSE_HASH_H -#define TSL_SPARSE_HASH_H - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "sparse_growth_policy.h" - -#ifdef __INTEL_COMPILER -#include // For _popcnt32 and _popcnt64 -#endif - -#ifdef _MSC_VER -#include // For __cpuid, __popcnt and __popcnt64 -#endif - -#ifdef TSL_DEBUG -#define tsl_sh_assert(expr) assert(expr) -#else -#define tsl_sh_assert(expr) (static_cast(0)) -#endif - -namespace tsl { - -namespace sh { -enum class probing { linear, quadratic }; - -enum class exception_safety { basic, strong }; - -enum class sparsity { high, medium, low }; -} // namespace sh - -namespace detail_popcount { -/** - * Define the popcount(ll) methods and pick-up the best depending on the - * compiler. - */ - -// From Wikipedia: https://en.wikipedia.org/wiki/Hamming_weight -inline int fallback_popcountll(unsigned long long int x) { - static_assert( - sizeof(unsigned long long int) == sizeof(std::uint64_t), - "sizeof(unsigned long long int) must be equal to sizeof(std::uint64_t). " - "Open a feature request if you need support for a platform where it " - "isn't the case."); - - const std::uint64_t m1 = 0x5555555555555555ull; - const std::uint64_t m2 = 0x3333333333333333ull; - const std::uint64_t m4 = 0x0f0f0f0f0f0f0f0full; - const std::uint64_t h01 = 0x0101010101010101ull; - - x -= (x >> 1ull) & m1; - x = (x & m2) + ((x >> 2ull) & m2); - x = (x + (x >> 4ull)) & m4; - return static_cast((x * h01) >> (64ull - 8ull)); -} - -inline int fallback_popcount(unsigned int x) { - static_assert(sizeof(unsigned int) == sizeof(std::uint32_t) || - sizeof(unsigned int) == sizeof(std::uint64_t), - "sizeof(unsigned int) must be equal to sizeof(std::uint32_t) " - "or sizeof(std::uint64_t). " - "Open a feature request if you need support for a platform " - "where it isn't the case."); - - if (sizeof(unsigned int) == sizeof(std::uint32_t)) { - const std::uint32_t m1 = 0x55555555; - const std::uint32_t m2 = 0x33333333; - const std::uint32_t m4 = 0x0f0f0f0f; - const std::uint32_t h01 = 0x01010101; - - x -= (x >> 1) & m1; - x = (x & m2) + ((x >> 2) & m2); - x = (x + (x >> 4)) & m4; - return static_cast((x * h01) >> (32 - 8)); - } else { - return fallback_popcountll(x); - } -} - -#if defined(__clang__) || defined(__GNUC__) -inline int popcountll(unsigned long long int value) { - return __builtin_popcountll(value); -} - -inline int popcount(unsigned int value) { return __builtin_popcount(value); } - -#elif defined(_MSC_VER) -/** - * We need to check for popcount support at runtime on Windows with __cpuid - * See https://msdn.microsoft.com/en-us/library/bb385231.aspx - */ -inline bool has_popcount_support() { - int cpu_infos[4]; - __cpuid(cpu_infos, 1); - return (cpu_infos[2] & (1 << 23)) != 0; -} - -inline int popcountll(unsigned long long int value) { -#ifdef _WIN64 - static_assert( - sizeof(unsigned long long int) == sizeof(std::int64_t), - "sizeof(unsigned long long int) must be equal to sizeof(std::int64_t). "); - - static const bool has_popcount = has_popcount_support(); - return has_popcount - ? static_cast(__popcnt64(static_cast(value))) - : fallback_popcountll(value); -#else - return fallback_popcountll(value); -#endif -} - -inline int popcount(unsigned int value) { - static_assert(sizeof(unsigned int) == sizeof(std::int32_t), - "sizeof(unsigned int) must be equal to sizeof(std::int32_t). "); - - static const bool has_popcount = has_popcount_support(); - return has_popcount - ? static_cast(__popcnt(static_cast(value))) - : fallback_popcount(value); -} - -#elif defined(__INTEL_COMPILER) -inline int popcountll(unsigned long long int value) { - static_assert(sizeof(unsigned long long int) == sizeof(__int64), ""); - return _popcnt64(static_cast<__int64>(value)); -} - -inline int popcount(unsigned int value) { - return _popcnt32(static_cast(value)); -} - -#else -inline int popcountll(unsigned long long int x) { - return fallback_popcountll(x); -} - -inline int popcount(unsigned int x) { return fallback_popcount(x); } - -#endif -} // namespace detail_popcount - -namespace detail_sparse_hash { - -template -struct make_void { - using type = void; -}; - -template -struct has_is_transparent : std::false_type {}; - -template -struct has_is_transparent::type> - : std::true_type {}; - -template -struct is_power_of_two_policy : std::false_type {}; - -template -struct is_power_of_two_policy> - : std::true_type {}; - -inline constexpr bool is_power_of_two(std::size_t value) { - return value != 0 && (value & (value - 1)) == 0; -} - -inline std::size_t round_up_to_power_of_two(std::size_t value) { - if (is_power_of_two(value)) { - return value; - } - - if (value == 0) { - return 1; - } - - --value; - for (std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) { - value |= value >> i; - } - - return value + 1; -} - -template -static T numeric_cast(U value, - const char *error_message = "numeric_cast() failed.") { - T ret = static_cast(value); - if (static_cast(ret) != value) { - throw std::runtime_error(error_message); - } - - const bool is_same_signedness = - (std::is_unsigned::value && std::is_unsigned::value) || - (std::is_signed::value && std::is_signed::value); - if (!is_same_signedness && (ret < T{}) != (value < U{})) { - throw std::runtime_error(error_message); - } - - return ret; -} - -/** - * Fixed size type used to represent size_type values on serialization. Need to - * be big enough to represent a std::size_t on 32 and 64 bits platforms, and - * must be the same size on both platforms. - */ -using slz_size_type = std::uint64_t; -static_assert(std::numeric_limits::max() >= - std::numeric_limits::max(), - "slz_size_type must be >= std::size_t"); - -template -static T deserialize_value(Deserializer &deserializer) { - // MSVC < 2017 is not conformant, circumvent the problem by removing the - // template keyword -#if defined(_MSC_VER) && _MSC_VER < 1910 - return deserializer.Deserializer::operator()(); -#else - return deserializer.Deserializer::template operator()(); -#endif -} - -/** - * WARNING: the sparse_array class doesn't free the ressources allocated through - * the allocator passed in parameter in each method. You have to manually call - * `clear(Allocator&)` when you don't need a sparse_array object anymore. - * - * The reason is that the sparse_array doesn't store the allocator to avoid - * wasting space in each sparse_array when the allocator has a size > 0. It only - * allocates/deallocates objects with the allocator that is passed in parameter. - * - * - * - * Index denotes a value between [0, BITMAP_NB_BITS), it is an index similar to - * std::vector. Offset denotes the real position in `m_values` corresponding to - * an index. - * - * We are using raw pointers instead of std::vector to avoid loosing - * 2*sizeof(size_t) bytes to store the capacity and size of the vector in each - * sparse_array. We know we can only store up to BITMAP_NB_BITS elements in the - * array, we don't need such big types. - * - * - * T must be nothrow move constructible and/or copy constructible. - * Behaviour is undefined if the destructor of T throws an exception. - * - * See https://smerity.com/articles/2015/google_sparsehash.html for details on - * the idea behinds the implementation. - * - * TODO Check to use std::realloc and std::memmove when possible - */ -template -class sparse_array { - public: - using value_type = T; - using size_type = std::uint_least8_t; - using allocator_type = Allocator; - using iterator = value_type *; - using const_iterator = const value_type *; - - private: - static const size_type CAPACITY_GROWTH_STEP = - (Sparsity == tsl::sh::sparsity::high) ? 2 - : (Sparsity == tsl::sh::sparsity::medium) - ? 4 - : 8; // (Sparsity == tsl::sh::sparsity::low) - - /** - * Bitmap size configuration. - * Use 32 bits for the bitmap on 32-bits or less environnement as popcount on - * 64 bits numbers is slow on these environnement. Use 64 bits bitmap - * otherwise. - */ -#if SIZE_MAX <= UINT32_MAX - using bitmap_type = std::uint_least32_t; - static const std::size_t BITMAP_NB_BITS = 32; - static const std::size_t BUCKET_SHIFT = 5; -#else - using bitmap_type = std::uint_least64_t; - static const std::size_t BITMAP_NB_BITS = 64; - static const std::size_t BUCKET_SHIFT = 6; -#endif - - static const std::size_t BUCKET_MASK = BITMAP_NB_BITS - 1; - - static_assert(is_power_of_two(BITMAP_NB_BITS), - "BITMAP_NB_BITS must be a power of two."); - static_assert(std::numeric_limits::digits >= BITMAP_NB_BITS, - "bitmap_type must be able to hold at least BITMAP_NB_BITS."); - static_assert((std::size_t(1) << BUCKET_SHIFT) == BITMAP_NB_BITS, - "(1 << BUCKET_SHIFT) must be equal to BITMAP_NB_BITS."); - static_assert(std::numeric_limits::max() >= BITMAP_NB_BITS, - "size_type must be big enough to hold BITMAP_NB_BITS."); - static_assert(std::is_unsigned::value, - "bitmap_type must be unsigned."); - static_assert((std::numeric_limits::max() & BUCKET_MASK) == - BITMAP_NB_BITS - 1, - ""); - - public: - /** - * Map an ibucket [0, bucket_count) in the hash table to a sparse_ibucket - * (a sparse_array holds multiple buckets, so there is less sparse_array than - * bucket_count). - * - * The bucket ibucket is in - * m_sparse_buckets[sparse_ibucket(ibucket)][index_in_sparse_bucket(ibucket)] - * instead of something like m_buckets[ibucket] in a classical hash table. - */ - static std::size_t sparse_ibucket(std::size_t ibucket) { - return ibucket >> BUCKET_SHIFT; - } - - /** - * Map an ibucket [0, bucket_count) in the hash table to an index in the - * sparse_array which corresponds to the bucket. - * - * The bucket ibucket is in - * m_sparse_buckets[sparse_ibucket(ibucket)][index_in_sparse_bucket(ibucket)] - * instead of something like m_buckets[ibucket] in a classical hash table. - */ - static typename sparse_array::size_type index_in_sparse_bucket( - std::size_t ibucket) { - return static_cast( - ibucket & sparse_array::BUCKET_MASK); - } - - static std::size_t nb_sparse_buckets(std::size_t bucket_count) noexcept { - if (bucket_count == 0) { - return 0; - } - - return std::max( - 1, sparse_ibucket(tsl::detail_sparse_hash::round_up_to_power_of_two( - bucket_count))); - } - - public: - sparse_array() noexcept - : m_values(nullptr), - m_bitmap_vals(0), - m_bitmap_deleted_vals(0), - m_nb_elements(0), - m_capacity(0), - m_last_array(false) {} - - explicit sparse_array(bool last_bucket) noexcept - : m_values(nullptr), - m_bitmap_vals(0), - m_bitmap_deleted_vals(0), - m_nb_elements(0), - m_capacity(0), - m_last_array(last_bucket) {} - - sparse_array(size_type capacity, Allocator &alloc) - : m_values(nullptr), - m_bitmap_vals(0), - m_bitmap_deleted_vals(0), - m_nb_elements(0), - m_capacity(capacity), - m_last_array(false) { - if (m_capacity > 0) { - m_values = alloc.allocate(m_capacity); - tsl_sh_assert(m_values != - nullptr); // allocate should throw if there is a failure - } - } - - sparse_array(const sparse_array &other, Allocator &alloc) - : m_values(nullptr), - m_bitmap_vals(other.m_bitmap_vals), - m_bitmap_deleted_vals(other.m_bitmap_deleted_vals), - m_nb_elements(0), - m_capacity(other.m_capacity), - m_last_array(other.m_last_array) { - tsl_sh_assert(other.m_capacity >= other.m_nb_elements); - if (m_capacity == 0) { - return; - } - - m_values = alloc.allocate(m_capacity); - tsl_sh_assert(m_values != - nullptr); // allocate should throw if there is a failure - try { - for (size_type i = 0; i < other.m_nb_elements; i++) { - construct_value(alloc, m_values + i, other.m_values[i]); - m_nb_elements++; - } - } catch (...) { - clear(alloc); - throw; - } - } - - sparse_array(sparse_array &&other) noexcept - : m_values(other.m_values), - m_bitmap_vals(other.m_bitmap_vals), - m_bitmap_deleted_vals(other.m_bitmap_deleted_vals), - m_nb_elements(other.m_nb_elements), - m_capacity(other.m_capacity), - m_last_array(other.m_last_array) { - other.m_values = nullptr; - other.m_bitmap_vals = 0; - other.m_bitmap_deleted_vals = 0; - other.m_nb_elements = 0; - other.m_capacity = 0; - } - - sparse_array(sparse_array &&other, Allocator &alloc) - : m_values(nullptr), - m_bitmap_vals(other.m_bitmap_vals), - m_bitmap_deleted_vals(other.m_bitmap_deleted_vals), - m_nb_elements(0), - m_capacity(other.m_capacity), - m_last_array(other.m_last_array) { - tsl_sh_assert(other.m_capacity >= other.m_nb_elements); - if (m_capacity == 0) { - return; - } - - m_values = alloc.allocate(m_capacity); - tsl_sh_assert(m_values != - nullptr); // allocate should throw if there is a failure - try { - for (size_type i = 0; i < other.m_nb_elements; i++) { - construct_value(alloc, m_values + i, std::move(other.m_values[i])); - m_nb_elements++; - } - } catch (...) { - clear(alloc); - throw; - } - } - - sparse_array &operator=(const sparse_array &) = delete; - sparse_array &operator=(sparse_array &&) = delete; - - ~sparse_array() noexcept { - // The code that manages the sparse_array must have called clear before - // destruction. See documentation of sparse_array for more details. - tsl_sh_assert(m_capacity == 0 && m_nb_elements == 0 && m_values == nullptr); - } - - iterator begin() noexcept { return m_values; } - iterator end() noexcept { return m_values + m_nb_elements; } - const_iterator begin() const noexcept { return cbegin(); } - const_iterator end() const noexcept { return cend(); } - const_iterator cbegin() const noexcept { return m_values; } - const_iterator cend() const noexcept { return m_values + m_nb_elements; } - - bool empty() const noexcept { return m_nb_elements == 0; } - - size_type size() const noexcept { return m_nb_elements; } - - void clear(allocator_type &alloc) noexcept { - destroy_and_deallocate_values(alloc, m_values, m_nb_elements, m_capacity); - - m_values = nullptr; - m_bitmap_vals = 0; - m_bitmap_deleted_vals = 0; - m_nb_elements = 0; - m_capacity = 0; - } - - bool last() const noexcept { return m_last_array; } - - void set_as_last() noexcept { m_last_array = true; } - - bool has_value(size_type index) const noexcept { - tsl_sh_assert(index < BITMAP_NB_BITS); - return (m_bitmap_vals & (bitmap_type(1) << index)) != 0; - } - - bool has_deleted_value(size_type index) const noexcept { - tsl_sh_assert(index < BITMAP_NB_BITS); - return (m_bitmap_deleted_vals & (bitmap_type(1) << index)) != 0; - } - - iterator value(size_type index) noexcept { - tsl_sh_assert(has_value(index)); - return m_values + index_to_offset(index); - } - - const_iterator value(size_type index) const noexcept { - tsl_sh_assert(has_value(index)); - return m_values + index_to_offset(index); - } - - /** - * Return iterator to set value. - */ - template - iterator set(allocator_type &alloc, size_type index, Args &&...value_args) { - tsl_sh_assert(!has_value(index)); - - const size_type offset = index_to_offset(index); - insert_at_offset(alloc, offset, std::forward(value_args)...); - - m_bitmap_vals = (m_bitmap_vals | (bitmap_type(1) << index)); - m_bitmap_deleted_vals = - (m_bitmap_deleted_vals & ~(bitmap_type(1) << index)); - - m_nb_elements++; - - tsl_sh_assert(has_value(index)); - tsl_sh_assert(!has_deleted_value(index)); - - return m_values + offset; - } - - iterator erase(allocator_type &alloc, iterator position) { - const size_type offset = - static_cast(std::distance(begin(), position)); - return erase(alloc, position, offset_to_index(offset)); - } - - // Return the next value or end if no next value - iterator erase(allocator_type &alloc, iterator position, size_type index) { - tsl_sh_assert(has_value(index)); - tsl_sh_assert(!has_deleted_value(index)); - - const size_type offset = - static_cast(std::distance(begin(), position)); - erase_at_offset(alloc, offset); - - m_bitmap_vals = (m_bitmap_vals & ~(bitmap_type(1) << index)); - m_bitmap_deleted_vals = (m_bitmap_deleted_vals | (bitmap_type(1) << index)); - - m_nb_elements--; - - tsl_sh_assert(!has_value(index)); - tsl_sh_assert(has_deleted_value(index)); - - return m_values + offset; - } - - void swap(sparse_array &other) { - using std::swap; - - swap(m_values, other.m_values); - swap(m_bitmap_vals, other.m_bitmap_vals); - swap(m_bitmap_deleted_vals, other.m_bitmap_deleted_vals); - swap(m_nb_elements, other.m_nb_elements); - swap(m_capacity, other.m_capacity); - swap(m_last_array, other.m_last_array); - } - - static iterator mutable_iterator(const_iterator pos) { - return const_cast(pos); - } - - template - void serialize(Serializer &serializer) const { - const slz_size_type sparse_bucket_size = m_nb_elements; - serializer(sparse_bucket_size); - - const slz_size_type bitmap_vals = m_bitmap_vals; - serializer(bitmap_vals); - - const slz_size_type bitmap_deleted_vals = m_bitmap_deleted_vals; - serializer(bitmap_deleted_vals); - - for (const value_type &value : *this) { - serializer(value); - } - } - - template - static sparse_array deserialize_hash_compatible(Deserializer &deserializer, - Allocator &alloc) { - const slz_size_type sparse_bucket_size = - deserialize_value(deserializer); - const slz_size_type bitmap_vals = - deserialize_value(deserializer); - const slz_size_type bitmap_deleted_vals = - deserialize_value(deserializer); - - if (sparse_bucket_size > BITMAP_NB_BITS) { - throw std::runtime_error( - "Deserialized sparse_bucket_size is too big for the platform. " - "Maximum should be BITMAP_NB_BITS."); - } - - sparse_array sarray; - if (sparse_bucket_size == 0) { - return sarray; - } - - sarray.m_bitmap_vals = numeric_cast( - bitmap_vals, "Deserialized bitmap_vals is too big."); - sarray.m_bitmap_deleted_vals = numeric_cast( - bitmap_deleted_vals, "Deserialized bitmap_deleted_vals is too big."); - - sarray.m_capacity = numeric_cast( - sparse_bucket_size, "Deserialized sparse_bucket_size is too big."); - sarray.m_values = alloc.allocate(sarray.m_capacity); - - try { - for (size_type ivalue = 0; ivalue < sarray.m_capacity; ivalue++) { - construct_value(alloc, sarray.m_values + ivalue, - deserialize_value(deserializer)); - sarray.m_nb_elements++; - } - } catch (...) { - sarray.clear(alloc); - throw; - } - - return sarray; - } - - /** - * Deserialize the values of the bucket and insert them all in sparse_hash - * through sparse_hash.insert(...). - */ - template - static void deserialize_values_into_sparse_hash(Deserializer &deserializer, - SparseHash &sparse_hash) { - const slz_size_type sparse_bucket_size = - deserialize_value(deserializer); - - const slz_size_type bitmap_vals = - deserialize_value(deserializer); - static_cast(bitmap_vals); // Ignore, not needed - - const slz_size_type bitmap_deleted_vals = - deserialize_value(deserializer); - static_cast(bitmap_deleted_vals); // Ignore, not needed - - for (slz_size_type ivalue = 0; ivalue < sparse_bucket_size; ivalue++) { - sparse_hash.insert(deserialize_value(deserializer)); - } - } - - private: - template - static void construct_value(allocator_type &alloc, value_type *value, - Args &&...value_args) { - std::allocator_traits::construct( - alloc, value, std::forward(value_args)...); - } - - static void destroy_value(allocator_type &alloc, value_type *value) noexcept { - std::allocator_traits::destroy(alloc, value); - } - - static void destroy_and_deallocate_values( - allocator_type &alloc, value_type *values, size_type nb_values, - size_type capacity_values) noexcept { - for (size_type i = 0; i < nb_values; i++) { - destroy_value(alloc, values + i); - } - - alloc.deallocate(values, capacity_values); - } - - static size_type popcount(bitmap_type val) noexcept { - if (sizeof(bitmap_type) <= sizeof(unsigned int)) { - return static_cast( - tsl::detail_popcount::popcount(static_cast(val))); - } else { - return static_cast(tsl::detail_popcount::popcountll(val)); - } - } - - size_type index_to_offset(size_type index) const noexcept { - tsl_sh_assert(index < BITMAP_NB_BITS); - return popcount(m_bitmap_vals & - ((bitmap_type(1) << index) - bitmap_type(1))); - } - - // TODO optimize - size_type offset_to_index(size_type offset) const noexcept { - tsl_sh_assert(offset < m_nb_elements); - - bitmap_type bitmap_vals = m_bitmap_vals; - size_type index = 0; - size_type nb_ones = 0; - - while (bitmap_vals != 0) { - if ((bitmap_vals & 0x1) == 1) { - if (nb_ones == offset) { - break; - } - - nb_ones++; - } - - index++; - bitmap_vals = bitmap_vals >> 1; - } - - return index; - } - - size_type next_capacity() const noexcept { - return static_cast(m_capacity + CAPACITY_GROWTH_STEP); - } - - /** - * Insertion - * - * Two situations: - * - Either we are in a situation where - * std::is_nothrow_move_constructible::value is true. In this - * case, on insertion we just reallocate m_values when we reach its capacity - * (i.e. m_nb_elements == m_capacity), otherwise we just put the new value at - * its appropriate place. We can easily keep the strong exception guarantee as - * moving the values around is safe. - * - Otherwise we are in a situation where - * std::is_nothrow_move_constructible::value is false. In this - * case on EACH insertion we allocate a new area of m_nb_elements + 1 where we - * copy the values of m_values into it and put the new value there. On - * success, we set m_values to this new area. Even if slower, it's the only - * way to preserve to strong exception guarantee. - */ - template ::value>::type * = nullptr> - void insert_at_offset(allocator_type &alloc, size_type offset, - Args &&...value_args) { - if (m_nb_elements < m_capacity) { - insert_at_offset_no_realloc(alloc, offset, - std::forward(value_args)...); - } else { - insert_at_offset_realloc(alloc, offset, next_capacity(), - std::forward(value_args)...); - } - } - - template ::value>::type * = nullptr> - void insert_at_offset(allocator_type &alloc, size_type offset, - Args &&...value_args) { - insert_at_offset_realloc(alloc, offset, m_nb_elements + 1, - std::forward(value_args)...); - } - - template ::value>::type * = nullptr> - void insert_at_offset_no_realloc(allocator_type &alloc, size_type offset, - Args &&...value_args) { - tsl_sh_assert(offset <= m_nb_elements); - tsl_sh_assert(m_nb_elements < m_capacity); - - for (size_type i = m_nb_elements; i > offset; i--) { - construct_value(alloc, m_values + i, std::move(m_values[i - 1])); - destroy_value(alloc, m_values + i - 1); - } - - try { - construct_value(alloc, m_values + offset, - std::forward(value_args)...); - } catch (...) { - for (size_type i = offset; i < m_nb_elements; i++) { - construct_value(alloc, m_values + i, std::move(m_values[i + 1])); - destroy_value(alloc, m_values + i + 1); - } - throw; - } - } - - template ::value>::type * = nullptr> - void insert_at_offset_realloc(allocator_type &alloc, size_type offset, - size_type new_capacity, Args &&...value_args) { - tsl_sh_assert(new_capacity > m_nb_elements); - - value_type *new_values = alloc.allocate(new_capacity); - // Allocate should throw if there is a failure - tsl_sh_assert(new_values != nullptr); - - try { - construct_value(alloc, new_values + offset, - std::forward(value_args)...); - } catch (...) { - alloc.deallocate(new_values, new_capacity); - throw; - } - - // Should not throw from here - for (size_type i = 0; i < offset; i++) { - construct_value(alloc, new_values + i, std::move(m_values[i])); - } - - for (size_type i = offset; i < m_nb_elements; i++) { - construct_value(alloc, new_values + i + 1, std::move(m_values[i])); - } - - destroy_and_deallocate_values(alloc, m_values, m_nb_elements, m_capacity); - - m_values = new_values; - m_capacity = new_capacity; - } - - template ::value>::type * = nullptr> - void insert_at_offset_realloc(allocator_type &alloc, size_type offset, - size_type new_capacity, Args &&...value_args) { - tsl_sh_assert(new_capacity > m_nb_elements); - - value_type *new_values = alloc.allocate(new_capacity); - // Allocate should throw if there is a failure - tsl_sh_assert(new_values != nullptr); - - size_type nb_new_values = 0; - try { - for (size_type i = 0; i < offset; i++) { - construct_value(alloc, new_values + i, m_values[i]); - nb_new_values++; - } - - construct_value(alloc, new_values + offset, - std::forward(value_args)...); - nb_new_values++; - - for (size_type i = offset; i < m_nb_elements; i++) { - construct_value(alloc, new_values + i + 1, m_values[i]); - nb_new_values++; - } - } catch (...) { - destroy_and_deallocate_values(alloc, new_values, nb_new_values, - new_capacity); - throw; - } - - tsl_sh_assert(nb_new_values == m_nb_elements + 1); - - destroy_and_deallocate_values(alloc, m_values, m_nb_elements, m_capacity); - - m_values = new_values; - m_capacity = new_capacity; - } - - /** - * Erasure - * - * Two situations: - * - Either we are in a situation where - * std::is_nothrow_move_constructible::value is true. Simply - * destroy the value and left-shift move the value on the right of offset. - * - Otherwise we are in a situation where - * std::is_nothrow_move_constructible::value is false. Copy all - * the values except the one at offset into a new heap area. On success, we - * set m_values to this new area. Even if slower, it's the only way to - * preserve to strong exception guarantee. - */ - template ::value>::type * = nullptr> - void erase_at_offset(allocator_type &alloc, size_type offset) noexcept { - tsl_sh_assert(offset < m_nb_elements); - - destroy_value(alloc, m_values + offset); - - for (size_type i = offset + 1; i < m_nb_elements; i++) { - construct_value(alloc, m_values + i - 1, std::move(m_values[i])); - destroy_value(alloc, m_values + i); - } - } - - template ::value>::type * = nullptr> - void erase_at_offset(allocator_type &alloc, size_type offset) { - tsl_sh_assert(offset < m_nb_elements); - - // Erasing the last element, don't need to reallocate. We keep the capacity. - if (offset + 1 == m_nb_elements) { - destroy_value(alloc, m_values + offset); - return; - } - - tsl_sh_assert(m_nb_elements > 1); - const size_type new_capacity = m_nb_elements - 1; - - value_type *new_values = alloc.allocate(new_capacity); - // Allocate should throw if there is a failure - tsl_sh_assert(new_values != nullptr); - - size_type nb_new_values = 0; - try { - for (size_type i = 0; i < m_nb_elements; i++) { - if (i != offset) { - construct_value(alloc, new_values + nb_new_values, m_values[i]); - nb_new_values++; - } - } - } catch (...) { - destroy_and_deallocate_values(alloc, new_values, nb_new_values, - new_capacity); - throw; - } - - tsl_sh_assert(nb_new_values == m_nb_elements - 1); - - destroy_and_deallocate_values(alloc, m_values, m_nb_elements, m_capacity); - - m_values = new_values; - m_capacity = new_capacity; - } - - private: - value_type *m_values; - - bitmap_type m_bitmap_vals; - bitmap_type m_bitmap_deleted_vals; - - size_type m_nb_elements; - size_type m_capacity; - bool m_last_array; -}; - -/** - * Internal common class used by `sparse_map` and `sparse_set`. - * - * `ValueType` is what will be stored by `sparse_hash` (usually `std::pair` for map and `Key` for set). - * - * `KeySelect` should be a `FunctionObject` which takes a `ValueType` in - * parameter and returns a reference to the key. - * - * `ValueSelect` should be a `FunctionObject` which takes a `ValueType` in - * parameter and returns a reference to the value. `ValueSelect` should be void - * if there is no value (in a set for example). - * - * The strong exception guarantee only holds if `ExceptionSafety` is set to - * `tsl::sh::exception_safety::strong`. - * - * `ValueType` must be nothrow move constructible and/or copy constructible. - * Behaviour is undefined if the destructor of `ValueType` throws. - * - * - * The class holds its buckets in a 2-dimensional fashion. Instead of having a - * linear `std::vector` for [0, bucket_count) where each bucket stores - * one value, we have a `std::vector` (m_sparse_buckets_data) - * where each `sparse_array` stores multiple values (up to - * `sparse_array::BITMAP_NB_BITS`). To convert a one dimensional `ibucket` - * position to a position in `std::vector` and a position in - * `sparse_array`, use respectively the methods - * `sparse_array::sparse_ibucket(ibucket)` and - * `sparse_array::index_in_sparse_bucket(ibucket)`. - */ -template -class sparse_hash : private Allocator, - private Hash, - private KeyEqual, - private GrowthPolicy { - private: - template - using has_mapped_type = - typename std::integral_constant::value>; - - static_assert( - noexcept(std::declval().bucket_for_hash(std::size_t(0))), - "GrowthPolicy::bucket_for_hash must be noexcept."); - static_assert(noexcept(std::declval().clear()), - "GrowthPolicy::clear must be noexcept."); - - public: - template - class sparse_iterator; - - using key_type = typename KeySelect::key_type; - using value_type = ValueType; - using size_type = std::size_t; - using difference_type = std::ptrdiff_t; - using hasher = Hash; - using key_equal = KeyEqual; - using allocator_type = Allocator; - using reference = value_type &; - using const_reference = const value_type &; - using pointer = value_type *; - using const_pointer = const value_type *; - using iterator = sparse_iterator; - using const_iterator = sparse_iterator; - - private: - using sparse_array = - tsl::detail_sparse_hash::sparse_array; - - using sparse_buckets_allocator = typename std::allocator_traits< - allocator_type>::template rebind_alloc; - using sparse_buckets_container = - std::vector; - - public: - /** - * The `operator*()` and `operator->()` methods return a const reference and - * const pointer respectively to the stored value type (`Key` for a set, - * `std::pair` for a map). - * - * In case of a map, to get a mutable reference to the value `T` associated to - * a key (the `.second` in the stored pair), you have to call `value()`. - */ - template - class sparse_iterator { - friend class sparse_hash; - - private: - using sparse_bucket_iterator = typename std::conditional< - IsConst, typename sparse_buckets_container::const_iterator, - typename sparse_buckets_container::iterator>::type; - - using sparse_array_iterator = - typename std::conditional::type; - - /** - * sparse_array_it should be nullptr if sparse_bucket_it == - * m_sparse_buckets_data.end(). (TODO better way?) - */ - sparse_iterator(sparse_bucket_iterator sparse_bucket_it, - sparse_array_iterator sparse_array_it) - : m_sparse_buckets_it(sparse_bucket_it), - m_sparse_array_it(sparse_array_it) {} - - public: - using iterator_category = std::forward_iterator_tag; - using value_type = const typename sparse_hash::value_type; - using difference_type = std::ptrdiff_t; - using reference = value_type &; - using pointer = value_type *; - - sparse_iterator() noexcept {} - - // Copy constructor from iterator to const_iterator. - template ::type * = nullptr> - sparse_iterator(const sparse_iterator &other) noexcept - : m_sparse_buckets_it(other.m_sparse_buckets_it), - m_sparse_array_it(other.m_sparse_array_it) {} - - sparse_iterator(const sparse_iterator &other) = default; - sparse_iterator(sparse_iterator &&other) = default; - sparse_iterator &operator=(const sparse_iterator &other) = default; - sparse_iterator &operator=(sparse_iterator &&other) = default; - - const typename sparse_hash::key_type &key() const { - return KeySelect()(*m_sparse_array_it); - } - - template ::value && - IsConst>::type * = nullptr> - const typename U::value_type &value() const { - return U()(*m_sparse_array_it); - } - - template ::value && - !IsConst>::type * = nullptr> - typename U::value_type &value() { - return U()(*m_sparse_array_it); - } - - reference operator*() const { return *m_sparse_array_it; } - - pointer operator->() const { return std::addressof(*m_sparse_array_it); } - - sparse_iterator &operator++() { - tsl_sh_assert(m_sparse_array_it != nullptr); - ++m_sparse_array_it; - - if (m_sparse_array_it == m_sparse_buckets_it->end()) { - do { - if (m_sparse_buckets_it->last()) { - ++m_sparse_buckets_it; - m_sparse_array_it = nullptr; - return *this; - } - - ++m_sparse_buckets_it; - } while (m_sparse_buckets_it->empty()); - - m_sparse_array_it = m_sparse_buckets_it->begin(); - } - - return *this; - } - - sparse_iterator operator++(int) { - sparse_iterator tmp(*this); - ++*this; - - return tmp; - } - - friend bool operator==(const sparse_iterator &lhs, - const sparse_iterator &rhs) { - return lhs.m_sparse_buckets_it == rhs.m_sparse_buckets_it && - lhs.m_sparse_array_it == rhs.m_sparse_array_it; - } - - friend bool operator!=(const sparse_iterator &lhs, - const sparse_iterator &rhs) { - return !(lhs == rhs); - } - - private: - sparse_bucket_iterator m_sparse_buckets_it; - sparse_array_iterator m_sparse_array_it; - }; - - public: - sparse_hash(size_type bucket_count, const Hash &hash, const KeyEqual &equal, - const Allocator &alloc, float max_load_factor) - : Allocator(alloc), - Hash(hash), - KeyEqual(equal), - GrowthPolicy(bucket_count), - m_sparse_buckets_data(alloc), - m_sparse_buckets(static_empty_sparse_bucket_ptr()), - m_bucket_count(bucket_count), - m_nb_elements(0), - m_nb_deleted_buckets(0) { - if (m_bucket_count > max_bucket_count()) { - throw std::length_error("The map exceeds its maximum size."); - } - - if (m_bucket_count > 0) { - /* - * We can't use the `vector(size_type count, const Allocator& alloc)` - * constructor as it's only available in C++14 and we need to support - * C++11. We thus must resize after using the `vector(const Allocator& - * alloc)` constructor. - * - * We can't use `vector(size_type count, const T& value, const Allocator& - * alloc)` as it requires the value T to be copyable. - */ - m_sparse_buckets_data.resize( - sparse_array::nb_sparse_buckets(bucket_count)); - m_sparse_buckets = m_sparse_buckets_data.data(); - - tsl_sh_assert(!m_sparse_buckets_data.empty()); - m_sparse_buckets_data.back().set_as_last(); - } - - this->max_load_factor(max_load_factor); - - // Check in the constructor instead of outside of a function to avoid - // compilation issues when value_type is not complete. - static_assert(std::is_nothrow_move_constructible::value || - std::is_copy_constructible::value, - "Key, and T if present, must be nothrow move constructible " - "and/or copy constructible."); - } - - ~sparse_hash() { clear(); } - - sparse_hash(const sparse_hash &other) - : Allocator(std::allocator_traits< - Allocator>::select_on_container_copy_construction(other)), - Hash(other), - KeyEqual(other), - GrowthPolicy(other), - m_sparse_buckets_data( - std::allocator_traits< - Allocator>::select_on_container_copy_construction(other)), - m_bucket_count(other.m_bucket_count), - m_nb_elements(other.m_nb_elements), - m_nb_deleted_buckets(other.m_nb_deleted_buckets), - m_load_threshold_rehash(other.m_load_threshold_rehash), - m_load_threshold_clear_deleted(other.m_load_threshold_clear_deleted), - m_max_load_factor(other.m_max_load_factor) { - copy_buckets_from(other), - m_sparse_buckets = m_sparse_buckets_data.empty() - ? static_empty_sparse_bucket_ptr() - : m_sparse_buckets_data.data(); - } - - sparse_hash(sparse_hash &&other) noexcept( - std::is_nothrow_move_constructible::value - &&std::is_nothrow_move_constructible::value - &&std::is_nothrow_move_constructible::value - &&std::is_nothrow_move_constructible::value - &&std::is_nothrow_move_constructible< - sparse_buckets_container>::value) - : Allocator(std::move(other)), - Hash(std::move(other)), - KeyEqual(std::move(other)), - GrowthPolicy(std::move(other)), - m_sparse_buckets_data(std::move(other.m_sparse_buckets_data)), - m_sparse_buckets(m_sparse_buckets_data.empty() - ? static_empty_sparse_bucket_ptr() - : m_sparse_buckets_data.data()), - m_bucket_count(other.m_bucket_count), - m_nb_elements(other.m_nb_elements), - m_nb_deleted_buckets(other.m_nb_deleted_buckets), - m_load_threshold_rehash(other.m_load_threshold_rehash), - m_load_threshold_clear_deleted(other.m_load_threshold_clear_deleted), - m_max_load_factor(other.m_max_load_factor) { - other.GrowthPolicy::clear(); - other.m_sparse_buckets_data.clear(); - other.m_sparse_buckets = static_empty_sparse_bucket_ptr(); - other.m_bucket_count = 0; - other.m_nb_elements = 0; - other.m_nb_deleted_buckets = 0; - other.m_load_threshold_rehash = 0; - other.m_load_threshold_clear_deleted = 0; - } - - sparse_hash &operator=(const sparse_hash &other) { - if (this != &other) { - clear(); - - if (std::allocator_traits< - Allocator>::propagate_on_container_copy_assignment::value) { - Allocator::operator=(other); - } - - Hash::operator=(other); - KeyEqual::operator=(other); - GrowthPolicy::operator=(other); - - if (std::allocator_traits< - Allocator>::propagate_on_container_copy_assignment::value) { - m_sparse_buckets_data = - sparse_buckets_container(static_cast(other)); - } else { - if (m_sparse_buckets_data.size() != - other.m_sparse_buckets_data.size()) { - m_sparse_buckets_data = - sparse_buckets_container(static_cast(*this)); - } else { - m_sparse_buckets_data.clear(); - } - } - - copy_buckets_from(other); - m_sparse_buckets = m_sparse_buckets_data.empty() - ? static_empty_sparse_bucket_ptr() - : m_sparse_buckets_data.data(); - - m_bucket_count = other.m_bucket_count; - m_nb_elements = other.m_nb_elements; - m_nb_deleted_buckets = other.m_nb_deleted_buckets; - m_load_threshold_rehash = other.m_load_threshold_rehash; - m_load_threshold_clear_deleted = other.m_load_threshold_clear_deleted; - m_max_load_factor = other.m_max_load_factor; - } - - return *this; - } - - sparse_hash &operator=(sparse_hash &&other) { - clear(); - - if (std::allocator_traits< - Allocator>::propagate_on_container_move_assignment::value) { - static_cast(*this) = - std::move(static_cast(other)); - m_sparse_buckets_data = std::move(other.m_sparse_buckets_data); - } else if (static_cast(*this) != - static_cast(other)) { - move_buckets_from(std::move(other)); - } else { - static_cast(*this) = - std::move(static_cast(other)); - m_sparse_buckets_data = std::move(other.m_sparse_buckets_data); - } - - m_sparse_buckets = m_sparse_buckets_data.empty() - ? static_empty_sparse_bucket_ptr() - : m_sparse_buckets_data.data(); - - static_cast(*this) = std::move(static_cast(other)); - static_cast(*this) = std::move(static_cast(other)); - static_cast(*this) = - std::move(static_cast(other)); - m_bucket_count = other.m_bucket_count; - m_nb_elements = other.m_nb_elements; - m_nb_deleted_buckets = other.m_nb_deleted_buckets; - m_load_threshold_rehash = other.m_load_threshold_rehash; - m_load_threshold_clear_deleted = other.m_load_threshold_clear_deleted; - m_max_load_factor = other.m_max_load_factor; - - other.GrowthPolicy::clear(); - other.m_sparse_buckets_data.clear(); - other.m_sparse_buckets = static_empty_sparse_bucket_ptr(); - other.m_bucket_count = 0; - other.m_nb_elements = 0; - other.m_nb_deleted_buckets = 0; - other.m_load_threshold_rehash = 0; - other.m_load_threshold_clear_deleted = 0; - - return *this; - } - - allocator_type get_allocator() const { - return static_cast(*this); - } - - /* - * Iterators - */ - iterator begin() noexcept { - auto begin = m_sparse_buckets_data.begin(); - while (begin != m_sparse_buckets_data.end() && begin->empty()) { - ++begin; - } - - return iterator(begin, (begin != m_sparse_buckets_data.end()) - ? begin->begin() - : nullptr); - } - - const_iterator begin() const noexcept { return cbegin(); } - - const_iterator cbegin() const noexcept { - auto begin = m_sparse_buckets_data.cbegin(); - while (begin != m_sparse_buckets_data.cend() && begin->empty()) { - ++begin; - } - - return const_iterator(begin, (begin != m_sparse_buckets_data.cend()) - ? begin->cbegin() - : nullptr); - } - - iterator end() noexcept { - return iterator(m_sparse_buckets_data.end(), nullptr); - } - - const_iterator end() const noexcept { return cend(); } - - const_iterator cend() const noexcept { - return const_iterator(m_sparse_buckets_data.cend(), nullptr); - } - - /* - * Capacity - */ - bool empty() const noexcept { return m_nb_elements == 0; } - - size_type size() const noexcept { return m_nb_elements; } - - size_type max_size() const noexcept { - return std::min(std::allocator_traits::max_size(), - m_sparse_buckets_data.max_size()); - } - - /* - * Modifiers - */ - void clear() noexcept { - for (auto &bucket : m_sparse_buckets_data) { - bucket.clear(*this); - } - - m_nb_elements = 0; - m_nb_deleted_buckets = 0; - } - - template - std::pair insert(P &&value) { - return insert_impl(KeySelect()(value), std::forward

(value)); - } - - template - iterator insert_hint(const_iterator hint, P &&value) { - if (hint != cend() && - compare_keys(KeySelect()(*hint), KeySelect()(value))) { - return mutable_iterator(hint); - } - - return insert(std::forward

(value)).first; - } - - template - void insert(InputIt first, InputIt last) { - if (std::is_base_of< - std::forward_iterator_tag, - typename std::iterator_traits::iterator_category>::value) { - const auto nb_elements_insert = std::distance(first, last); - const size_type nb_free_buckets = m_load_threshold_rehash - size(); - tsl_sh_assert(m_load_threshold_rehash >= size()); - - if (nb_elements_insert > 0 && - nb_free_buckets < size_type(nb_elements_insert)) { - reserve(size() + size_type(nb_elements_insert)); - } - } - - for (; first != last; ++first) { - insert(*first); - } - } - - template - std::pair insert_or_assign(K &&key, M &&obj) { - auto it = try_emplace(std::forward(key), std::forward(obj)); - if (!it.second) { - it.first.value() = std::forward(obj); - } - - return it; - } - - template - iterator insert_or_assign(const_iterator hint, K &&key, M &&obj) { - if (hint != cend() && compare_keys(KeySelect()(*hint), key)) { - auto it = mutable_iterator(hint); - it.value() = std::forward(obj); - - return it; - } - - return insert_or_assign(std::forward(key), std::forward(obj)).first; - } - - template - std::pair emplace(Args &&...args) { - return insert(value_type(std::forward(args)...)); - } - - template - iterator emplace_hint(const_iterator hint, Args &&...args) { - return insert_hint(hint, value_type(std::forward(args)...)); - } - - template - std::pair try_emplace(K &&key, Args &&...args) { - return insert_impl(key, std::piecewise_construct, - std::forward_as_tuple(std::forward(key)), - std::forward_as_tuple(std::forward(args)...)); - } - - template - iterator try_emplace_hint(const_iterator hint, K &&key, Args &&...args) { - if (hint != cend() && compare_keys(KeySelect()(*hint), key)) { - return mutable_iterator(hint); - } - - return try_emplace(std::forward(key), std::forward(args)...).first; - } - - /** - * Here to avoid `template size_type erase(const K& key)` being used - * when we use an iterator instead of a const_iterator. - */ - iterator erase(iterator pos) { - tsl_sh_assert(pos != end() && m_nb_elements > 0); - auto it_sparse_array_next = - pos.m_sparse_buckets_it->erase(*this, pos.m_sparse_array_it); - m_nb_elements--; - m_nb_deleted_buckets++; - - if (it_sparse_array_next == pos.m_sparse_buckets_it->end()) { - auto it_sparse_buckets_next = pos.m_sparse_buckets_it; - do { - ++it_sparse_buckets_next; - } while (it_sparse_buckets_next != m_sparse_buckets_data.end() && - it_sparse_buckets_next->empty()); - - if (it_sparse_buckets_next == m_sparse_buckets_data.end()) { - return end(); - } else { - return iterator(it_sparse_buckets_next, - it_sparse_buckets_next->begin()); - } - } else { - return iterator(pos.m_sparse_buckets_it, it_sparse_array_next); - } - } - - iterator erase(const_iterator pos) { return erase(mutable_iterator(pos)); } - - iterator erase(const_iterator first, const_iterator last) { - if (first == last) { - return mutable_iterator(first); - } - - // TODO Optimize, could avoid the call to std::distance. - const size_type nb_elements_to_erase = - static_cast(std::distance(first, last)); - auto to_delete = mutable_iterator(first); - for (size_type i = 0; i < nb_elements_to_erase; i++) { - to_delete = erase(to_delete); - } - - return to_delete; - } - - template - size_type erase(const K &key) { - return erase(key, hash_key(key)); - } - - template - size_type erase(const K &key, std::size_t hash) { - return erase_impl(key, hash); - } - - void swap(sparse_hash &other) { - using std::swap; - - if (std::allocator_traits::propagate_on_container_swap::value) { - swap(static_cast(*this), static_cast(other)); - } else { - tsl_sh_assert(static_cast(*this) == - static_cast(other)); - } - - swap(static_cast(*this), static_cast(other)); - swap(static_cast(*this), static_cast(other)); - swap(static_cast(*this), - static_cast(other)); - swap(m_sparse_buckets_data, other.m_sparse_buckets_data); - swap(m_sparse_buckets, other.m_sparse_buckets); - swap(m_bucket_count, other.m_bucket_count); - swap(m_nb_elements, other.m_nb_elements); - swap(m_nb_deleted_buckets, other.m_nb_deleted_buckets); - swap(m_load_threshold_rehash, other.m_load_threshold_rehash); - swap(m_load_threshold_clear_deleted, other.m_load_threshold_clear_deleted); - swap(m_max_load_factor, other.m_max_load_factor); - } - - /* - * Lookup - */ - template < - class K, class U = ValueSelect, - typename std::enable_if::value>::type * = nullptr> - typename U::value_type &at(const K &key) { - return at(key, hash_key(key)); - } - - template < - class K, class U = ValueSelect, - typename std::enable_if::value>::type * = nullptr> - typename U::value_type &at(const K &key, std::size_t hash) { - return const_cast( - static_cast(this)->at(key, hash)); - } - - template < - class K, class U = ValueSelect, - typename std::enable_if::value>::type * = nullptr> - const typename U::value_type &at(const K &key) const { - return at(key, hash_key(key)); - } - - template < - class K, class U = ValueSelect, - typename std::enable_if::value>::type * = nullptr> - const typename U::value_type &at(const K &key, std::size_t hash) const { - auto it = find(key, hash); - if (it != cend()) { - return it.value(); - } else { - throw std::out_of_range("Couldn't find key."); - } - } - - template < - class K, class U = ValueSelect, - typename std::enable_if::value>::type * = nullptr> - typename U::value_type &operator[](K &&key) { - return try_emplace(std::forward(key)).first.value(); - } - - template - bool contains(const K &key) const { - return contains(key, hash_key(key)); - } - - template - bool contains(const K &key, std::size_t hash) const { - return count(key, hash) != 0; - } - - template - size_type count(const K &key) const { - return count(key, hash_key(key)); - } - - template - size_type count(const K &key, std::size_t hash) const { - if (find(key, hash) != cend()) { - return 1; - } else { - return 0; - } - } - - template - iterator find(const K &key) { - return find_impl(key, hash_key(key)); - } - - template - iterator find(const K &key, std::size_t hash) { - return find_impl(key, hash); - } - - template - const_iterator find(const K &key) const { - return find_impl(key, hash_key(key)); - } - - template - const_iterator find(const K &key, std::size_t hash) const { - return find_impl(key, hash); - } - - template - std::pair equal_range(const K &key) { - return equal_range(key, hash_key(key)); - } - - template - std::pair equal_range(const K &key, std::size_t hash) { - iterator it = find(key, hash); - return std::make_pair(it, (it == end()) ? it : std::next(it)); - } - - template - std::pair equal_range(const K &key) const { - return equal_range(key, hash_key(key)); - } - - template - std::pair equal_range( - const K &key, std::size_t hash) const { - const_iterator it = find(key, hash); - return std::make_pair(it, (it == cend()) ? it : std::next(it)); - } - - /* - * Bucket interface - */ - size_type bucket_count() const { return m_bucket_count; } - - size_type max_bucket_count() const { - return m_sparse_buckets_data.max_size(); - } - - /* - * Hash policy - */ - float load_factor() const { - if (bucket_count() == 0) { - return 0; - } - - return float(m_nb_elements) / float(bucket_count()); - } - - float max_load_factor() const { return m_max_load_factor; } - - void max_load_factor(float ml) { - m_max_load_factor = std::max(0.1f, std::min(ml, 0.8f)); - m_load_threshold_rehash = - size_type(float(bucket_count()) * m_max_load_factor); - - const float max_load_factor_with_deleted_buckets = - m_max_load_factor + 0.5f * (1.0f - m_max_load_factor); - tsl_sh_assert(max_load_factor_with_deleted_buckets > 0.0f && - max_load_factor_with_deleted_buckets <= 1.0f); - m_load_threshold_clear_deleted = - size_type(float(bucket_count()) * max_load_factor_with_deleted_buckets); - } - - void rehash(size_type count) { - count = std::max(count, - size_type(std::ceil(float(size()) / max_load_factor()))); - rehash_impl(count); - } - - void reserve(size_type count) { - rehash(size_type(std::ceil(float(count) / max_load_factor()))); - } - - /* - * Observers - */ - hasher hash_function() const { return static_cast(*this); } - - key_equal key_eq() const { return static_cast(*this); } - - /* - * Other - */ - iterator mutable_iterator(const_iterator pos) { - auto it_sparse_buckets = - m_sparse_buckets_data.begin() + - std::distance(m_sparse_buckets_data.cbegin(), pos.m_sparse_buckets_it); - - return iterator(it_sparse_buckets, - sparse_array::mutable_iterator(pos.m_sparse_array_it)); - } - - template - void serialize(Serializer &serializer) const { - serialize_impl(serializer); - } - - template - void deserialize(Deserializer &deserializer, bool hash_compatible) { - deserialize_impl(deserializer, hash_compatible); - } - - private: - template - std::size_t hash_key(const K &key) const { - return Hash::operator()(key); - } - - template - bool compare_keys(const K1 &key1, const K2 &key2) const { - return KeyEqual::operator()(key1, key2); - } - - size_type bucket_for_hash(std::size_t hash) const { - const std::size_t bucket = GrowthPolicy::bucket_for_hash(hash); - tsl_sh_assert(sparse_array::sparse_ibucket(bucket) < - m_sparse_buckets_data.size() || - (bucket == 0 && m_sparse_buckets_data.empty())); - - return bucket; - } - - template ::value>::type * = - nullptr> - size_type next_bucket(size_type ibucket, size_type iprobe) const { - (void)iprobe; - if (Probing == tsl::sh::probing::linear) { - return (ibucket + 1) & this->m_mask; - } else { - tsl_sh_assert(Probing == tsl::sh::probing::quadratic); - return (ibucket + iprobe) & this->m_mask; - } - } - - template ::value>::type * = - nullptr> - size_type next_bucket(size_type ibucket, size_type iprobe) const { - (void)iprobe; - if (Probing == tsl::sh::probing::linear) { - ibucket++; - return (ibucket != bucket_count()) ? ibucket : 0; - } else { - tsl_sh_assert(Probing == tsl::sh::probing::quadratic); - ibucket += iprobe; - return (ibucket < bucket_count()) ? ibucket : ibucket % bucket_count(); - } - } - - // TODO encapsulate m_sparse_buckets_data to avoid the managing the allocator - void copy_buckets_from(const sparse_hash &other) { - m_sparse_buckets_data.reserve(other.m_sparse_buckets_data.size()); - - try { - for (const auto &bucket : other.m_sparse_buckets_data) { - m_sparse_buckets_data.emplace_back(bucket, - static_cast(*this)); - } - } catch (...) { - clear(); - throw; - } - - tsl_sh_assert(m_sparse_buckets_data.empty() || - m_sparse_buckets_data.back().last()); - } - - void move_buckets_from(sparse_hash &&other) { - m_sparse_buckets_data.reserve(other.m_sparse_buckets_data.size()); - - try { - for (auto &&bucket : other.m_sparse_buckets_data) { - m_sparse_buckets_data.emplace_back(std::move(bucket), - static_cast(*this)); - } - } catch (...) { - clear(); - throw; - } - - tsl_sh_assert(m_sparse_buckets_data.empty() || - m_sparse_buckets_data.back().last()); - } - - template - std::pair insert_impl(const K &key, - Args &&...value_type_args) { - if (size() >= m_load_threshold_rehash) { - rehash_impl(GrowthPolicy::next_bucket_count()); - } else if (size() + m_nb_deleted_buckets >= - m_load_threshold_clear_deleted) { - clear_deleted_buckets(); - } - tsl_sh_assert(!m_sparse_buckets_data.empty()); - - /** - * We must insert the value in the first empty or deleted bucket we find. If - * we first find a deleted bucket, we still have to continue the search - * until we find an empty bucket or until we have searched all the buckets - * to be sure that the value is not in the hash table. We thus remember the - * position, if any, of the first deleted bucket we have encountered so we - * can insert it there if needed. - */ - bool found_first_deleted_bucket = false; - std::size_t sparse_ibucket_first_deleted = 0; - typename sparse_array::size_type index_in_sparse_bucket_first_deleted = 0; - - const std::size_t hash = hash_key(key); - std::size_t ibucket = bucket_for_hash(hash); - - std::size_t probe = 0; - while (true) { - std::size_t sparse_ibucket = sparse_array::sparse_ibucket(ibucket); - auto index_in_sparse_bucket = - sparse_array::index_in_sparse_bucket(ibucket); - - if (m_sparse_buckets[sparse_ibucket].has_value(index_in_sparse_bucket)) { - auto value_it = - m_sparse_buckets[sparse_ibucket].value(index_in_sparse_bucket); - if (compare_keys(key, KeySelect()(*value_it))) { - return std::make_pair( - iterator(m_sparse_buckets_data.begin() + sparse_ibucket, - value_it), - false); - } - } else if (m_sparse_buckets[sparse_ibucket].has_deleted_value( - index_in_sparse_bucket) && - probe < m_bucket_count) { - if (!found_first_deleted_bucket) { - found_first_deleted_bucket = true; - sparse_ibucket_first_deleted = sparse_ibucket; - index_in_sparse_bucket_first_deleted = index_in_sparse_bucket; - } - } else if (found_first_deleted_bucket) { - auto it = insert_in_bucket(sparse_ibucket_first_deleted, - index_in_sparse_bucket_first_deleted, - std::forward(value_type_args)...); - m_nb_deleted_buckets--; - - return it; - } else { - return insert_in_bucket(sparse_ibucket, index_in_sparse_bucket, - std::forward(value_type_args)...); - } - - probe++; - ibucket = next_bucket(ibucket, probe); - } - } - - template - std::pair insert_in_bucket( - std::size_t sparse_ibucket, - typename sparse_array::size_type index_in_sparse_bucket, - Args &&...value_type_args) { - auto value_it = m_sparse_buckets[sparse_ibucket].set( - *this, index_in_sparse_bucket, std::forward(value_type_args)...); - m_nb_elements++; - - return std::make_pair( - iterator(m_sparse_buckets_data.begin() + sparse_ibucket, value_it), - true); - } - - template - size_type erase_impl(const K &key, std::size_t hash) { - std::size_t ibucket = bucket_for_hash(hash); - - std::size_t probe = 0; - while (true) { - const std::size_t sparse_ibucket = sparse_array::sparse_ibucket(ibucket); - const auto index_in_sparse_bucket = - sparse_array::index_in_sparse_bucket(ibucket); - - if (m_sparse_buckets[sparse_ibucket].has_value(index_in_sparse_bucket)) { - auto value_it = - m_sparse_buckets[sparse_ibucket].value(index_in_sparse_bucket); - if (compare_keys(key, KeySelect()(*value_it))) { - m_sparse_buckets[sparse_ibucket].erase(*this, value_it, - index_in_sparse_bucket); - m_nb_elements--; - m_nb_deleted_buckets++; - - return 1; - } - } else if (!m_sparse_buckets[sparse_ibucket].has_deleted_value( - index_in_sparse_bucket) || - probe >= m_bucket_count) { - return 0; - } - - probe++; - ibucket = next_bucket(ibucket, probe); - } - } - - template - iterator find_impl(const K &key, std::size_t hash) { - return mutable_iterator( - static_cast(this)->find(key, hash)); - } - - template - const_iterator find_impl(const K &key, std::size_t hash) const { - std::size_t ibucket = bucket_for_hash(hash); - - std::size_t probe = 0; - while (true) { - const std::size_t sparse_ibucket = sparse_array::sparse_ibucket(ibucket); - const auto index_in_sparse_bucket = - sparse_array::index_in_sparse_bucket(ibucket); - - if (m_sparse_buckets[sparse_ibucket].has_value(index_in_sparse_bucket)) { - auto value_it = - m_sparse_buckets[sparse_ibucket].value(index_in_sparse_bucket); - if (compare_keys(key, KeySelect()(*value_it))) { - return const_iterator(m_sparse_buckets_data.cbegin() + sparse_ibucket, - value_it); - } - } else if (!m_sparse_buckets[sparse_ibucket].has_deleted_value( - index_in_sparse_bucket) || - probe >= m_bucket_count) { - return cend(); - } - - probe++; - ibucket = next_bucket(ibucket, probe); - } - } - - void clear_deleted_buckets() { - // TODO could be optimized, we could do it in-place instead of allocating a - // new bucket array. - rehash_impl(m_bucket_count); - tsl_sh_assert(m_nb_deleted_buckets == 0); - } - - template ::type - * = nullptr> - void rehash_impl(size_type count) { - sparse_hash new_table(count, static_cast(*this), - static_cast(*this), - static_cast(*this), m_max_load_factor); - - for (auto &bucket : m_sparse_buckets_data) { - for (auto &val : bucket) { - new_table.insert_on_rehash(std::move(val)); - } - - // TODO try to reuse some of the memory - bucket.clear(*this); - } - - new_table.swap(*this); - } - - /** - * TODO: For now we copy each element into the new map. We could move - * them if they are nothrow_move_constructible without triggering - * any exception if we reserve enough space in the sparse arrays beforehand. - */ - template ::type * = nullptr> - void rehash_impl(size_type count) { - sparse_hash new_table(count, static_cast(*this), - static_cast(*this), - static_cast(*this), m_max_load_factor); - - for (const auto &bucket : m_sparse_buckets_data) { - for (const auto &val : bucket) { - new_table.insert_on_rehash(val); - } - } - - new_table.swap(*this); - } - - template - void insert_on_rehash(K &&key_value) { - const key_type &key = KeySelect()(key_value); - - const std::size_t hash = hash_key(key); - std::size_t ibucket = bucket_for_hash(hash); - - std::size_t probe = 0; - while (true) { - std::size_t sparse_ibucket = sparse_array::sparse_ibucket(ibucket); - auto index_in_sparse_bucket = - sparse_array::index_in_sparse_bucket(ibucket); - - if (!m_sparse_buckets[sparse_ibucket].has_value(index_in_sparse_bucket)) { - m_sparse_buckets[sparse_ibucket].set(*this, index_in_sparse_bucket, - std::forward(key_value)); - m_nb_elements++; - - return; - } else { - tsl_sh_assert(!compare_keys( - key, KeySelect()(*m_sparse_buckets[sparse_ibucket].value( - index_in_sparse_bucket)))); - } - - probe++; - ibucket = next_bucket(ibucket, probe); - } - } - - template - void serialize_impl(Serializer &serializer) const { - const slz_size_type version = SERIALIZATION_PROTOCOL_VERSION; - serializer(version); - - const slz_size_type bucket_count = m_bucket_count; - serializer(bucket_count); - - const slz_size_type nb_sparse_buckets = m_sparse_buckets_data.size(); - serializer(nb_sparse_buckets); - - const slz_size_type nb_elements = m_nb_elements; - serializer(nb_elements); - - const slz_size_type nb_deleted_buckets = m_nb_deleted_buckets; - serializer(nb_deleted_buckets); - - const float max_load_factor = m_max_load_factor; - serializer(max_load_factor); - - for (const auto &bucket : m_sparse_buckets_data) { - bucket.serialize(serializer); - } - } - - template - void deserialize_impl(Deserializer &deserializer, bool hash_compatible) { - tsl_sh_assert( - m_bucket_count == 0 && - m_sparse_buckets_data.empty()); // Current hash table must be empty - - const slz_size_type version = - deserialize_value(deserializer); - // For now we only have one version of the serialization protocol. - // If it doesn't match there is a problem with the file. - if (version != SERIALIZATION_PROTOCOL_VERSION) { - throw std::runtime_error( - "Can't deserialize the sparse_map/set. The " - "protocol version header is invalid."); - } - - const slz_size_type bucket_count_ds = - deserialize_value(deserializer); - const slz_size_type nb_sparse_buckets = - deserialize_value(deserializer); - const slz_size_type nb_elements = - deserialize_value(deserializer); - const slz_size_type nb_deleted_buckets = - deserialize_value(deserializer); - const float max_load_factor = deserialize_value(deserializer); - - if (!hash_compatible) { - this->max_load_factor(max_load_factor); - reserve(numeric_cast(nb_elements, - "Deserialized nb_elements is too big.")); - for (slz_size_type ibucket = 0; ibucket < nb_sparse_buckets; ibucket++) { - sparse_array::deserialize_values_into_sparse_hash(deserializer, *this); - } - } else { - m_bucket_count = numeric_cast( - bucket_count_ds, "Deserialized bucket_count is too big."); - - GrowthPolicy::operator=(GrowthPolicy(m_bucket_count)); - // GrowthPolicy should not modify the bucket count we got from - // deserialization - if (m_bucket_count != bucket_count_ds) { - throw std::runtime_error( - "The GrowthPolicy is not the same even though " - "hash_compatible is true."); - } - - if (nb_sparse_buckets != - sparse_array::nb_sparse_buckets(m_bucket_count)) { - throw std::runtime_error("Deserialized nb_sparse_buckets is invalid."); - } - - m_nb_elements = numeric_cast( - nb_elements, "Deserialized nb_elements is too big."); - m_nb_deleted_buckets = numeric_cast( - nb_deleted_buckets, "Deserialized nb_deleted_buckets is too big."); - - m_sparse_buckets_data.reserve(numeric_cast( - nb_sparse_buckets, "Deserialized nb_sparse_buckets is too big.")); - for (slz_size_type ibucket = 0; ibucket < nb_sparse_buckets; ibucket++) { - m_sparse_buckets_data.emplace_back( - sparse_array::deserialize_hash_compatible( - deserializer, static_cast(*this))); - } - - if (!m_sparse_buckets_data.empty()) { - m_sparse_buckets_data.back().set_as_last(); - m_sparse_buckets = m_sparse_buckets_data.data(); - } - - this->max_load_factor(max_load_factor); - if (load_factor() > this->max_load_factor()) { - throw std::runtime_error( - "Invalid max_load_factor. Check that the serializer and " - "deserializer support " - "floats correctly as they can be converted implicitely to ints."); - } - } - } - - public: - static const size_type DEFAULT_INIT_BUCKET_COUNT = 0; - static constexpr float DEFAULT_MAX_LOAD_FACTOR = 0.5f; - - /** - * Protocol version currenlty used for serialization. - */ - static const slz_size_type SERIALIZATION_PROTOCOL_VERSION = 1; - - /** - * Return an always valid pointer to an static empty bucket_entry with - * last_bucket() == true. - */ - sparse_array *static_empty_sparse_bucket_ptr() { - static sparse_array empty_sparse_bucket(true); - return &empty_sparse_bucket; - } - - private: - sparse_buckets_container m_sparse_buckets_data; - - /** - * Points to m_sparse_buckets_data.data() if !m_sparse_buckets_data.empty() - * otherwise points to static_empty_sparse_bucket_ptr. This variable is useful - * to avoid the cost of checking if m_sparse_buckets_data is empty when trying - * to find an element. - * - * TODO Remove m_sparse_buckets_data and only use a pointer instead of a - * pointer+vector to save some space in the sparse_hash object. - */ - sparse_array *m_sparse_buckets; - - size_type m_bucket_count; - size_type m_nb_elements; - size_type m_nb_deleted_buckets; - - /** - * Maximum that m_nb_elements can reach before a rehash occurs automatically - * to grow the hash table. - */ - size_type m_load_threshold_rehash; - - /** - * Maximum that m_nb_elements + m_nb_deleted_buckets can reach before cleaning - * up the buckets marked as deleted. - */ - size_type m_load_threshold_clear_deleted; - float m_max_load_factor; -}; - -} // namespace detail_sparse_hash -} // namespace tsl - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_map.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_map.h deleted file mode 100644 index 601742d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_map.h +++ /dev/null @@ -1,800 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Thibaut Goetghebuer-Planchon - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_SPARSE_MAP_H -#define TSL_SPARSE_MAP_H - -#include -#include -#include -#include -#include -#include - -#include "sparse_hash.h" - -namespace tsl { - -/** - * Implementation of a sparse hash map using open-addressing with quadratic - * probing. The goal on the hash map is to be the most memory efficient - * possible, even at low load factor, while keeping reasonable performances. - * - * `GrowthPolicy` defines how the map grows and consequently how a hash value is - * mapped to a bucket. By default the map uses - * `tsl::sh::power_of_two_growth_policy`. This policy keeps the number of - * buckets to a power of two and uses a mask to map the hash to a bucket instead - * of the slow modulo. Other growth policies are available and you may define - * your own growth policy, check `tsl::sh::power_of_two_growth_policy` for the - * interface. - * - * `ExceptionSafety` defines the exception guarantee provided by the class. By - * default only the basic exception safety is guaranteed which mean that all - * resources used by the hash map will be freed (no memory leaks) but the hash - * map may end-up in an undefined state if an exception is thrown (undefined - * here means that some elements may be missing). This can ONLY happen on rehash - * (either on insert or if `rehash` is called explicitly) and will occur if the - * Allocator can't allocate memory (`std::bad_alloc`) or if the copy constructor - * (when a nothrow move constructor is not available) throws an exception. This - * can be avoided by calling `reserve` beforehand. This basic guarantee is - * similar to the one of `google::sparse_hash_map` and `spp::sparse_hash_map`. - * It is possible to ask for the strong exception guarantee with - * `tsl::sh::exception_safety::strong`, the drawback is that the map will be - * slower on rehashes and will also need more memory on rehashes. - * - * `Sparsity` defines how much the hash set will compromise between insertion - * speed and memory usage. A high sparsity means less memory usage but longer - * insertion times, and vice-versa for low sparsity. The default - * `tsl::sh::sparsity::medium` sparsity offers a good compromise. It doesn't - * change the lookup speed. - * - * `Key` and `T` must be nothrow move constructible and/or copy constructible. - * - * If the destructor of `Key` or `T` throws an exception, the behaviour of the - * class is undefined. - * - * Iterators invalidation: - * - clear, operator=, reserve, rehash: always invalidate the iterators. - * - insert, emplace, emplace_hint, operator[]: if there is an effective - * insert, invalidate the iterators. - * - erase: always invalidate the iterators. - */ -template , - class KeyEqual = std::equal_to, - class Allocator = std::allocator>, - class GrowthPolicy = tsl::sh::power_of_two_growth_policy<2>, - tsl::sh::exception_safety ExceptionSafety = - tsl::sh::exception_safety::basic, - tsl::sh::sparsity Sparsity = tsl::sh::sparsity::medium> -class sparse_map { - private: - template - using has_is_transparent = tsl::detail_sparse_hash::has_is_transparent; - - class KeySelect { - public: - using key_type = Key; - - const key_type &operator()( - const std::pair &key_value) const noexcept { - return key_value.first; - } - - key_type &operator()(std::pair &key_value) noexcept { - return key_value.first; - } - }; - - class ValueSelect { - public: - using value_type = T; - - const value_type &operator()( - const std::pair &key_value) const noexcept { - return key_value.second; - } - - value_type &operator()(std::pair &key_value) noexcept { - return key_value.second; - } - }; - - using ht = detail_sparse_hash::sparse_hash< - std::pair, KeySelect, ValueSelect, Hash, KeyEqual, Allocator, - GrowthPolicy, ExceptionSafety, Sparsity, tsl::sh::probing::quadratic>; - - public: - using key_type = typename ht::key_type; - using mapped_type = T; - using value_type = typename ht::value_type; - using size_type = typename ht::size_type; - using difference_type = typename ht::difference_type; - using hasher = typename ht::hasher; - using key_equal = typename ht::key_equal; - using allocator_type = typename ht::allocator_type; - using reference = typename ht::reference; - using const_reference = typename ht::const_reference; - using pointer = typename ht::pointer; - using const_pointer = typename ht::const_pointer; - using iterator = typename ht::iterator; - using const_iterator = typename ht::const_iterator; - - public: - /* - * Constructors - */ - sparse_map() : sparse_map(ht::DEFAULT_INIT_BUCKET_COUNT) {} - - explicit sparse_map(size_type bucket_count, const Hash &hash = Hash(), - const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR) {} - - sparse_map(size_type bucket_count, const Allocator &alloc) - : sparse_map(bucket_count, Hash(), KeyEqual(), alloc) {} - - sparse_map(size_type bucket_count, const Hash &hash, const Allocator &alloc) - : sparse_map(bucket_count, hash, KeyEqual(), alloc) {} - - explicit sparse_map(const Allocator &alloc) - : sparse_map(ht::DEFAULT_INIT_BUCKET_COUNT, alloc) {} - - template - sparse_map(InputIt first, InputIt last, - size_type bucket_count = ht::DEFAULT_INIT_BUCKET_COUNT, - const Hash &hash = Hash(), const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : sparse_map(bucket_count, hash, equal, alloc) { - insert(first, last); - } - - template - sparse_map(InputIt first, InputIt last, size_type bucket_count, - const Allocator &alloc) - : sparse_map(first, last, bucket_count, Hash(), KeyEqual(), alloc) {} - - template - sparse_map(InputIt first, InputIt last, size_type bucket_count, - const Hash &hash, const Allocator &alloc) - : sparse_map(first, last, bucket_count, hash, KeyEqual(), alloc) {} - - sparse_map(std::initializer_list init, - size_type bucket_count = ht::DEFAULT_INIT_BUCKET_COUNT, - const Hash &hash = Hash(), const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : sparse_map(init.begin(), init.end(), bucket_count, hash, equal, alloc) { - } - - sparse_map(std::initializer_list init, size_type bucket_count, - const Allocator &alloc) - : sparse_map(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), - alloc) {} - - sparse_map(std::initializer_list init, size_type bucket_count, - const Hash &hash, const Allocator &alloc) - : sparse_map(init.begin(), init.end(), bucket_count, hash, KeyEqual(), - alloc) {} - - sparse_map &operator=(std::initializer_list ilist) { - m_ht.clear(); - - m_ht.reserve(ilist.size()); - m_ht.insert(ilist.begin(), ilist.end()); - - return *this; - } - - allocator_type get_allocator() const { return m_ht.get_allocator(); } - - /* - * Iterators - */ - iterator begin() noexcept { return m_ht.begin(); } - const_iterator begin() const noexcept { return m_ht.begin(); } - const_iterator cbegin() const noexcept { return m_ht.cbegin(); } - - iterator end() noexcept { return m_ht.end(); } - const_iterator end() const noexcept { return m_ht.end(); } - const_iterator cend() const noexcept { return m_ht.cend(); } - - /* - * Capacity - */ - bool empty() const noexcept { return m_ht.empty(); } - size_type size() const noexcept { return m_ht.size(); } - size_type max_size() const noexcept { return m_ht.max_size(); } - - /* - * Modifiers - */ - void clear() noexcept { m_ht.clear(); } - - std::pair insert(const value_type &value) { - return m_ht.insert(value); - } - - template ::value>::type * = nullptr> - std::pair insert(P &&value) { - return m_ht.emplace(std::forward

(value)); - } - - std::pair insert(value_type &&value) { - return m_ht.insert(std::move(value)); - } - - iterator insert(const_iterator hint, const value_type &value) { - return m_ht.insert_hint(hint, value); - } - - template ::value>::type * = nullptr> - iterator insert(const_iterator hint, P &&value) { - return m_ht.emplace_hint(hint, std::forward

(value)); - } - - iterator insert(const_iterator hint, value_type &&value) { - return m_ht.insert_hint(hint, std::move(value)); - } - - template - void insert(InputIt first, InputIt last) { - m_ht.insert(first, last); - } - - void insert(std::initializer_list ilist) { - m_ht.insert(ilist.begin(), ilist.end()); - } - - template - std::pair insert_or_assign(const key_type &k, M &&obj) { - return m_ht.insert_or_assign(k, std::forward(obj)); - } - - template - std::pair insert_or_assign(key_type &&k, M &&obj) { - return m_ht.insert_or_assign(std::move(k), std::forward(obj)); - } - - template - iterator insert_or_assign(const_iterator hint, const key_type &k, M &&obj) { - return m_ht.insert_or_assign(hint, k, std::forward(obj)); - } - - template - iterator insert_or_assign(const_iterator hint, key_type &&k, M &&obj) { - return m_ht.insert_or_assign(hint, std::move(k), std::forward(obj)); - } - - /** - * Due to the way elements are stored, emplace will need to move or copy the - * key-value once. The method is equivalent to - * `insert(value_type(std::forward(args)...));`. - * - * Mainly here for compatibility with the `std::unordered_map` interface. - */ - template - std::pair emplace(Args &&...args) { - return m_ht.emplace(std::forward(args)...); - } - - /** - * Due to the way elements are stored, emplace_hint will need to move or copy - * the key-value once. The method is equivalent to `insert(hint, - * value_type(std::forward(args)...));`. - * - * Mainly here for compatibility with the `std::unordered_map` interface. - */ - template - iterator emplace_hint(const_iterator hint, Args &&...args) { - return m_ht.emplace_hint(hint, std::forward(args)...); - } - - template - std::pair try_emplace(const key_type &k, Args &&...args) { - return m_ht.try_emplace(k, std::forward(args)...); - } - - template - std::pair try_emplace(key_type &&k, Args &&...args) { - return m_ht.try_emplace(std::move(k), std::forward(args)...); - } - - template - iterator try_emplace(const_iterator hint, const key_type &k, Args &&...args) { - return m_ht.try_emplace_hint(hint, k, std::forward(args)...); - } - - template - iterator try_emplace(const_iterator hint, key_type &&k, Args &&...args) { - return m_ht.try_emplace_hint(hint, std::move(k), - std::forward(args)...); - } - - iterator erase(iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator first, const_iterator last) { - return m_ht.erase(first, last); - } - size_type erase(const key_type &key) { return m_ht.erase(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - size_type erase(const key_type &key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type erase(const K &key) { - return m_ht.erase(key); - } - - /** - * @copydoc erase(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type erase(const K &key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - void swap(sparse_map &other) { other.m_ht.swap(m_ht); } - - /* - * Lookup - */ - T &at(const Key &key) { return m_ht.at(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - T &at(const Key &key, std::size_t precalculated_hash) { - return m_ht.at(key, precalculated_hash); - } - - const T &at(const Key &key) const { return m_ht.at(key); } - - /** - * @copydoc at(const Key& key, std::size_t precalculated_hash) - */ - const T &at(const Key &key, std::size_t precalculated_hash) const { - return m_ht.at(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - T &at(const K &key) { - return m_ht.at(key); - } - - /** - * @copydoc at(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - T &at(const K &key, std::size_t precalculated_hash) { - return m_ht.at(key, precalculated_hash); - } - - /** - * @copydoc at(const K& key) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const T &at(const K &key) const { - return m_ht.at(key); - } - - /** - * @copydoc at(const K& key, std::size_t precalculated_hash) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const T &at(const K &key, std::size_t precalculated_hash) const { - return m_ht.at(key, precalculated_hash); - } - - T &operator[](const Key &key) { return m_ht[key]; } - T &operator[](Key &&key) { return m_ht[std::move(key)]; } - - size_type count(const Key &key) const { return m_ht.count(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - size_type count(const Key &key, std::size_t precalculated_hash) const { - return m_ht.count(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type count(const K &key) const { - return m_ht.count(key); - } - - /** - * @copydoc count(const K& key) const - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type count(const K &key, std::size_t precalculated_hash) const { - return m_ht.count(key, precalculated_hash); - } - - iterator find(const Key &key) { return m_ht.find(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - iterator find(const Key &key, std::size_t precalculated_hash) { - return m_ht.find(key, precalculated_hash); - } - - const_iterator find(const Key &key) const { return m_ht.find(key); } - - /** - * @copydoc find(const Key& key, std::size_t precalculated_hash) - */ - const_iterator find(const Key &key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - iterator find(const K &key) { - return m_ht.find(key); - } - - /** - * @copydoc find(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - iterator find(const K &key, std::size_t precalculated_hash) { - return m_ht.find(key, precalculated_hash); - } - - /** - * @copydoc find(const K& key) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const_iterator find(const K &key) const { - return m_ht.find(key); - } - - /** - * @copydoc find(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const_iterator find(const K &key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - bool contains(const Key &key) const { return m_ht.contains(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The - * hash value should be the same as hash_function()(key). Useful to speed-up - * the lookup if you already have the hash. - */ - bool contains(const Key &key, std::size_t precalculated_hash) const { - return m_ht.contains(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * KeyEqual::is_transparent exists. If so, K must be hashable and comparable - * to Key. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - bool contains(const K &key) const { - return m_ht.contains(key); - } - - /** - * @copydoc contains(const K& key) const - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The - * hash value should be the same as hash_function()(key). Useful to speed-up - * the lookup if you already have the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - bool contains(const K &key, std::size_t precalculated_hash) const { - return m_ht.contains(key, precalculated_hash); - } - - std::pair equal_range(const Key &key) { - return m_ht.equal_range(key); - } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - std::pair equal_range(const Key &key, - std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - std::pair equal_range(const Key &key) const { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) - */ - std::pair equal_range( - const Key &key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key) { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key, - std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * @copydoc equal_range(const K& key) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key) const { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const K& key, std::size_t precalculated_hash) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range( - const K &key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /* - * Bucket interface - */ - size_type bucket_count() const { return m_ht.bucket_count(); } - size_type max_bucket_count() const { return m_ht.max_bucket_count(); } - - /* - * Hash policy - */ - float load_factor() const { return m_ht.load_factor(); } - float max_load_factor() const { return m_ht.max_load_factor(); } - void max_load_factor(float ml) { m_ht.max_load_factor(ml); } - - void rehash(size_type count) { m_ht.rehash(count); } - void reserve(size_type count) { m_ht.reserve(count); } - - /* - * Observers - */ - hasher hash_function() const { return m_ht.hash_function(); } - key_equal key_eq() const { return m_ht.key_eq(); } - - /* - * Other - */ - - /** - * Convert a `const_iterator` to an `iterator`. - */ - iterator mutable_iterator(const_iterator pos) { - return m_ht.mutable_iterator(pos); - } - - /** - * Serialize the map through the `serializer` parameter. - * - * The `serializer` parameter must be a function object that supports the - * following call: - * - `template void operator()(const U& value);` where the types - * `std::uint64_t`, `float` and `std::pair` must be supported for U. - * - * The implementation leaves binary compatibility (endianness, IEEE 754 for - * floats, ...) of the types it serializes in the hands of the `Serializer` - * function object if compatibility is required. - */ - template - void serialize(Serializer &serializer) const { - m_ht.serialize(serializer); - } - - /** - * Deserialize a previously serialized map through the `deserializer` - * parameter. - * - * The `deserializer` parameter must be a function object that supports the - * following calls: - * - `template U operator()();` where the types `std::uint64_t`, - * `float` and `std::pair` must be supported for U. - * - * If the deserialized hash map type is hash compatible with the serialized - * map, the deserialization process can be sped up by setting - * `hash_compatible` to true. To be hash compatible, the Hash, KeyEqual and - * GrowthPolicy must behave the same way than the ones used on the serialized - * map. The `std::size_t` must also be of the same size as the one on the - * platform used to serialize the map. If these criteria are not met, the - * behaviour is undefined with `hash_compatible` sets to true. - * - * The behaviour is undefined if the type `Key` and `T` of the `sparse_map` - * are not the same as the types used during serialization. - * - * The implementation leaves binary compatibility (endianness, IEEE 754 for - * floats, size of int, ...) of the types it deserializes in the hands of the - * `Deserializer` function object if compatibility is required. - */ - template - static sparse_map deserialize(Deserializer &deserializer, - bool hash_compatible = false) { - sparse_map map(0); - map.m_ht.deserialize(deserializer, hash_compatible); - - return map; - } - - friend bool operator==(const sparse_map &lhs, const sparse_map &rhs) { - if (lhs.size() != rhs.size()) { - return false; - } - - for (const auto &element_lhs : lhs) { - const auto it_element_rhs = rhs.find(element_lhs.first); - if (it_element_rhs == rhs.cend() || - element_lhs.second != it_element_rhs->second) { - return false; - } - } - - return true; - } - - friend bool operator!=(const sparse_map &lhs, const sparse_map &rhs) { - return !operator==(lhs, rhs); - } - - friend void swap(sparse_map &lhs, sparse_map &rhs) { lhs.swap(rhs); } - - private: - ht m_ht; -}; - -/** - * Same as `tsl::sparse_map`. - */ -template , - class KeyEqual = std::equal_to, - class Allocator = std::allocator>> -using sparse_pg_map = - sparse_map; - -} // end namespace tsl - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_set.h b/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_set.h deleted file mode 100644 index 3ce6a58..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/tsl/sparse_set.h +++ /dev/null @@ -1,655 +0,0 @@ -/** - * MIT License - * - * Copyright (c) 2017 Thibaut Goetghebuer-Planchon - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ -#ifndef TSL_SPARSE_SET_H -#define TSL_SPARSE_SET_H - -#include -#include -#include -#include -#include -#include - -#include "sparse_hash.h" - -namespace tsl { - -/** - * Implementation of a sparse hash set using open-addressing with quadratic - * probing. The goal on the hash set is to be the most memory efficient - * possible, even at low load factor, while keeping reasonable performances. - * - * `GrowthPolicy` defines how the set grows and consequently how a hash value is - * mapped to a bucket. By default the set uses - * `tsl::sh::power_of_two_growth_policy`. This policy keeps the number of - * buckets to a power of two and uses a mask to map the hash to a bucket instead - * of the slow modulo. Other growth policies are available and you may define - * your own growth policy, check `tsl::sh::power_of_two_growth_policy` for the - * interface. - * - * `ExceptionSafety` defines the exception guarantee provided by the class. By - * default only the basic exception safety is guaranteed which mean that all - * resources used by the hash set will be freed (no memory leaks) but the hash - * set may end-up in an undefined state if an exception is thrown (undefined - * here means that some elements may be missing). This can ONLY happen on rehash - * (either on insert or if `rehash` is called explicitly) and will occur if the - * Allocator can't allocate memory (`std::bad_alloc`) or if the copy constructor - * (when a nothrow move constructor is not available) throws an exception. This - * can be avoided by calling `reserve` beforehand. This basic guarantee is - * similar to the one of `google::sparse_hash_map` and `spp::sparse_hash_map`. - * It is possible to ask for the strong exception guarantee with - * `tsl::sh::exception_safety::strong`, the drawback is that the set will be - * slower on rehashes and will also need more memory on rehashes. - * - * `Sparsity` defines how much the hash set will compromise between insertion - * speed and memory usage. A high sparsity means less memory usage but longer - * insertion times, and vice-versa for low sparsity. The default - * `tsl::sh::sparsity::medium` sparsity offers a good compromise. It doesn't - * change the lookup speed. - * - * `Key` must be nothrow move constructible and/or copy constructible. - * - * If the destructor of `Key` throws an exception, the behaviour of the class is - * undefined. - * - * Iterators invalidation: - * - clear, operator=, reserve, rehash: always invalidate the iterators. - * - insert, emplace, emplace_hint: if there is an effective insert, invalidate - * the iterators. - * - erase: always invalidate the iterators. - */ -template , - class KeyEqual = std::equal_to, - class Allocator = std::allocator, - class GrowthPolicy = tsl::sh::power_of_two_growth_policy<2>, - tsl::sh::exception_safety ExceptionSafety = - tsl::sh::exception_safety::basic, - tsl::sh::sparsity Sparsity = tsl::sh::sparsity::medium> -class sparse_set { - private: - template - using has_is_transparent = tsl::detail_sparse_hash::has_is_transparent; - - class KeySelect { - public: - using key_type = Key; - - const key_type &operator()(const Key &key) const noexcept { return key; } - - key_type &operator()(Key &key) noexcept { return key; } - }; - - using ht = - detail_sparse_hash::sparse_hash; - - public: - using key_type = typename ht::key_type; - using value_type = typename ht::value_type; - using size_type = typename ht::size_type; - using difference_type = typename ht::difference_type; - using hasher = typename ht::hasher; - using key_equal = typename ht::key_equal; - using allocator_type = typename ht::allocator_type; - using reference = typename ht::reference; - using const_reference = typename ht::const_reference; - using pointer = typename ht::pointer; - using const_pointer = typename ht::const_pointer; - using iterator = typename ht::iterator; - using const_iterator = typename ht::const_iterator; - - /* - * Constructors - */ - sparse_set() : sparse_set(ht::DEFAULT_INIT_BUCKET_COUNT) {} - - explicit sparse_set(size_type bucket_count, const Hash &hash = Hash(), - const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR) {} - - sparse_set(size_type bucket_count, const Allocator &alloc) - : sparse_set(bucket_count, Hash(), KeyEqual(), alloc) {} - - sparse_set(size_type bucket_count, const Hash &hash, const Allocator &alloc) - : sparse_set(bucket_count, hash, KeyEqual(), alloc) {} - - explicit sparse_set(const Allocator &alloc) - : sparse_set(ht::DEFAULT_INIT_BUCKET_COUNT, alloc) {} - - template - sparse_set(InputIt first, InputIt last, - size_type bucket_count = ht::DEFAULT_INIT_BUCKET_COUNT, - const Hash &hash = Hash(), const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : sparse_set(bucket_count, hash, equal, alloc) { - insert(first, last); - } - - template - sparse_set(InputIt first, InputIt last, size_type bucket_count, - const Allocator &alloc) - : sparse_set(first, last, bucket_count, Hash(), KeyEqual(), alloc) {} - - template - sparse_set(InputIt first, InputIt last, size_type bucket_count, - const Hash &hash, const Allocator &alloc) - : sparse_set(first, last, bucket_count, hash, KeyEqual(), alloc) {} - - sparse_set(std::initializer_list init, - size_type bucket_count = ht::DEFAULT_INIT_BUCKET_COUNT, - const Hash &hash = Hash(), const KeyEqual &equal = KeyEqual(), - const Allocator &alloc = Allocator()) - : sparse_set(init.begin(), init.end(), bucket_count, hash, equal, alloc) { - } - - sparse_set(std::initializer_list init, size_type bucket_count, - const Allocator &alloc) - : sparse_set(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), - alloc) {} - - sparse_set(std::initializer_list init, size_type bucket_count, - const Hash &hash, const Allocator &alloc) - : sparse_set(init.begin(), init.end(), bucket_count, hash, KeyEqual(), - alloc) {} - - sparse_set &operator=(std::initializer_list ilist) { - m_ht.clear(); - - m_ht.reserve(ilist.size()); - m_ht.insert(ilist.begin(), ilist.end()); - - return *this; - } - - allocator_type get_allocator() const { return m_ht.get_allocator(); } - - /* - * Iterators - */ - iterator begin() noexcept { return m_ht.begin(); } - const_iterator begin() const noexcept { return m_ht.begin(); } - const_iterator cbegin() const noexcept { return m_ht.cbegin(); } - - iterator end() noexcept { return m_ht.end(); } - const_iterator end() const noexcept { return m_ht.end(); } - const_iterator cend() const noexcept { return m_ht.cend(); } - - /* - * Capacity - */ - bool empty() const noexcept { return m_ht.empty(); } - size_type size() const noexcept { return m_ht.size(); } - size_type max_size() const noexcept { return m_ht.max_size(); } - - /* - * Modifiers - */ - void clear() noexcept { m_ht.clear(); } - - std::pair insert(const value_type &value) { - return m_ht.insert(value); - } - - std::pair insert(value_type &&value) { - return m_ht.insert(std::move(value)); - } - - iterator insert(const_iterator hint, const value_type &value) { - return m_ht.insert_hint(hint, value); - } - - iterator insert(const_iterator hint, value_type &&value) { - return m_ht.insert_hint(hint, std::move(value)); - } - - template - void insert(InputIt first, InputIt last) { - m_ht.insert(first, last); - } - - void insert(std::initializer_list ilist) { - m_ht.insert(ilist.begin(), ilist.end()); - } - - /** - * Due to the way elements are stored, emplace will need to move or copy the - * key-value once. The method is equivalent to - * `insert(value_type(std::forward(args)...));`. - * - * Mainly here for compatibility with the `std::unordered_map` interface. - */ - template - std::pair emplace(Args &&...args) { - return m_ht.emplace(std::forward(args)...); - } - - /** - * Due to the way elements are stored, emplace_hint will need to move or copy - * the key-value once. The method is equivalent to `insert(hint, - * value_type(std::forward(args)...));`. - * - * Mainly here for compatibility with the `std::unordered_map` interface. - */ - template - iterator emplace_hint(const_iterator hint, Args &&...args) { - return m_ht.emplace_hint(hint, std::forward(args)...); - } - - iterator erase(iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator pos) { return m_ht.erase(pos); } - iterator erase(const_iterator first, const_iterator last) { - return m_ht.erase(first, last); - } - size_type erase(const key_type &key) { return m_ht.erase(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - size_type erase(const key_type &key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type erase(const K &key) { - return m_ht.erase(key); - } - - /** - * @copydoc erase(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type erase(const K &key, std::size_t precalculated_hash) { - return m_ht.erase(key, precalculated_hash); - } - - void swap(sparse_set &other) { other.m_ht.swap(m_ht); } - - /* - * Lookup - */ - size_type count(const Key &key) const { return m_ht.count(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - size_type count(const Key &key, std::size_t precalculated_hash) const { - return m_ht.count(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type count(const K &key) const { - return m_ht.count(key); - } - - /** - * @copydoc count(const K& key) const - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - size_type count(const K &key, std::size_t precalculated_hash) const { - return m_ht.count(key, precalculated_hash); - } - - iterator find(const Key &key) { return m_ht.find(key); } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - iterator find(const Key &key, std::size_t precalculated_hash) { - return m_ht.find(key, precalculated_hash); - } - - const_iterator find(const Key &key) const { return m_ht.find(key); } - - /** - * @copydoc find(const Key& key, std::size_t precalculated_hash) - */ - const_iterator find(const Key &key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - iterator find(const K &key) { - return m_ht.find(key); - } - - /** - * @copydoc find(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - iterator find(const K &key, std::size_t precalculated_hash) { - return m_ht.find(key, precalculated_hash); - } - - /** - * @copydoc find(const K& key) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const_iterator find(const K &key) const { - return m_ht.find(key); - } - - /** - * @copydoc find(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - const_iterator find(const K &key, std::size_t precalculated_hash) const { - return m_ht.find(key, precalculated_hash); - } - - bool contains(const Key &key) const { return m_ht.contains(key); } - - /** - * Use the hash value 'precalculated_hash' instead of hashing the key. The - * hash value should be the same as hash_function()(key). Useful to speed-up - * the lookup if you already have the hash. - */ - bool contains(const Key &key, std::size_t precalculated_hash) const { - return m_ht.contains(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * KeyEqual::is_transparent exists. If so, K must be hashable and comparable - * to Key. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - bool contains(const K &key) const { - return m_ht.contains(key); - } - - /** - * @copydoc contains(const K& key) const - * - * Use the hash value 'precalculated_hash' instead of hashing the key. The - * hash value should be the same as hash_function()(key). Useful to speed-up - * the lookup if you already have the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - bool contains(const K &key, std::size_t precalculated_hash) const { - return m_ht.contains(key, precalculated_hash); - } - - std::pair equal_range(const Key &key) { - return m_ht.equal_range(key); - } - - /** - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - std::pair equal_range(const Key &key, - std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - std::pair equal_range(const Key &key) const { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) - */ - std::pair equal_range( - const Key &key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * This overload only participates in the overload resolution if the typedef - * `KeyEqual::is_transparent` exists. If so, `K` must be hashable and - * comparable to `Key`. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key) { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const K& key) - * - * Use the hash value `precalculated_hash` instead of hashing the key. The - * hash value should be the same as `hash_function()(key)`, otherwise the - * behaviour is undefined. Useful to speed-up the lookup if you already have - * the hash. - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key, - std::size_t precalculated_hash) { - return m_ht.equal_range(key, precalculated_hash); - } - - /** - * @copydoc equal_range(const K& key) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range(const K &key) const { - return m_ht.equal_range(key); - } - - /** - * @copydoc equal_range(const K& key, std::size_t precalculated_hash) - */ - template < - class K, class KE = KeyEqual, - typename std::enable_if::value>::type * = nullptr> - std::pair equal_range( - const K &key, std::size_t precalculated_hash) const { - return m_ht.equal_range(key, precalculated_hash); - } - - /* - * Bucket interface - */ - size_type bucket_count() const { return m_ht.bucket_count(); } - size_type max_bucket_count() const { return m_ht.max_bucket_count(); } - - /* - * Hash policy - */ - float load_factor() const { return m_ht.load_factor(); } - float max_load_factor() const { return m_ht.max_load_factor(); } - void max_load_factor(float ml) { m_ht.max_load_factor(ml); } - - void rehash(size_type count) { m_ht.rehash(count); } - void reserve(size_type count) { m_ht.reserve(count); } - - /* - * Observers - */ - hasher hash_function() const { return m_ht.hash_function(); } - key_equal key_eq() const { return m_ht.key_eq(); } - - /* - * Other - */ - - /** - * Convert a `const_iterator` to an `iterator`. - */ - iterator mutable_iterator(const_iterator pos) { - return m_ht.mutable_iterator(pos); - } - - /** - * Serialize the set through the `serializer` parameter. - * - * The `serializer` parameter must be a function object that supports the - * following call: - * - `void operator()(const U& value);` where the types `std::uint64_t`, - * `float` and `Key` must be supported for U. - * - * The implementation leaves binary compatibility (endianness, IEEE 754 for - * floats, ...) of the types it serializes in the hands of the `Serializer` - * function object if compatibility is required. - */ - template - void serialize(Serializer &serializer) const { - m_ht.serialize(serializer); - } - - /** - * Deserialize a previously serialized set through the `deserializer` - * parameter. - * - * The `deserializer` parameter must be a function object that supports the - * following calls: - * - `template U operator()();` where the types `std::uint64_t`, - * `float` and `Key` must be supported for U. - * - * If the deserialized hash set type is hash compatible with the serialized - * set, the deserialization process can be sped up by setting - * `hash_compatible` to true. To be hash compatible, the Hash, KeyEqual and - * GrowthPolicy must behave the same way than the ones used on the serialized - * set. The `std::size_t` must also be of the same size as the one on the - * platform used to serialize the set. If these criteria are not met, the - * behaviour is undefined with `hash_compatible` sets to true. - * - * The behaviour is undefined if the type `Key` of the `sparse_set` is not the - * same as the type used during serialization. - * - * The implementation leaves binary compatibility (endianness, IEEE 754 for - * floats, size of int, ...) of the types it deserializes in the hands of the - * `Deserializer` function object if compatibility is required. - */ - template - static sparse_set deserialize(Deserializer &deserializer, - bool hash_compatible = false) { - sparse_set set(0); - set.m_ht.deserialize(deserializer, hash_compatible); - - return set; - } - - friend bool operator==(const sparse_set &lhs, const sparse_set &rhs) { - if (lhs.size() != rhs.size()) { - return false; - } - - for (const auto &element_lhs : lhs) { - const auto it_element_rhs = rhs.find(element_lhs); - if (it_element_rhs == rhs.cend()) { - return false; - } - } - - return true; - } - - friend bool operator!=(const sparse_set &lhs, const sparse_set &rhs) { - return !operator==(lhs, rhs); - } - - friend void swap(sparse_set &lhs, sparse_set &rhs) { lhs.swap(rhs); } - - private: - ht m_ht; -}; - -/** - * Same as `tsl::sparse_set`. - */ -template , - class KeyEqual = std::equal_to, - class Allocator = std::allocator> -using sparse_pg_set = - sparse_set; - -} // end namespace tsl - -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/types.h b/packages/leann-backend-diskann/third_party/DiskANN/include/types.h deleted file mode 100644 index 953d59a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/types.h +++ /dev/null @@ -1,22 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include -#include -#include "any_wrappers.h" - -namespace diskann -{ -typedef uint32_t location_t; - -using DataType = std::any; -using TagType = std::any; -using LabelType = std::any; -using TagVector = AnyWrapper::AnyVector; -using DataVector = AnyWrapper::AnyVector; -using Labelvector = AnyWrapper::AnyVector; -using TagRobinSet = AnyWrapper::AnyRobinSet; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/utils.h b/packages/leann-backend-diskann/third_party/DiskANN/include/utils.h deleted file mode 100644 index 355a613..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/utils.h +++ /dev/null @@ -1,1455 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include - -#include "common_includes.h" - -#ifdef __APPLE__ -#include -#else -#include -#endif - -#ifdef _WINDOWS -#include -typedef HANDLE FileHandle; -#else -#include -typedef int FileHandle; -#endif - -#include "distance.h" -#include "logger.h" -#include "cached_io.h" -#include "ann_exception.h" -#include "windows_customizations.h" -#include "tsl/robin_set.h" -#include "types.h" -#include "tag_uint128.h" -#include - -#ifdef EXEC_ENV_OLS -#include "content_buf.h" -#include "memory_mapped_files.h" -#endif - -#ifdef __APPLE__ -#ifdef __arm64__ -#define _MM_HINT_T0 1 -#define _MM_HINT_T1 2 - -static inline __attribute__((always_inline)) void _mm_prefetch(char const *p, int i) -{ - switch (i) - { - case _MM_HINT_T0: - __builtin_prefetch(p, 0, 3); - break; - case _MM_HINT_T1: - __builtin_prefetch(p, 0, 2); - break; - } -} -#endif - -#define LAPACK_COL_MAJOR 1 -#define LAPACK_ROW_MAJOR 0 -#ifdef __APPLE__ -typedef int clp_int; -#else -typedef __CLPK_integer clp_int; -#endif - -inline void _sge_trans(int matrix_layout, clp_int m, clp_int n, const float *in, clp_int ldin, float *out, - clp_int ldout) -{ - clp_int i, j, x, y; - - if (matrix_layout == LAPACK_COL_MAJOR) - { - x = n; - y = m; - } - else - { - x = m; - y = n; - } - for (i = 0; i < MIN(y, ldin); i++) - { - for (j = 0; j < MIN(x, ldout); j++) - { - out[(size_t)i * ldout + j] = in[(size_t)j * ldin + i]; - } - } -} -inline clp_int sgesdd_rm_work(char jobz, clp_int m, clp_int n, float *a, clp_int lda, float *s, float *u, clp_int ldu, - float *vt, clp_int ldvt, float *work, clp_int lwork, clp_int *iwork) -{ - clp_int info = 0; - clp_int nrows_u = ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && m < n)) ? m : 1; - clp_int ncols_u = ((jobz == 'a') || ((jobz == 'o') && m < n)) ? m : ((jobz == 's') ? MIN(m, n) : 1); - clp_int nrows_vt = ((jobz == 'a') || ((jobz == 'o') && m >= n)) ? n : ((jobz == 's') ? MIN(m, n) : 1); - - clp_int lda_t = MAX(1, m); - clp_int ldu_t = MAX(1, nrows_u); - clp_int ldvt_t = MAX(1, nrows_vt); - float *a_t = NULL; - float *u_t = NULL; - float *vt_t = NULL; - - // check leading dimensions - if (lda < n) - { - info = -6; - return info; - } - if (ldu < ncols_u) - { - info = -9; - return info; - } - if (ldvt < n) - { - info = -11; - return info; - } - - // query for optimal work size if lwork = -1 - if (lwork == -1) - { - sgesdd_(&jobz, &m, &n, a, &lda_t, s, u, &ldu_t, vt, &ldvt_t, work, &lwork, iwork, &info); - return (info < 0) ? (info - 1) : info; - } - - // setup temp arrays - a_t = (float *)malloc(sizeof(float) * lda_t * MAX(1, n)); - if (a_t == NULL) - { - info = -1011; - return info; - } - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m < n))) - { - u_t = (float *)malloc(sizeof(float) * ldu_t * MAX(1, ncols_u)); - if (u_t == NULL) - { - info = -1011; - free(a_t); - return info; - } - } - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m >= n))) - { - vt_t = (float *)malloc(sizeof(float) * ldvt_t * MAX(1, n)); - if (vt_t == NULL) - { - info = -1011; - free(a_t); - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m < n))) - { - free(u_t); - } - return info; - } - } - - _sge_trans(LAPACK_ROW_MAJOR, m, n, a, lda, a_t, lda_t); - sgesdd_(&jobz, &m, &n, a_t, &lda_t, s, u_t, &ldu_t, vt_t, &ldvt_t, work, &lwork, iwork, &info); - - if (info < 0) - { - info = info - 1; - } - /* Transpose output matrices */ - _sge_trans(LAPACK_COL_MAJOR, m, n, a_t, lda_t, a, lda); - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m < n))) - { - _sge_trans(LAPACK_COL_MAJOR, nrows_u, ncols_u, u_t, ldu_t, u, ldu); - } - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m >= n))) - { - _sge_trans(LAPACK_COL_MAJOR, nrows_vt, n, vt_t, ldvt_t, vt, ldvt); - } - /* Release memory and exit */ - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m >= n))) - { - free(vt_t); - } - if ((jobz == 'a') || (jobz == 's') || ((jobz == 'o') && (m < n))) - { - free(u_t); - } - free(a_t); - return info; -} - -inline clp_int LAPACKE_sgesdd(int matrix_layout, char jobz, clp_int m, clp_int n, float *a, clp_int lda, float *s, - float *u, clp_int ldu, float *vt, clp_int ldvt) -{ - // internal SGESDD vars - clp_int info = 0; - clp_int lwork = -1; - clp_int *iwork = NULL; - float *work = NULL; - float work_query; - - // allocate space for iwork - iwork = (clp_int *)malloc(sizeof(clp_int) * MAX(1, 8 * MIN(m, n))); - if (iwork == NULL) - throw; - /* Query optimal working array(s) size */ - info = sgesdd_rm_work(jobz, m, n, a, lda, s, u, ldu, vt, ldvt, &work_query, lwork, iwork); - if (info != 0) - { - free(iwork); - info = -1010; - return info; - } - - lwork = (clp_int)work_query; - /* Allocate memory for work arrays */ - work = (float *)malloc(sizeof(float) * lwork); - if (work == NULL) - throw; - - /* Call middle-level interface */ - info = sgesdd_rm_work(jobz, m, n, a, lda, s, u, ldu, vt, ldvt, work, lwork, iwork); - /* Release memory and exit */ - free(work); - free(iwork); - return info; -} -#endif - -// taken from -// https://github.com/Microsoft/BLAS-on-flash/blob/master/include/utils.h -// round up X to the nearest multiple of Y -#define ROUND_UP(X, Y) ((((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) * (Y)) - -#define DIV_ROUND_UP(X, Y) (((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) - -// round down X to the nearest multiple of Y -#define ROUND_DOWN(X, Y) (((uint64_t)(X) / (Y)) * (Y)) - -// alignment tests -#define IS_ALIGNED(X, Y) ((uint64_t)(X) % (uint64_t)(Y) == 0) -#define IS_512_ALIGNED(X) IS_ALIGNED(X, 512) -#define IS_4096_ALIGNED(X) IS_ALIGNED(X, 4096) -#define METADATA_SIZE \ - 4096 // all metadata of individual sub-component files is written in first - // 4KB for unified files - -#define BUFFER_SIZE_FOR_CACHED_IO (size_t)1024 * (size_t)1048576 - -#define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||" -#define PBWIDTH 60 - -inline bool file_exists_impl(const std::string &name, bool dirCheck = false) -{ - int val; -#ifndef _WINDOWS - struct stat buffer; - val = stat(name.c_str(), &buffer); -#else - // It is the 21st century but Windows API still thinks in 32-bit terms. - // Turns out calling stat() on a file > 4GB results in errno = 132 - // (OVERFLOW). How silly is this!? So calling _stat64() - struct _stat64 buffer; - val = _stat64(name.c_str(), &buffer); -#endif - - if (val != 0) - { - switch (errno) - { - case EINVAL: - diskann::cout << "Invalid argument passed to stat()" << std::endl; - break; - case ENOENT: - // file is not existing, not an issue, so we won't cout anything. - break; - default: - diskann::cout << "Unexpected error in stat():" << errno << std::endl; - break; - } - return false; - } - else - { - // the file entry exists. If reqd, check if this is a directory. - return dirCheck ? buffer.st_mode & S_IFDIR : true; - } -} - -inline bool file_exists(const std::string &name, bool dirCheck = false) -{ -#ifdef EXEC_ENV_OLS - bool exists = file_exists_impl(name, dirCheck); - if (exists) - { - return true; - } - if (!dirCheck) - { - // try with .enc extension - std::string enc_name = name + ENCRYPTED_EXTENSION; - return file_exists_impl(enc_name, dirCheck); - } - else - { - return exists; - } -#else - return file_exists_impl(name, dirCheck); -#endif -} - -inline void open_file_to_write(std::ofstream &writer, const std::string &filename) -{ - writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); - if (!file_exists(filename)) - writer.open(filename, std::ios::binary | std::ios::out); - else - writer.open(filename, std::ios::binary | std::ios::in | std::ios::out); - - if (writer.fail()) - { - char buff[1024]; -#ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); -#elif __APPLE__ - auto ret = std::to_string(strerror_r(errno, buff, 1024)); -#else - auto ret = std::string(strerror_r(errno, buff, 1024)); -#endif - auto message = std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; - diskann::cerr << message << std::endl; - throw diskann::ANNException(message, -1); - } -} - -inline size_t get_file_size(const std::string &fname) -{ - std::ifstream reader(fname, std::ios::binary | std::ios::ate); - if (!reader.fail() && reader.is_open()) - { - size_t end_pos = reader.tellg(); - reader.close(); - return end_pos; - } - else - { - diskann::cerr << "Could not open file: " << fname << std::endl; - return 0; - } -} - -inline int delete_file(const std::string &fileName) -{ - if (file_exists(fileName)) - { - auto rc = ::remove(fileName.c_str()); - if (rc != 0) - { - diskann::cerr << "Could not delete file: " << fileName - << " even though it exists. This might indicate a permissions " - "issue. " - "If you see this message, please contact the diskann team." - << std::endl; - } - return rc; - } - else - { - return 0; - } -} - -// generates formatted_label and _labels_map file. -inline void convert_labels_string_to_int(const std::string &inFileName, const std::string &outFileName, - const std::string &mapFileName, const std::string &unv_label) -{ - std::unordered_map string_int_map; - std::ofstream label_writer(outFileName); - std::ifstream label_reader(inFileName); - if (unv_label != "") - string_int_map[unv_label] = 0; // if universal label is provided map it to 0 always - std::string line, token; - while (std::getline(label_reader, line)) - { - std::istringstream new_iss(line); - std::vector lbls; - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (string_int_map.find(token) == string_int_map.end()) - { - uint32_t nextId = (uint32_t)string_int_map.size() + 1; - string_int_map[token] = nextId; // nextId can never be 0 - } - lbls.push_back(string_int_map[token]); - } - if (lbls.size() <= 0) - { - std::cout << "No label found"; - exit(-1); - } - for (size_t j = 0; j < lbls.size(); j++) - { - if (j != lbls.size() - 1) - label_writer << lbls[j] << ","; - else - label_writer << lbls[j] << std::endl; - } - } - label_writer.close(); - - std::ofstream map_writer(mapFileName); - for (auto mp : string_int_map) - { - map_writer << mp.first << "\t" << mp.second << std::endl; - } - map_writer.close(); -} - -#ifdef EXEC_ENV_OLS -class AlignedFileReader; -#endif - -namespace diskann -{ -static const size_t MAX_SIZE_OF_STREAMBUF = 2LL * 1024 * 1024 * 1024; - -inline void print_error_and_terminate(std::stringstream &error_stream) -{ - diskann::cerr << error_stream.str() << std::endl; - throw diskann::ANNException(error_stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); -} - -inline void report_memory_allocation_failure() -{ - std::stringstream stream; - stream << "Memory Allocation Failed."; - print_error_and_terminate(stream); -} - -inline void report_misalignment_of_requested_size(size_t align) -{ - std::stringstream stream; - stream << "Requested memory size is not a multiple of " << align << ". Can not be allocated."; - print_error_and_terminate(stream); -} - -inline void alloc_aligned(void **ptr, size_t size, size_t align) -{ - *ptr = nullptr; - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); -#ifdef _WINDOWS - *ptr = ::_aligned_malloc(size, align); // note the swapped arguments! -#elif __APPLE__ - int err = posix_memalign(ptr, align, size); - if (err) - { - std::cout << err << std::endl; - throw; - } -#else - *ptr = ::aligned_alloc(align, size); -#endif - if (*ptr == nullptr) - report_memory_allocation_failure(); -} - -inline void realloc_aligned(void **ptr, size_t size, size_t align) -{ - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); -#ifdef _WINDOWS - *ptr = ::_aligned_realloc(*ptr, size, align); -#else - diskann::cerr << "No aligned realloc on GCC. Must malloc and mem_align, " - "left it out for now." - << std::endl; -#endif - if (*ptr == nullptr) - report_memory_allocation_failure(); -} - -inline void check_stop(std::string arnd) -{ - int brnd; - diskann::cout << arnd << std::endl; - std::cin >> brnd; -} - -inline void aligned_free(void *ptr) -{ - // Gopal. Must have a check here if the pointer was actually allocated by - // _alloc_aligned - if (ptr == nullptr) - { - return; - } -#ifndef _WINDOWS - free(ptr); -#else - ::_aligned_free(ptr); -#endif -} - -inline void GenRandom(std::mt19937 &rng, unsigned *addr, unsigned size, unsigned N) -{ - for (unsigned i = 0; i < size; ++i) - { - addr[i] = rng() % (N - size); - } - - std::sort(addr, addr + size); - for (unsigned i = 1; i < size; ++i) - { - if (addr[i] <= addr[i - 1]) - { - addr[i] = addr[i - 1] + 1; - } - } - unsigned off = rng() % N; - for (unsigned i = 0; i < size; ++i) - { - addr[i] = (addr[i] + off) % N; - } -} - -// get_bin_metadata functions START -inline void get_bin_metadata_impl(std::basic_istream &reader, size_t &nrows, size_t &ncols, size_t offset = 0) -{ - int nrows_32, ncols_32; - reader.seekg(offset, reader.beg); - reader.read((char *)&nrows_32, sizeof(int)); - reader.read((char *)&ncols_32, sizeof(int)); - nrows = nrows_32; - ncols = ncols_32; -} - -#ifdef EXEC_ENV_OLS -inline void get_bin_metadata(MemoryMappedFiles &files, const std::string &bin_file, size_t &nrows, size_t &ncols, - size_t offset = 0) -{ - diskann::cout << "Getting metadata for file: " << bin_file << std::endl; - auto fc = files.getContent(bin_file); - // auto cb = ContentBuf((char*) fc._content, fc._size); - // std::basic_istream reader(&cb); - // get_bin_metadata_impl(reader, nrows, ncols, offset); - - int nrows_32, ncols_32; - int32_t *metadata_ptr = (int32_t *)((char *)fc._content + offset); - nrows_32 = *metadata_ptr; - ncols_32 = *(metadata_ptr + 1); - nrows = nrows_32; - ncols = ncols_32; -} -#endif - -inline void get_bin_metadata(const std::string &bin_file, size_t &nrows, size_t &ncols, size_t offset = 0) -{ - std::ifstream reader(bin_file.c_str(), std::ios::binary); - get_bin_metadata_impl(reader, nrows, ncols, offset); -} -// get_bin_metadata functions END - -#ifndef EXEC_ENV_OLS -inline size_t get_graph_num_frozen_points(const std::string &graph_file) -{ - size_t expected_file_size; - uint32_t max_observed_degree, start; - size_t file_frozen_pts; - - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); - - in.open(graph_file, std::ios::binary); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&max_observed_degree, sizeof(uint32_t)); - in.read((char *)&start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); - - return file_frozen_pts; -} -#endif - -template inline std::string getValues(T *data, size_t num) -{ - std::stringstream stream; - stream << "["; - for (size_t i = 0; i < num; i++) - { - stream << std::to_string(data[i]) << ","; - } - stream << "]" << std::endl; - - return stream.str(); -} - -// load_bin functions START -template -inline void load_bin_impl(std::basic_istream &reader, T *&data, size_t &npts, size_t &dim, size_t file_offset = 0) -{ - int npts_i32, dim_i32; - - reader.seekg(file_offset, reader.beg); - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - dim = (unsigned)dim_i32; - - std::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." << std::endl; - - data = new T[npts * dim]; - reader.read((char *)data, npts * dim * sizeof(T)); -} - -#ifdef EXEC_ENV_OLS -template -inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, T *&data, size_t &npts, size_t &dim, - size_t offset = 0) -{ - diskann::cout << "Reading bin file " << bin_file.c_str() << " at offset: " << offset << "..." << std::endl; - auto fc = files.getContent(bin_file); - - uint32_t t_npts, t_dim; - uint32_t *contentAsIntPtr = (uint32_t *)((char *)fc._content + offset); - t_npts = *(contentAsIntPtr); - t_dim = *(contentAsIntPtr + 1); - - npts = t_npts; - dim = t_dim; - - data = (T *)((char *)fc._content + offset + 2 * sizeof(uint32_t)); // No need to copy! -} - -DISKANN_DLLEXPORT void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset = 0); -template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset = 0); -template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, - size_t offset = 0); - -template -DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset = 0); - -// Unlike load_bin, assumes that data is already allocated 'size' entries -template -DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset = 0); - -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, T &value, size_t offset = 0); -#endif - -template -inline void load_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t offset = 0) -{ - diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." << std::endl; - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try - { - diskann::cout << "Opening bin file " << bin_file.c_str() << "... " << std::endl; - reader.open(bin_file, std::ios::binary | std::ios::ate); - reader.seekg(0); - load_bin_impl(reader, data, npts, dim, offset); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::cout << "done." << std::endl; -} - -inline void wait_for_keystroke() -{ - int a; - std::cout << "Press any number to continue.." << std::endl; - std::cin >> a; -} -// load_bin functions END - -inline void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - dim = (unsigned)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } -} - -inline void prune_truthset_for_range(const std::string &bin_file, float range, - std::vector> &groundtruth, size_t &npts) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - uint64_t dim = (unsigned)dim_i32; - uint32_t *ids; - float *dists; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } - float min_dist = std::numeric_limits::max(); - float max_dist = 0; - groundtruth.resize(npts); - for (uint32_t i = 0; i < npts; i++) - { - groundtruth[i].clear(); - for (uint32_t j = 0; j < dim; j++) - { - if (dists[i * dim + j] <= range) - { - groundtruth[i].emplace_back(ids[i * dim + j]); - } - min_dist = min_dist > dists[i * dim + j] ? dists[i * dim + j] : min_dist; - max_dist = max_dist < dists[i * dim + j] ? dists[i * dim + j] : max_dist; - } - // std::cout<> &groundtruth, - size_t >_num) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::flush; - size_t actual_file_size = reader.get_file_size(); - - int nptsuint32_t, totaluint32_t; - reader.read((char *)&nptsuint32_t, sizeof(int)); - reader.read((char *)&totaluint32_t, sizeof(int)); - - gt_num = (uint64_t)nptsuint32_t; - uint64_t total_res = (uint64_t)totaluint32_t; - - diskann::cout << "Metadata: #pts = " << gt_num << ", #total_results = " << total_res << "..." << std::endl; - - size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + total_res * sizeof(uint32_t); - - if (actual_file_size != expected_file_size) - { - std::stringstream stream; - stream << "Error. File size mismatch in range truthset. actual size: " << actual_file_size - << ", expected: " << expected_file_size; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - groundtruth.clear(); - groundtruth.resize(gt_num); - std::vector gt_count(gt_num); - - reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); - - std::vector gt_stats(gt_count); - std::sort(gt_stats.begin(), gt_stats.end()); - - std::cout << "GT count percentiles:" << std::endl; - for (uint32_t p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] - << std::endl; - std::cout << "percentile 100" - << ": " << gt_stats[gt_num - 1] << std::endl; - - for (uint32_t i = 0; i < gt_num; i++) - { - groundtruth[i].clear(); - groundtruth[i].resize(gt_count[i]); - if (gt_count[i] != 0) - reader.read((char *)groundtruth[i].data(), sizeof(uint32_t) * gt_count[i]); - } -} - -#ifdef EXEC_ENV_OLS -template -inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, std::unique_ptr &data, size_t &npts, - size_t &dim, size_t offset = 0) -{ - T *ptr; - load_bin(files, bin_file, ptr, npts, dim, offset); - data.reset(ptr); -} -#endif - -inline void copy_file(std::string in_file, std::string out_file) -{ - std::ifstream source(in_file, std::ios::binary); - std::ofstream dest(out_file, std::ios::binary); - - std::istreambuf_iterator begin_source(source); - std::istreambuf_iterator end_source; - std::ostreambuf_iterator begin_dest(dest); - std::copy(begin_source, end_source, begin_dest); - - source.close(); - dest.close(); -} - -DISKANN_DLLEXPORT double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, - unsigned *our_results, unsigned dim_or, unsigned recall_at); - -DISKANN_DLLEXPORT double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, - unsigned *our_results, unsigned dim_or, unsigned recall_at, - const tsl::robin_set &active_tags); - -DISKANN_DLLEXPORT double calculate_range_search_recall(unsigned num_queries, - std::vector> &groundtruth, - std::vector> &our_results); - -template -inline void load_bin(const std::string &bin_file, std::unique_ptr &data, size_t &npts, size_t &dim, - size_t offset = 0) -{ - T *ptr; - load_bin(bin_file, ptr, npts, dim, offset); - data.reset(ptr); -} - -inline void open_file_to_write(std::ofstream &writer, const std::string &filename) -{ - writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); - if (!file_exists(filename)) - writer.open(filename, std::ios::binary | std::ios::out); - else - writer.open(filename, std::ios::binary | std::ios::in | std::ios::out); - - if (writer.fail()) - { - char buff[1024]; -#ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); -#elif __APPLE__ - auto ret = std::to_string(strerror_r(errno, buff, 1024)); -#else - auto ret = std::string(strerror_r(errno, buff, 1024)); -#endif - - std::string error_message = - std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; - diskann::cerr << error_message << std::endl; - throw diskann::ANNException(error_message, -1); - } -} - -template -inline size_t save_bin(const std::string &filename, T *data, size_t npts, size_t ndims, size_t offset = 0) -{ - std::ofstream writer; - open_file_to_write(writer, filename); - - diskann::cout << "Writing bin: " << filename.c_str() << std::endl; - writer.seekp(offset, writer.beg); - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - size_t bytes_written = npts * ndims * sizeof(T) + 2 * sizeof(uint32_t); - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - diskann::cout << "bin: #pts = " << npts << ", #dims = " << ndims << ", size = " << bytes_written << "B" - << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - diskann::cout << "Finished writing bin." << std::endl; - return bytes_written; -} - -inline void print_progress(double percentage) -{ - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); -} - -// load_aligned_bin functions START - -template -inline void load_aligned_bin_impl(std::basic_istream &reader, size_t actual_file_size, T *&data, size_t &npts, - size_t &dim, size_t &rounded_dim) -{ - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - dim = (unsigned)dim_i32; - - size_t expected_actual_file_size = npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != expected_actual_file_size) - { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << actual_file_size << " while expected size is " - << expected_actual_file_size << " npts = " << npts << " dim = " << dim << " size of = " << sizeof(T) - << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - rounded_dim = ROUND_UP(dim, 8); - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << ", aligned_dim = " << rounded_dim << "... " - << std::flush; - size_t allocSize = npts * rounded_dim * sizeof(T); - diskann::cout << "allocating aligned memory of " << allocSize << " bytes... " << std::flush; - alloc_aligned(((void **)&data), allocSize, 8 * sizeof(T)); - diskann::cout << "done. Copying data to mem_aligned buffer..." << std::flush; - - for (size_t i = 0; i < npts; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - diskann::cout << " done." << std::endl; -} - -#ifdef EXEC_ENV_OLS -template -inline void load_aligned_bin(MemoryMappedFiles &files, const std::string &bin_file, T *&data, size_t &npts, size_t &dim, - size_t &rounded_dim) -{ - try - { - diskann::cout << "Opening bin file " << bin_file << " ..." << std::flush; - FileContent fc = files.getContent(bin_file); - ContentBuf buf((char *)fc._content, fc._size); - std::basic_istream reader(&buf); - - size_t actual_file_size = fc._size; - load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, rounded_dim); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } -} -#endif - -template -inline void load_aligned_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t &rounded_dim) -{ - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try - { - diskann::cout << "Reading (with alignment) bin file " << bin_file << " ..." << std::flush; - reader.open(bin_file, std::ios::binary | std::ios::ate); - - uint64_t fsize = reader.tellg(); - reader.seekg(0); - load_aligned_bin_impl(reader, fsize, data, npts, dim, rounded_dim); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template -void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t dim) -{ -#pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (uint64_t j = 0; j < dim; j++) - { - destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; - } - } -} - -// this function will take in_file of n*d dimensions and save the output as a -// floating point matrix -// with n*(d+1) dimensions. All vectors are scaled by a large value M so that -// the norms are <=1 and the final coordinate is set so that the resulting -// norm (in d+1 coordinates) is equal to 1 this is a classical transformation -// from MIPS to L2 search from "On Symmetric and Asymmetric LSHs for Inner -// Product Search" by Neyshabur and Srebro - -template float prepare_base_for_inner_products(const std::string in_file, const std::string out_file) -{ - std::cout << "Pre-processing base file by adding extra coordinate" << std::endl; - std::ifstream in_reader(in_file.c_str(), std::ios::binary); - std::ofstream out_writer(out_file.c_str(), std::ios::binary); - uint64_t npts, in_dims, out_dims; - float max_norm = 0; - - uint32_t npts32, dims32; - in_reader.read((char *)&npts32, sizeof(uint32_t)); - in_reader.read((char *)&dims32, sizeof(uint32_t)); - - npts = npts32; - in_dims = dims32; - out_dims = in_dims + 1; - uint32_t outdims32 = (uint32_t)out_dims; - - out_writer.write((char *)&npts32, sizeof(uint32_t)); - out_writer.write((char *)&outdims32, sizeof(uint32_t)); - - size_t BLOCK_SIZE = 100000; - size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; - std::unique_ptr in_block_data = std::make_unique(block_size * in_dims); - std::unique_ptr out_block_data = std::make_unique(block_size * out_dims); - - std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); - uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); - - std::vector norms(npts, 0); - - for (uint64_t b = 0; b < num_blocks; b++) - { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) - { - for (uint64_t j = 0; j < in_dims; j++) - { - norms[start_id + p] += in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; - } - max_norm = max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; - } - } - - max_norm = std::sqrt(max_norm); - - in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); - for (uint64_t b = 0; b < num_blocks; b++) - { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) - { - for (uint64_t j = 0; j < in_dims; j++) - { - out_block_data[p * out_dims + j] = in_block_data[p * in_dims + j] / max_norm; - } - float res = 1 - (norms[start_id + p] / (max_norm * max_norm)); - res = res <= 0 ? 0 : std::sqrt(res); - out_block_data[p * out_dims + out_dims - 1] = res; - } - out_writer.write((char *)out_block_data.get(), block_pts * out_dims * sizeof(float)); - } - out_writer.close(); - return max_norm; -} - -// plain saves data as npts X ndims array into filename -template void save_Tvecs(const char *filename, T *data, size_t npts, size_t ndims) -{ - std::string fname(filename); - - // create cached ofstream with 64MB cache - cached_ofstream writer(fname, 64 * 1048576); - - unsigned dims_u32 = (unsigned)ndims; - - // start writing - for (size_t i = 0; i < npts; i++) - { - // write dims in u32 - writer.write((char *)&dims_u32, sizeof(unsigned)); - - // get cur point in data - T *cur_pt = data + i * ndims; - writer.write((char *)cur_pt, ndims * sizeof(T)); - } -} -template -inline size_t save_data_in_base_dimensions(const std::string &filename, T *data, size_t npts, size_t ndims, - size_t aligned_dim, size_t offset = 0) -{ - std::ofstream writer; //(filename, std::ios::binary | std::ios::out); - open_file_to_write(writer, filename); - int npts_i32 = (int)npts, ndims_i32 = (int)ndims; - size_t bytes_written = 2 * sizeof(uint32_t) + npts * ndims * sizeof(T); - writer.seekp(offset, writer.beg); - writer.write((char *)&npts_i32, sizeof(int)); - writer.write((char *)&ndims_i32, sizeof(int)); - for (size_t i = 0; i < npts; i++) - { - writer.write((char *)(data + i * aligned_dim), ndims * sizeof(T)); - } - writer.close(); - return bytes_written; -} - -template -inline void copy_aligned_data_from_file(const char *bin_file, T *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset = 0) -{ - if (data == nullptr) - { - diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file function", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - std::ifstream reader; - reader.exceptions(std::ios::badbit | std::ios::failbit); - reader.open(bin_file, std::ios::binary); - reader.seekg(offset, reader.beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - dim = (unsigned)dim_i32; - - for (size_t i = 0; i < npts; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } -} - -// NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector(const char *vec, size_t vecsize) -{ - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T0); -} - -// NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector_l2(const char *vec, size_t vecsize) -{ - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T1); -} - -// NOTE: Implementation in utils.cpp. -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, uint64_t npts, uint64_t ndims); - -DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, const std::string &outFileName); - -inline std::string get_tag_string(std::uint64_t tag) -{ - return std::to_string(tag); -} - -inline std::string get_tag_string(const tag_uint128 &tag) -{ - std::string str = std::to_string(tag._data2) + "_" + std::to_string(tag._data1); - return str; -} - -}; // namespace diskann - -struct PivotContainer -{ - PivotContainer() = default; - - PivotContainer(size_t pivo_id, float pivo_dist) : piv_id{pivo_id}, piv_dist{pivo_dist} - { - } - - bool operator<(const PivotContainer &p) const - { - return p.piv_dist < piv_dist; - } - - bool operator>(const PivotContainer &p) const - { - return p.piv_dist > piv_dist; - } - - size_t piv_id; - float piv_dist; -}; - -inline bool validate_index_file_size(std::ifstream &in) -{ - if (!in.is_open()) - throw diskann::ANNException("Index file size check called on unopened file stream", -1, __FUNCSIG__, __FILE__, - __LINE__); - in.seekg(0, in.end); - size_t actual_file_size = in.tellg(); - in.seekg(0, in.beg); - size_t expected_file_size; - in.read((char *)&expected_file_size, sizeof(uint64_t)); - in.seekg(0, in.beg); - if (actual_file_size != expected_file_size) - { - diskann::cerr << "Index file size error. Expected size (metadata): " << expected_file_size - << ", actual file size : " << actual_file_size << "." << std::endl; - return false; - } - return true; -} - -template inline float get_norm(T *arr, const size_t dim) -{ - float sum = 0.0f; - for (uint32_t i = 0; i < dim; i++) - { - sum += arr[i] * arr[i]; - } - return sqrt(sum); -} - -// This function is valid only for float data type. -template inline void normalize(T *arr, const size_t dim) -{ - float norm = get_norm(arr, dim); - for (uint32_t i = 0; i < dim; i++) - { - arr[i] = (T)(arr[i] / norm); - } -} - -inline std::vector read_file_to_vector_of_strings(const std::string &filename, bool unique = false) -{ - std::vector result; - std::set elementSet; - if (filename != "") - { - std::ifstream file(filename); - if (file.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); - } - std::string line; - while (std::getline(file, line)) - { - if (line.empty()) - { - break; - } - if (line.find(',') != std::string::npos) - { - std::cerr << "Every query must have exactly one filter" << std::endl; - exit(-1); - } - if (!line.empty() && (line.back() == '\r' || line.back() == '\n')) - { - line.erase(line.size() - 1); - } - if (!elementSet.count(line)) - { - result.push_back(line); - } - if (unique) - { - elementSet.insert(line); - } - } - file.close(); - } - else - { - throw diskann::ANNException(std::string("Failed to open file. filename can not be blank"), -1); - } - return result; -} - -inline void clean_up_artifacts(tsl::robin_set paths_to_clean, tsl::robin_set path_suffixes) -{ - try - { - for (const auto &path : paths_to_clean) - { - for (const auto &suffix : path_suffixes) - { - std::string curr_path_to_clean(path + "_" + suffix); - if (std::remove(curr_path_to_clean.c_str()) != 0) - diskann::cout << "Warning: Unable to remove file :" << curr_path_to_clean << std::endl; - } - } - diskann::cout << "Cleaned all artifacts" << std::endl; - } - catch (const std::exception &e) - { - diskann::cout << "Warning: Unable to clean all artifacts " << e.what() << std::endl; - } -} - -template inline const char *diskann_type_to_name() = delete; -template <> inline const char *diskann_type_to_name() -{ - return "float"; -} -template <> inline const char *diskann_type_to_name() -{ - return "uint8"; -} -template <> inline const char *diskann_type_to_name() -{ - return "int8"; -} -template <> inline const char *diskann_type_to_name() -{ - return "uint16"; -} -template <> inline const char *diskann_type_to_name() -{ - return "int16"; -} -template <> inline const char *diskann_type_to_name() -{ - return "uint32"; -} -template <> inline const char *diskann_type_to_name() -{ - return "int32"; -} -template <> inline const char *diskann_type_to_name() -{ - return "uint64"; -} -template <> inline const char *diskann_type_to_name() -{ - return "int64"; -} - -#ifdef _WINDOWS -#include -#include - -extern bool AvxSupportedCPU; -extern bool Avx2SupportedCPU; - -inline size_t getMemoryUsage() -{ - PROCESS_MEMORY_COUNTERS_EX pmc; - GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS *)&pmc, sizeof(pmc)); - return pmc.PrivateUsage; -} - -inline std::string getWindowsErrorMessage(DWORD lastError) -{ - char *errorText; - FormatMessageA( - // use system message tables to retrieve error text - FORMAT_MESSAGE_FROM_SYSTEM - // allocate buffer on local heap for error text - | FORMAT_MESSAGE_ALLOCATE_BUFFER - // Important! will fail otherwise, since we're not - // (and CANNOT) pass insertion parameters - | FORMAT_MESSAGE_IGNORE_INSERTS, - NULL, // unused with FORMAT_MESSAGE_FROM_SYSTEM - lastError, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), - (LPSTR)&errorText, // output - 0, // minimum size for output buffer - NULL); // arguments - see note - - return errorText != nullptr ? std::string(errorText) : std::string(); -} - -inline void printProcessMemory(const char *message) -{ - PROCESS_MEMORY_COUNTERS counters; - HANDLE h = GetCurrentProcess(); - GetProcessMemoryInfo(h, &counters, sizeof(counters)); - diskann::cout << message - << " [Peaking Working Set size: " << counters.PeakWorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) - << "GB Working set size: " << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) - << "GB Private bytes " << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" << std::endl; -} -#else - -// need to check and change this -inline bool avx2Supported() -{ - return true; -} -inline void printProcessMemory(const char *) -{ -} - -inline size_t getMemoryUsage() -{ // for non-windows, we have not implemented this function - return 0; -} - -#endif - -extern bool AvxSupportedCPU; -extern bool Avx2SupportedCPU; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_aligned_file_reader.h b/packages/leann-backend-diskann/third_party/DiskANN/include/windows_aligned_file_reader.h deleted file mode 100644 index 0d9a317..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_aligned_file_reader.h +++ /dev/null @@ -1,57 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once -#ifdef _WINDOWS -#ifndef USE_BING_INFRA -#include -#include -#include -#include - -#include -#include -#include -#include "aligned_file_reader.h" -#include "tsl/robin_map.h" -#include "utils.h" -#include "windows_customizations.h" - -class WindowsAlignedFileReader : public AlignedFileReader -{ - private: -#ifdef UNICODE - std::wstring m_filename; -#else - std::string m_filename; -#endif - - protected: - // virtual IOContext createContext(); - - public: - DISKANN_DLLEXPORT WindowsAlignedFileReader(){}; - DISKANN_DLLEXPORT virtual ~WindowsAlignedFileReader(){}; - - // Open & close ops - // Blocking calls - DISKANN_DLLEXPORT virtual void open(const std::string &fname) override; - DISKANN_DLLEXPORT virtual void close() override; - - DISKANN_DLLEXPORT virtual void register_thread() override; - DISKANN_DLLEXPORT virtual void deregister_thread() override - { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual void deregister_all_threads() override - { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual IOContext &get_ctx() override; - - // process batch of aligned requests in parallel - // NOTE :: blocking call for the calling thread, but can thread-safe - DISKANN_DLLEXPORT virtual void read(std::vector &read_reqs, IOContext &ctx, bool async) override; -}; -#endif // USE_BING_INFRA -#endif //_WINDOWS diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_customizations.h b/packages/leann-backend-diskann/third_party/DiskANN/include/windows_customizations.h deleted file mode 100644 index e6c5846..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_customizations.h +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#ifdef _WINDOWS - -#ifdef _WINDLL -#define DISKANN_DLLEXPORT __declspec(dllexport) -#else -#define DISKANN_DLLEXPORT __declspec(dllimport) -#endif - -#else -#define DISKANN_DLLEXPORT -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_slim_lock.h b/packages/leann-backend-diskann/third_party/DiskANN/include/windows_slim_lock.h deleted file mode 100644 index 7fc09b8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/include/windows_slim_lock.h +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. -#pragma once - -#ifndef WIN32_LEAN_AND_MEAN -#define WIN32_LEAN_AND_MEAN -#endif -#include "Windows.h" - -namespace diskann -{ -// A thin C++ wrapper around Windows exclusive functionality of Windows -// SlimReaderWriterLock. -// -// The SlimReaderWriterLock is simpler/more lightweight than std::mutex -// (8 bytes vs 80 bytes), which is useful in the scenario where DiskANN has -// one lock per vector in the index. It does not support recursive locking and -// requires Windows Vista or later. -// -// Full documentation can be found at. -// https://msdn.microsoft.com/en-us/library/windows/desktop/aa904937(v=vs.85).aspx -class windows_exclusive_slim_lock -{ - public: - windows_exclusive_slim_lock() : _lock(SRWLOCK_INIT) - { - } - - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock(const windows_exclusive_slim_lock &) = delete; - windows_exclusive_slim_lock &operator=(const windows_exclusive_slim_lock &) = delete; - - void lock() - { - return AcquireSRWLockExclusive(&_lock); - } - - bool try_lock() - { - return TryAcquireSRWLockExclusive(&_lock) != FALSE; - } - - void unlock() - { - return ReleaseSRWLockExclusive(&_lock); - } - - private: - SRWLOCK _lock; -}; - -// An exclusive lock over a SlimReaderWriterLock. -class windows_exclusive_slim_lock_guard -{ - public: - windows_exclusive_slim_lock_guard(windows_exclusive_slim_lock &p_lock) : _lock(p_lock) - { - _lock.lock(); - } - - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock_guard(const windows_exclusive_slim_lock_guard &) = delete; - windows_exclusive_slim_lock_guard &operator=(const windows_exclusive_slim_lock_guard &) = delete; - - ~windows_exclusive_slim_lock_guard() - { - _lock.unlock(); - } - - private: - windows_exclusive_slim_lock &_lock; -}; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/pyproject.toml b/packages/leann-backend-diskann/third_party/DiskANN/pyproject.toml deleted file mode 100644 index 3871c71..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/pyproject.toml +++ /dev/null @@ -1,58 +0,0 @@ -[build-system] -requires = [ - "setuptools>=59.6", - "pybind11>=2.10.0", - "cmake>=3.22", - "numpy==1.25", # this is important to keep fixed. It also means anyone using something other than 1.25 won't be able to use this library - "wheel", - "ninja" -] -build-backend = "setuptools.build_meta" - -[project] -name = "diskannpy" -version = "0.7.1" - -description = "DiskANN Python extension module" -readme = "python/README.md" -requires-python = ">=3.9" -license = {text = "MIT License"} -dependencies = [ - "numpy==1.25" -] -authors = [ - {name = "Harsha Vardhan Simhadri", email = "harshasi@microsoft.com"}, - {name = "Dax Pryce", email = "daxpryce@microsoft.com"} -] - -[project.optional-dependencies] -dev = ["black", "isort", "mypy"] - -[tool.setuptools] -package-dir = {"" = "python/src"} - -[tool.isort] -profile = "black" -multi_line_output = 3 - -[tool.mypy] -plugins = "numpy.typing.mypy_plugin" - -[tool.cibuildwheel] -manylinux-x86_64-image = "manylinux_2_28" -test-requires = ["scikit-learn~=1.2"] -build-frontend = "build" -skip = ["pp*", "*-win32", "*-manylinux_i686", "*-musllinux*"] -test-command = "python -m unittest discover {project}/python/tests" - -[tool.cibuildwheel.linux] -before-build = [ - "rpm --import https://repo.almalinux.org/almalinux/RPM-GPG-KEY-AlmaLinux", - "dnf makecache --refresh", - "dnf upgrade -y almalinux-release", - "dnf install -y epel-release", - "dnf config-manager -y --add-repo https://yum.repos.intel.com/mkl/setup/intel-mkl.repo", - "rpm --import https://yum.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB", - "dnf makecache --refresh -y", - "dnf install -y wget make cmake gcc-c++ libaio-devel gperftools-libs libunwind-devel clang-tools-extra boost-devel boost-program-options intel-mkl-2020.4-912" -] diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/python/CMakeLists.txt deleted file mode 100644 index 66a5ba3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/CMakeLists.txt +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -cmake_minimum_required(VERSION 3.18...3.22) - -set(CMAKE_CXX_STANDARD 17) - -if (PYTHON_EXECUTABLE) - set(Python3_EXECUTABLE ${PYTHON_EXECUTABLE}) -endif() - -find_package(Python3 COMPONENTS Interpreter Development.Module NumPy REQUIRED) - -execute_process(COMMAND ${Python3_EXECUTABLE} -c "import pybind11; print(pybind11.get_cmake_dir())" - OUTPUT_VARIABLE _tmp_dir - OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ECHO STDOUT) -list(APPEND CMAKE_PREFIX_PATH "${_tmp_dir}") - -# Now we can find pybind11 -find_package(pybind11 CONFIG REQUIRED) - -execute_process(COMMAND ${Python3_EXECUTABLE} -c "import numpy; print(numpy.get_include())" - OUTPUT_VARIABLE _numpy_include - OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ECHO STDOUT) - -# pybind11_add_module(diskannpy MODULE src/diskann_bindings.cpp) -# the following is fairly synonymous with pybind11_add_module, but we need more target_link_libraries -# see https://pybind11.readthedocs.io/en/latest/compiling.html#advanced-interface-library-targets for more details -add_library(_diskannpy MODULE - src/module.cpp - src/builder.cpp - src/dynamic_memory_index.cpp - src/static_memory_index.cpp - src/static_disk_index.cpp -) - -target_include_directories(_diskannpy AFTER PRIVATE include) - -if (MSVC) - target_compile_options(_diskannpy PRIVATE /U_WINDLL) -endif() - -target_link_libraries( - _diskannpy - PRIVATE - pybind11::module - pybind11::lto - pybind11::windows_extras - ${PROJECT_NAME} - ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} - ${DISKANN_ASYNC_LIB} -) - -pybind11_extension(_diskannpy) -if(NOT MSVC AND NOT ${CMAKE_BUILD_TYPE} MATCHES Debug|RelWithDebInfo) - # Strip unnecessary sections of the binary on Linux/macOS - pybind11_strip(_diskannpy) -endif() - -set_target_properties(_diskannpy PROPERTIES CXX_VISIBILITY_PRESET "hidden" - CUDA_VISIBILITY_PRESET "hidden") - -# generally, the VERSION_INFO flag is set by pyproject.toml, by way of setup.py. -# attempts to locate the version within CMake fail because the version has to be available -# to pyproject.toml for the sdist to work after we build it. - -if(NOT VERSION_INFO) - set(VERSION_INFO "0.0.0dev") -endif() -target_compile_definitions(_diskannpy PRIVATE VERSION_INFO="${VERSION_INFO}") - -# Add a post-build command to automatically copy the compiled Python module -if(UPDATE_EDITABLE_INSTALL) -add_custom_command( -TARGET _diskannpy -POST_BUILD -COMMAND ${CMAKE_COMMAND} -E copy - ${CMAKE_CURRENT_BINARY_DIR}/_diskannpy.cpython-*.so - ${CMAKE_SOURCE_DIR}/python/src/ -COMMENT "Copying Python module to python/src directory" -) -endif() \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/README.md b/packages/leann-backend-diskann/third_party/DiskANN/python/README.md deleted file mode 100644 index a0c9475..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/README.md +++ /dev/null @@ -1,55 +0,0 @@ -# diskannpy - -[![DiskANN Paper](https://img.shields.io/badge/Paper-NeurIPS%3A_DiskANN-blue)](https://papers.nips.cc/paper/9527-rand-nsg-fast-accurate-billion-point-nearest-neighbor-search-on-a-single-node.pdf) -[![DiskANN Paper](https://img.shields.io/badge/Paper-Arxiv%3A_Fresh--DiskANN-blue)](https://arxiv.org/abs/2105.09613) -[![DiskANN Paper](https://img.shields.io/badge/Paper-Filtered--DiskANN-blue)](https://harsha-simhadri.org/pubs/Filtered-DiskANN23.pdf) -[![DiskANN Main](https://github.com/microsoft/DiskANN/actions/workflows/push-test.yml/badge.svg?branch=main)](https://github.com/microsoft/DiskANN/actions/workflows/push-test.yml) -[![PyPI version](https://img.shields.io/pypi/v/diskannpy.svg)](https://pypi.org/project/diskannpy/) -[![Downloads shield](https://pepy.tech/badge/diskannpy)](https://pepy.tech/project/diskannpy) -[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) - -## Installation -Packages published to PyPI will always be built using the latest numpy major.minor release (at this time, 1.25). - -Conda distributions for versions 1.19-1.25 will be completed as a future effort. In the meantime, feel free to -clone this repository and build it yourself. - -## Local Build Instructions -Please see the [Project README](https://github.com/microsoft/DiskANN/blob/main/README.md) for system dependencies and requirements. - -After ensuring you've followed the directions to build the project library and executables, you will be ready to also -build `diskannpy` with these additional instructions. - -### Changing Numpy Version -In the root folder of DiskANN, there is a file `pyproject.toml`. You will need to edit the version of numpy in both the -`[build-system.requires]` section, as well as the `[project.dependencies]` section. The version numbers must match. - -#### Linux -```bash -python3.11 -m venv venv # versions from python3.9 and up should work -source venv/bin/activate -pip install build -python -m build -``` - -#### Windows -```powershell -py -3.11 -m venv venv # versions from python3.9 and up should work -venv\Scripts\Activate.ps1 -pip install build -python -m build -``` - -The built wheel will be placed in the `dist` directory in your DiskANN root. Install it using `pip install dist/.whl` - -## Citations -Please cite this software in your work as: -``` -@misc{diskann-github, - author = {Simhadri, Harsha Vardhan and Krishnaswamy, Ravishankar and Srinivasa, Gopal and Subramanya, Suhas Jayaram and Antonijevic, Andrija and Pryce, Dax and Kaczynski, David and Williams, Shane and Gollapudi, Siddarth and Sivashankar, Varun and Karia, Neel and Singh, Aditi and Jaiswal, Shikhar and Mahapatro, Neelam and Adams, Philip and Tower, Bryan and Patel, Yash}}, - title = {{DiskANN: Graph-structured Indices for Scalable, Fast, Fresh and Filtered Approximate Nearest Neighbor Search}}, - url = {https://github.com/Microsoft/DiskANN}, - version = {0.6.1}, - year = {2023} -} -``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cli/__main__.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cli/__main__.py deleted file mode 100644 index d2c9990..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cli/__main__.py +++ /dev/null @@ -1,152 +0,0 @@ -import diskannpy as dap -import numpy as np -import numpy.typing as npt - -import fire - -from contextlib import contextmanager -from time import perf_counter - -from typing import Tuple - - -def _basic_setup( - dtype: str, - query_vectors_file: str -) -> Tuple[dap.VectorDType, npt.NDArray[dap.VectorDType]]: - _dtype = dap.valid_dtype(dtype) - vectors_to_query = dap.vectors_from_binary(query_vectors_file, dtype=_dtype) - return _dtype, vectors_to_query - - -def dynamic( - dtype: str, - index_vectors_file: str, - query_vectors_file: str, - build_complexity: int, - graph_degree: int, - K: int, - search_complexity: int, - num_insert_threads: int, - num_search_threads: int, - gt_file: str = "", -): - _dtype, vectors_to_query = _basic_setup(dtype, query_vectors_file) - vectors_to_index = dap.vectors_from_binary(index_vectors_file, dtype=_dtype) - - npts, ndims = vectors_to_index.shape - index = dap.DynamicMemoryIndex( - "l2", _dtype, ndims, npts, build_complexity, graph_degree - ) - - tags = np.arange(1, npts+1, dtype=np.uintc) - timer = Timer() - - with timer.time("batch insert"): - index.batch_insert(vectors_to_index, tags, num_insert_threads) - - delete_tags = np.random.choice( - np.array(range(1, npts + 1, 1), dtype=np.uintc), - size=int(0.5 * npts), - replace=False - ) - with timer.time("mark deletion"): - for tag in delete_tags: - index.mark_deleted(tag) - - with timer.time("consolidation"): - index.consolidate_delete() - - deleted_data = vectors_to_index[delete_tags - 1, :] - - with timer.time("re-insertion"): - index.batch_insert(deleted_data, delete_tags, num_insert_threads) - - with timer.time("batch searched"): - tags, dists = index.batch_search(vectors_to_query, K, search_complexity, num_search_threads) - - # res_ids = tags - 1 - # if gt_file != "": - # recall = utils.calculate_recall_from_gt_file(K, res_ids, gt_file) - # print(f"recall@{K} is {recall}") - -def static( - dtype: str, - index_directory: str, - index_vectors_file: str, - query_vectors_file: str, - build_complexity: int, - graph_degree: int, - K: int, - search_complexity: int, - num_threads: int, - gt_file: str = "", - index_prefix: str = "ann" -): - _dtype, vectors_to_query = _basic_setup(dtype, query_vectors_file) - timer = Timer() - with timer.time("build static index"): - # build index - dap.build_memory_index( - data=index_vectors_file, - metric="l2", - vector_dtype=_dtype, - index_directory=index_directory, - complexity=build_complexity, - graph_degree=graph_degree, - num_threads=num_threads, - index_prefix=index_prefix, - alpha=1.2, - use_pq_build=False, - num_pq_bytes=8, - use_opq=False, - ) - - with timer.time("load static index"): - # ready search object - index = dap.StaticMemoryIndex( - metric="l2", - vector_dtype=_dtype, - data_path=index_vectors_file, - index_directory=index_directory, - num_threads=num_threads, # this can be different at search time if you would like - initial_search_complexity=search_complexity, - index_prefix=index_prefix - ) - - ids, dists = index.batch_search(vectors_to_query, K, search_complexity, num_threads) - - # if gt_file != "": - # recall = utils.calculate_recall_from_gt_file(K, ids, gt_file) - # print(f"recall@{K} is {recall}") - -def dynamic_clustered(): - pass - -def generate_clusters(): - pass - - -class Timer: - def __init__(self): - self._start = -1 - - @contextmanager - def time(self, message: str): - start = perf_counter() - if self._start == -1: - self._start = start - yield - now = perf_counter() - print(f"Operation {message} completed in {(now - start):.3f}s, total: {(now - self._start):.3f}s") - - - - -if __name__ == "__main__": - fire.Fire({ - "in-mem-dynamic": dynamic, - "in-mem-static": static, - "in-mem-dynamic-clustered": dynamic_clustered, - "generate-clusters": generate_clusters - }, name="cli") diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cluster.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cluster.py deleted file mode 100644 index 27a34bb..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/cluster.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import argparse -import utils - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="cluster", description="kmeans cluster points in a file" - ) - - parser.add_argument("-d", "--data_type", required=True) - parser.add_argument("-i", "--indexdata_file", required=True) - parser.add_argument("-k", "--num_clusters", type=int, required=True) - args = parser.parse_args() - - npts, ndims = get_bin_metadata(indexdata_file) - - data = utils.bin_to_numpy(args.data_type, args.indexdata_file) - - offsets, permutation = utils.cluster_and_permute( - args.data_type, npts, ndims, data, args.num_clusters - ) - - permuted_data = data[permutation] - - utils.numpy_to_bin(permuted_data, args.indexdata_file + ".cluster") diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-dynamic.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-dynamic.py deleted file mode 100644 index f97e131..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-dynamic.py +++ /dev/null @@ -1,161 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import argparse - -import diskannpy -import numpy as np -import utils - -def insert_and_search( - dtype_str, - indexdata_file, - querydata_file, - Lb, - graph_degree, - K, - Ls, - num_insert_threads, - num_search_threads, - gt_file, -) -> dict[str, float]: - """ - - :param dtype_str: - :param indexdata_file: - :param querydata_file: - :param Lb: - :param graph_degree: - :param K: - :param Ls: - :param num_insert_threads: - :param num_search_threads: - :param gt_file: - :return: Dictionary of timings. Key is the event and value is the number of seconds the event took - """ - timer_results: dict[str, float] = {} - - method_timer: utils.Timer = utils.Timer() - - npts, ndims = utils.get_bin_metadata(indexdata_file) - - if dtype_str == "float": - dtype = np.float32 - elif dtype_str == "int8": - dtype = np.int8 - elif dtype_str == "uint8": - dtype = np.uint8 - else: - raise ValueError("data_type must be float, int8 or uint8") - - index = diskannpy.DynamicMemoryIndex( - distance_metric="l2", - vector_dtype=dtype, - dimensions=ndims, - max_vectors=npts, - complexity=Lb, - graph_degree=graph_degree - ) - queries = diskannpy.vectors_from_file(querydata_file, dtype) - data = diskannpy.vectors_from_file(indexdata_file, dtype) - - tags = np.zeros(npts, dtype=np.uintc) - timer = utils.Timer() - for i in range(npts): - tags[i] = i + 1 - index.batch_insert(data, tags, num_insert_threads) - compute_seconds = timer.elapsed() - print('batch_insert complete in', compute_seconds, 's') - timer_results["batch_insert_seconds"] = compute_seconds - - delete_tags = np.random.choice( - np.array(range(1, npts + 1, 1), dtype=np.uintc), - size=int(0.5 * npts), - replace=False - ) - - timer.reset() - for tag in delete_tags: - index.mark_deleted(tag) - compute_seconds = timer.elapsed() - timer_results['mark_deletion_seconds'] = compute_seconds - print('mark deletion completed in', compute_seconds, 's') - - timer.reset() - index.consolidate_delete() - compute_seconds = timer.elapsed() - print('consolidation completed in', compute_seconds, 's') - timer_results['consolidation_completed_seconds'] = compute_seconds - - deleted_data = data[delete_tags - 1, :] - - timer.reset() - index.batch_insert(deleted_data, delete_tags, num_insert_threads) - compute_seconds = timer.elapsed() - print('re-insertion completed in', compute_seconds, 's') - timer_results['re-insertion_seconds'] = compute_seconds - - timer.reset() - tags, dists = index.batch_search(queries, K, Ls, num_search_threads) - compute_seconds = timer.elapsed() - print('Batch searched', queries.shape[0], ' queries in ', compute_seconds, 's') - timer_results['batch_searched_seconds'] = compute_seconds - - res_ids = tags - 1 - if gt_file != "": - timer.reset() - recall = utils.calculate_recall_from_gt_file(K, res_ids, gt_file) - print(f"recall@{K} is {recall}") - timer_results['recall_computed_seconds'] = timer.elapsed() - - timer_results['total_time_seconds'] = method_timer.elapsed() - - return timer_results - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="in-mem-dynamic", - description="Inserts points dynamically in a clustered order and search from vectors in a file.", - ) - - parser.add_argument("-d", "--data_type", required=True) - parser.add_argument("-i", "--indexdata_file", required=True) - parser.add_argument("-q", "--querydata_file", required=True) - parser.add_argument("-Lb", "--Lbuild", default=50, type=int) - parser.add_argument("-Ls", "--Lsearch", default=50, type=int) - parser.add_argument("-R", "--graph_degree", default=32, type=int) - parser.add_argument("-TI", "--num_insert_threads", default=8, type=int) - parser.add_argument("-TS", "--num_search_threads", default=8, type=int) - parser.add_argument("-K", default=10, type=int) - parser.add_argument("--gt_file", default="") - parser.add_argument("--json_timings_output", required=False, default=None, help="File to write out timings to as JSON. If not specified, timings will not be written out.") - args = parser.parse_args() - - timings = insert_and_search( - args.data_type, - args.indexdata_file, - args.querydata_file, - args.Lbuild, - args.graph_degree, # Build args - args.K, - args.Lsearch, - args.num_insert_threads, - args.num_search_threads, # search args - args.gt_file, - ) - - if args.json_timings_output is not None: - import json - timings['log_file'] = args.json_timings_output - with open(args.json_timings_output, "w") as f: - json.dump(timings, f) - -""" -An ingest optimized example with SIFT1M -source venv/bin/activate -python python/apps/in-mem-dynamic.py -d float \ --i "$HOME/data/sift/sift_base.fbin" -q "$HOME/data/sift/sift_query.fbin" --gt_file "$HOME/data/sift/gt100_base" \ --Lb 10 -R 30 -Ls 200 -""" - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-static.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-static.py deleted file mode 100644 index 9fb9a2c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/in-mem-static.py +++ /dev/null @@ -1,149 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import argparse -from xml.dom.pulldom import default_bufsize - -import diskannpy -import numpy as np -import utils - -def build_and_search( - metric, - dtype_str, - index_directory, - indexdata_file, - querydata_file, - Lb, - graph_degree, - K, - Ls, - num_threads, - gt_file, - index_prefix, - search_only -) -> dict[str, float]: - """ - - :param metric: - :param dtype_str: - :param index_directory: - :param indexdata_file: - :param querydata_file: - :param Lb: - :param graph_degree: - :param K: - :param Ls: - :param num_threads: - :param gt_file: - :param index_prefix: - :param search_only: - :return: Dictionary of timings. Key is the event and value is the number of seconds the event took - in wall-clock-time. - """ - timer_results: dict[str, float] = {} - - method_timer: utils.Timer = utils.Timer() - - if dtype_str == "float": - dtype = np.single - elif dtype_str == "int8": - dtype = np.byte - elif dtype_str == "uint8": - dtype = np.ubyte - else: - raise ValueError("data_type must be float, int8 or uint8") - - # build index - if not search_only: - build_index_timer = utils.Timer() - diskannpy.build_memory_index( - data=indexdata_file, - distance_metric=metric, - vector_dtype=dtype, - index_directory=index_directory, - complexity=Lb, - graph_degree=graph_degree, - num_threads=num_threads, - index_prefix=index_prefix, - alpha=1.2, - use_pq_build=False, - num_pq_bytes=8, - use_opq=False, - ) - timer_results["build_index_seconds"] = build_index_timer.elapsed() - - # ready search object - load_index_timer = utils.Timer() - index = diskannpy.StaticMemoryIndex( - distance_metric=metric, - vector_dtype=dtype, - index_directory=index_directory, - num_threads=num_threads, # this can be different at search time if you would like - initial_search_complexity=Ls, - index_prefix=index_prefix - ) - timer_results["load_index_seconds"] = load_index_timer.elapsed() - - queries = utils.bin_to_numpy(dtype, querydata_file) - - query_timer = utils.Timer() - ids, dists = index.batch_search(queries, 10, Ls, num_threads) - query_time = query_timer.elapsed() - qps = round(queries.shape[0]/query_time, 1) - print('Batch searched', queries.shape[0], 'in', query_time, 's @', qps, 'QPS') - timer_results["query_seconds"] = query_time - - if gt_file != "": - recall_timer = utils.Timer() - recall = utils.calculate_recall_from_gt_file(K, ids, gt_file) - print(f"recall@{K} is {recall}") - timer_results["recall_seconds"] = recall_timer.elapsed() - - timer_results['total_time_seconds'] = method_timer.elapsed() - - return timer_results - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="in-mem-static", - description="Static in-memory build and search from vectors in a file", - ) - - parser.add_argument("-m", "--metric", required=False, default="l2") - parser.add_argument("-d", "--data_type", required=True) - parser.add_argument("-id", "--index_directory", required=False, default=".") - parser.add_argument("-i", "--indexdata_file", required=True) - parser.add_argument("-q", "--querydata_file", required=True) - parser.add_argument("-Lb", "--Lbuild", default=50, type=int) - parser.add_argument("-Ls", "--Lsearch", default=50, type=int) - parser.add_argument("-R", "--graph_degree", default=32, type=int) - parser.add_argument("-T", "--num_threads", default=8, type=int) - parser.add_argument("-K", default=10, type=int) - parser.add_argument("-G", "--gt_file", default="") - parser.add_argument("-ip", "--index_prefix", required=False, default="ann") - parser.add_argument("--search_only", required=False, default=False) - parser.add_argument("--json_timings_output", required=False, default=None, help="File to write out timings to as JSON. If not specified, timings will not be written out.") - args = parser.parse_args() - - timings: dict[str, float] = build_and_search( - args.metric, - args.data_type, - args.index_directory.strip(), - args.indexdata_file.strip(), - args.querydata_file.strip(), - args.Lbuild, - args.graph_degree, # Build args - args.K, - args.Lsearch, - args.num_threads, # search args - args.gt_file, - args.index_prefix, - args.search_only - ) - - if args.json_timings_output is not None: - import json - timings['log_file'] = args.json_timings_output - with open(args.json_timings_output, "w") as f: - json.dump(timings, f) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/insert-in-clustered-order.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/insert-in-clustered-order.py deleted file mode 100644 index 25cb9d5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/insert-in-clustered-order.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import argparse - -import diskannpy -import numpy as np -import utils - - -def insert_and_search( - dtype_str, - indexdata_file, - querydata_file, - Lb, - graph_degree, - num_clusters, - num_insert_threads, - K, - Ls, - num_search_threads, - gt_file, -): - npts, ndims = utils.get_bin_metadata(indexdata_file) - - if dtype_str == "float": - dtype = np.float32 - elif dtype_str == "int8": - dtype = np.int8 - elif dtype_str == "uint8": - dtype = np.uint8 - else: - raise ValueError("data_type must be float, int8 or uint8") - - index = diskannpy.DynamicMemoryIndex( - distance_metric="l2", - vector_dtype=dtype, - dimensions=ndims, - max_vectors=npts, - complexity=Lb, - graph_degree=graph_degree - ) - queries = diskannpy.vectors_from_file(querydata_file, dtype) - data = diskannpy.vectors_from_file(indexdata_file, dtype) - - offsets, permutation = utils.cluster_and_permute( - dtype_str, npts, ndims, data, num_clusters - ) - - i = 0 - timer = utils.Timer() - for c in range(num_clusters): - cluster_index_range = range(offsets[c], offsets[c + 1]) - cluster_indices = np.array(permutation[cluster_index_range], dtype=np.uint32) - cluster_data = data[cluster_indices, :] - index.batch_insert(cluster_data, cluster_indices + 1, num_insert_threads) - print('Inserted cluster', c, 'in', timer.elapsed(), 's') - tags, dists = index.batch_search(queries, K, Ls, num_search_threads) - print('Batch searched', queries.shape[0], 'queries in', timer.elapsed(), 's') - res_ids = tags - 1 - - if gt_file != "": - recall = utils.calculate_recall_from_gt_file(K, res_ids, gt_file) - print(f"recall@{K} is {recall}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="in-mem-dynamic", - description="Inserts points dynamically in a clustered order and search from vectors in a file.", - ) - - parser.add_argument("-d", "--data_type", required=True) - parser.add_argument("-i", "--indexdata_file", required=True) - parser.add_argument("-q", "--querydata_file", required=True) - parser.add_argument("-Lb", "--Lbuild", default=50, type=int) - parser.add_argument("-Ls", "--Lsearch", default=50, type=int) - parser.add_argument("-R", "--graph_degree", default=32, type=int) - parser.add_argument("-TI", "--num_insert_threads", default=8, type=int) - parser.add_argument("-TS", "--num_search_threads", default=8, type=int) - parser.add_argument("-C", "--num_clusters", default=32, type=int) - parser.add_argument("-K", default=10, type=int) - parser.add_argument("--gt_file", default="") - args = parser.parse_args() - - insert_and_search( - args.data_type, - args.indexdata_file, - args.querydata_file, - args.Lbuild, - args.graph_degree, # Build args - args.num_clusters, - args.num_insert_threads, - args.K, - args.Lsearch, - args.num_search_threads, # search args - args.gt_file, - ) - -# An ingest optimized example with SIFT1M -# python3 ~/DiskANN/python/apps/insert-in-clustered-order.py -d float \ -# -i sift_base.fbin -q sift_query.fbin --gt_file gt100_base \ -# -Lb 10 -R 30 -Ls 200 -C 32 \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/utils.py b/packages/leann-backend-diskann/third_party/DiskANN/python/apps/utils.py deleted file mode 100644 index a526984..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/apps/utils.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import numpy as np -from scipy.cluster.vq import vq, kmeans2 -from typing import Tuple -from time import perf_counter - - -def get_bin_metadata(bin_file) -> Tuple[int, int]: - array = np.fromfile(file=bin_file, dtype=np.uint32, count=2) - return array[0], array[1] - - -def bin_to_numpy(dtype, bin_file) -> np.ndarray: - npts, ndims = get_bin_metadata(bin_file) - return np.fromfile(file=bin_file, dtype=dtype, offset=8).reshape(npts, ndims) - - -class Timer: - last = perf_counter() - - def reset(self): - new = perf_counter() - self.last = new - - def elapsed(self, round_digit:int = 3): - new = perf_counter() - elapsed_time = new - self.last - self.last = new - return round(elapsed_time, round_digit) - - -def numpy_to_bin(array, out_file): - shape = np.shape(array) - npts = shape[0].astype(np.uint32) - ndims = shape[1].astype(np.uint32) - f = open(out_file, "wb") - f.write(npts.tobytes()) - f.write(ndims.tobytes()) - f.write(array.tobytes()) - f.close() - - -def read_gt_file(gt_file) -> Tuple[np.ndarray[int], np.ndarray[float]]: - """ - Return ids and distances to queries - """ - nq, K = get_bin_metadata(gt_file) - ids = np.fromfile(file=gt_file, dtype=np.uint32, offset=8, count=nq * K).reshape( - nq, K - ) - dists = np.fromfile( - file=gt_file, dtype=np.float32, offset=8 + nq * K * 4, count=nq * K - ).reshape(nq, K) - return ids, dists - - -def calculate_recall( - result_set_indices: np.ndarray[int], - truth_set_indices: np.ndarray[int], - recall_at: int = 5, -) -> float: - """ - result_set_indices and truth_set_indices correspond by row index. the columns in each row contain the indices of - the nearest neighbors, with result_set_indices being the approximate nearest neighbor results and truth_set_indices - being the brute force nearest neighbor calculation via sklearn's NearestNeighbor class. - :param result_set_indices: - :param truth_set_indices: - :param recall_at: - :return: - """ - found = 0 - for i in range(0, result_set_indices.shape[0]): - result_set_set = set(result_set_indices[i][0:recall_at]) - truth_set_set = set(truth_set_indices[i][0:recall_at]) - found += len(result_set_set.intersection(truth_set_set)) - return found / (result_set_indices.shape[0] * recall_at) - - -def calculate_recall_from_gt_file(K: int, ids: np.ndarray[int], gt_file: str) -> float: - """ - Calculate recall from ids returned from search and those read from file - """ - gt_ids, gt_dists = read_gt_file(gt_file) - return calculate_recall(ids, gt_ids, K) - - -def cluster_and_permute( - dtype_str, npts, ndims, data, num_clusters -) -> Tuple[np.ndarray[int], np.ndarray[int]]: - """ - Cluster the data and return permutation of row indices - that would group indices of the same cluster together - """ - sample_size = min(100000, npts) - sample_indices = np.random.choice(range(npts), size=sample_size, replace=False) - sampled_data = data[sample_indices, :] - centroids, sample_labels = kmeans2(sampled_data, num_clusters, minit="++", iter=10) - labels, dist = vq(data, centroids) - - count = np.zeros(num_clusters) - for i in range(npts): - count[labels[i]] += 1 - print("Cluster counts") - print(count) - - offsets = np.zeros(num_clusters + 1, dtype=int) - for i in range(0, num_clusters, 1): - offsets[i + 1] = offsets[i] + count[i] - - permutation = np.zeros(npts, dtype=int) - counters = np.zeros(num_clusters, dtype=int) - for i in range(npts): - label = labels[i] - row = offsets[label] + counters[label] - counters[label] += 1 - permutation[row] = i - - return offsets, permutation diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/include/builder.h b/packages/leann-backend-diskann/third_party/DiskANN/python/include/builder.h deleted file mode 100644 index 56677ac..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/include/builder.h +++ /dev/null @@ -1,27 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include "common.h" -#include "distance.h" - -namespace diskannpy -{ -template -void build_disk_index(diskann::Metric metric, const std::string &data_file_path, const std::string &index_prefix_path, - uint32_t complexity, uint32_t graph_degree, double final_index_ram_limit, - double indexing_ram_budget, uint32_t num_threads, uint32_t pq_disk_bytes, - const std::string &codebook_prefix); - -template -void build_memory_index(diskann::Metric metric, const std::string &vector_bin_path, - const std::string &index_output_path, uint32_t graph_degree, uint32_t complexity, float alpha, - uint32_t num_threads, bool use_pq_build, size_t num_pq_bytes, bool use_opq, - bool use_tags = false, const std::string &filter_labels_file = "", - const std::string &universal_label = "", uint32_t filter_complexity = 0); - -} // namespace diskannpy diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/include/common.h b/packages/leann-backend-diskann/third_party/DiskANN/python/include/common.h deleted file mode 100644 index 7c63534..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/include/common.h +++ /dev/null @@ -1,24 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include -#include - -namespace py = pybind11; - -namespace diskannpy -{ - -typedef uint32_t filterT; - -typedef uint32_t StaticIdType; -typedef uint32_t DynamicIdType; - -template using NeighborsAndDistances = std::pair, py::array_t>; - -}; // namespace diskannpy diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/include/dynamic_memory_index.h b/packages/leann-backend-diskann/third_party/DiskANN/python/include/dynamic_memory_index.h deleted file mode 100644 index 02d6b8c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/include/dynamic_memory_index.h +++ /dev/null @@ -1,53 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include -#include - -#include "common.h" -#include "index.h" -#include "parameters.h" - -namespace py = pybind11; - -namespace diskannpy -{ - -template -class DynamicMemoryIndex -{ - public: - DynamicMemoryIndex(diskann::Metric m, size_t dimensions, size_t max_vectors, uint32_t complexity, - uint32_t graph_degree, bool saturate_graph, uint32_t max_occlusion_size, float alpha, - uint32_t num_threads, uint32_t filter_complexity, uint32_t num_frozen_points, - uint32_t initial_search_complexity, uint32_t initial_search_threads, - bool concurrent_consolidation); - - void load(const std::string &index_path); - int insert(const py::array_t &vector, DynamicIdType id); - py::array_t batch_insert(py::array_t &vectors, - py::array_t &ids, int32_t num_inserts, - int num_threads = 0); - int mark_deleted(DynamicIdType id); - void save(const std::string &save_path, bool compact_before_save = false); - NeighborsAndDistances search(py::array_t &query, uint64_t knn, - uint64_t complexity); - NeighborsAndDistances batch_search(py::array_t &queries, - uint64_t num_queries, uint64_t knn, uint64_t complexity, - uint32_t num_threads); - void consolidate_delete(); - size_t num_points(); - - - private: - const uint32_t _initial_search_complexity; - const diskann::IndexWriteParameters _write_parameters; - diskann::Index _index; -}; - -}; // namespace diskannpy \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_disk_index.h b/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_disk_index.h deleted file mode 100644 index a3b79c4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_disk_index.h +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include -#include - -#ifdef _WINDOWS -#include "windows_aligned_file_reader.h" -#elif __APPLE__ -#include "apple_aligned_file_reader.h" -#else -#include "linux_aligned_file_reader.h" -#endif - -#include "common.h" -#include "pq_flash_index.h" - -namespace py = pybind11; - -namespace diskannpy -{ - -#ifdef _WINDOWS -typedef WindowsAlignedFileReader PlatformSpecificAlignedFileReader; -#elif __APPLE__ -typedef AppleAlignedFileReader PlatformSpecificAlignedFileReader; -#else -typedef LinuxAlignedFileReader PlatformSpecificAlignedFileReader; -#endif - -template class StaticDiskIndex -{ - public: - StaticDiskIndex(diskann::Metric metric, const std::string &index_path_prefix, uint32_t num_threads, - size_t num_nodes_to_cache, uint32_t cache_mechanism, const std::string &pq_prefix, - const std::string &partition_prefix); - - void cache_bfs_levels(size_t num_nodes_to_cache); - - void cache_sample_paths(size_t num_nodes_to_cache, const std::string &warmup_query_file, uint32_t num_threads); - - NeighborsAndDistances search(py::array_t &query, - uint64_t knn, uint64_t complexity, uint64_t beam_width, - bool USE_DEFERRED_FETCH = false, bool skip_search_reorder = false, - bool recompute_beighbor_embeddings = false, bool dedup_node_dis = false, - float prune_ratio = 0, bool batch_recompute = false, - bool global_pruning = false); - - NeighborsAndDistances batch_search( - py::array_t &queries, uint64_t num_queries, uint64_t knn, - uint64_t complexity, uint64_t beam_width, uint32_t num_threads, bool USE_DEFERRED_FETCH = false, - bool skip_search_reorder = false, bool recompute_beighbor_embeddings = false, bool dedup_node_dis = false, - float prune_ratio = 0, bool batch_recompute = false, bool global_pruning = false); - - private: - std::shared_ptr _reader; - std::shared_ptr _graph_reader; - diskann::PQFlashIndex

_index; -}; -} // namespace diskannpy diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_memory_index.h b/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_memory_index.h deleted file mode 100644 index 6ed5a08..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/include/static_memory_index.h +++ /dev/null @@ -1,40 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#pragma once - -#include -#include - -#include -#include - -#include "common.h" -#include "index.h" - -namespace py = pybind11; - -namespace diskannpy -{ - -template class StaticMemoryIndex -{ - public: - StaticMemoryIndex(diskann::Metric m, const std::string &index_prefix, size_t num_points, size_t dimensions, - uint32_t num_threads, uint32_t initial_search_complexity); - - NeighborsAndDistances search(py::array_t &query, - uint64_t knn, uint64_t complexity); - - NeighborsAndDistances search_with_filter( - py::array_t &query, uint64_t knn, uint64_t complexity, - filterT filter); - - NeighborsAndDistances batch_search( - py::array_t &queries, uint64_t num_queries, uint64_t knn, - uint64_t complexity, uint32_t num_threads); - - private: - diskann::Index _index; -}; -} // namespace diskannpy \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/__init__.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/__init__.py deleted file mode 100644 index c2e1b07..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/__init__.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -""" -# Documentation Overview -`diskannpy` is mostly structured around 2 distinct processes: [Index Builder Functions](#index-builders) and [Search Classes](#search-classes) - -It also includes a few nascent [utilities](#utilities). - -And lastly, it makes substantial use of type hints, with various shorthand [type aliases](#parameter-and-response-type-aliases) documented. -When reading the `diskannpy` code we refer to the type aliases, though `pdoc` helpfully expands them. - -## Index Builders -- `build_disk_index` - To build an index that cannot fully fit into memory when searching -- `build_memory_index` - To build an index that can fully fit into memory when searching - -## Search Classes -- `StaticMemoryIndex` - for indices that can fully fit in memory and won't be changed during the search operations -- `StaticDiskIndex` - for indices that cannot fully fit in memory, thus relying on disk IO to search, and also won't be changed during search operations -- `DynamicMemoryIndex` - for indices that can fully fit in memory and will be mutated via insert/deletion operations as well as search operations - -## Parameter Defaults -- `diskannpy.defaults` - Default values exported from the C++ extension for Python users - -## Parameter and Response Type Aliases -- `DistanceMetric` - What distance metrics does `diskannpy` support? -- `VectorDType` - What vector datatypes does `diskannpy` support? -- `QueryResponse` - What can I expect as a response to my search? -- `QueryResponseBatch` - What can I expect as a response to my batch search? -- `VectorIdentifier` - What types do `diskannpy` support as vector identifiers? -- `VectorIdentifierBatch` - A batch of identifiers of the exact same type. The type can change, but they must **all** change. -- `VectorLike` - How does a vector look to `diskannpy`, to be inserted or searched with. -- `VectorLikeBatch` - A batch of those vectors, to be inserted or searched with. -- `Metadata` - DiskANN vector binary file metadata (num_points, vector_dim) - -## Utilities -- `vectors_to_file` - Turns a 2 dimensional `numpy.typing.NDArray[VectorDType]` with shape `(number_of_points, vector_dim)` into a DiskANN vector bin file. -- `vectors_from_file` - Reads a DiskANN vector bin file representing stored vectors into a numpy ndarray. -- `vectors_metadata_from_file` - Reads metadata stored in a DiskANN vector bin file without reading the entire file -- `tags_to_file` - Turns a 1 dimensional `numpy.typing.NDArray[VectorIdentifier]` into a DiskANN tags bin file. -- `tags_from_file` - Reads a DiskANN tags bin file representing stored tags into a numpy ndarray. -- `valid_dtype` - Checks if a given vector dtype is supported by `diskannpy` -""" - -from typing import Any, Literal, NamedTuple, Type, Union - -import numpy as np -from numpy import typing as npt - -DistanceMetric = Literal["l2", "mips", "cosine"] -""" Type alias for one of {"l2", "mips", "cosine"} """ -VectorDType = Union[Type[np.float32], Type[np.int8], Type[np.uint8]] -""" Type alias for one of {`numpy.float32`, `numpy.int8`, `numpy.uint8`} """ -VectorLike = npt.NDArray[VectorDType] -""" Type alias for something that can be treated as a vector """ -VectorLikeBatch = npt.NDArray[VectorDType] -""" Type alias for a batch of VectorLikes """ -VectorIdentifier = np.uint32 -""" -Type alias for a vector identifier, whether it be an implicit array index identifier from StaticMemoryIndex or -StaticDiskIndex, or an explicit tag identifier from DynamicMemoryIndex -""" -VectorIdentifierBatch = npt.NDArray[np.uint32] -""" Type alias for a batch of VectorIdentifiers """ - - -class QueryResponse(NamedTuple): - """ - Tuple with two values, identifiers and distances. Both are 1d arrays, positionally correspond, and will contain the - nearest neighbors from [0..k_neighbors) - """ - - identifiers: npt.NDArray[VectorIdentifier] - """ A `numpy.typing.NDArray[VectorIdentifier]` array of vector identifiers, 1 dimensional """ - distances: npt.NDArray[np.float32] - """ - A `numpy.typing.NDAarray[numpy.float32]` of distances as calculated by the distance metric function, 1 dimensional - """ - - -class QueryResponseBatch(NamedTuple): - """ - Tuple with two values, identifiers and distances. Both are 2d arrays, with dimensionality determined by the - rows corresponding to the number of queries made, and the columns corresponding to the k neighbors - requested. The two 2d arrays have an implicit, position-based relationship - """ - - identifiers: npt.NDArray[VectorIdentifier] - """ - A `numpy.typing.NDArray[VectorIdentifier]` array of vector identifiers, 2 dimensional. The row corresponds to index - of the query, and the column corresponds to the k neighbors requested - """ - distances: np.ndarray[np.float32] - """ - A `numpy.typing.NDAarray[numpy.float32]` of distances as calculated by the distance metric function, 2 dimensional. - The row corresponds to the index of the query, and the column corresponds to the distance of the query to the - *k-th* neighbor - """ - - -from . import defaults -from ._builder import build_disk_index, build_memory_index -from ._common import valid_dtype -from ._dynamic_memory_index import DynamicMemoryIndex -from ._files import ( - Metadata, - tags_from_file, - tags_to_file, - vectors_from_file, - vectors_metadata_from_file, - vectors_to_file, -) -from ._static_disk_index import StaticDiskIndex -from ._static_memory_index import StaticMemoryIndex - -__all__ = [ - "build_disk_index", - "build_memory_index", - "StaticDiskIndex", - "StaticMemoryIndex", - "DynamicMemoryIndex", - "defaults", - "DistanceMetric", - "VectorDType", - "QueryResponse", - "QueryResponseBatch", - "VectorIdentifier", - "VectorIdentifierBatch", - "VectorLike", - "VectorLikeBatch", - "Metadata", - "vectors_metadata_from_file", - "vectors_to_file", - "vectors_from_file", - "tags_to_file", - "tags_from_file", - "valid_dtype", -] diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.py deleted file mode 100644 index 6567020..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.py +++ /dev/null @@ -1,349 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import json -import os -import shutil -from pathlib import Path -from typing import Optional, Tuple, Union - -import numpy as np - -from . import DistanceMetric, VectorDType, VectorIdentifierBatch, VectorLikeBatch -from . import _diskannpy as _native_dap -from ._common import ( - _assert, - _assert_is_nonnegative_uint32, - _assert_is_positive_uint32, - _castable_dtype_or_raise, - _valid_metric, - _write_index_metadata, - valid_dtype, -) -from ._diskannpy import defaults -from ._files import tags_to_file, vectors_metadata_from_file, vectors_to_file - - -def _valid_path_and_dtype( - data: Union[str, VectorLikeBatch], - vector_dtype: VectorDType, - index_path: str, - index_prefix: str, -) -> Tuple[str, VectorDType]: - if isinstance(data, str): - vector_bin_path = data - _assert( - Path(data).exists() and Path(data).is_file(), - "if data is of type `str`, it must both exist and be a file", - ) - vector_dtype_actual = valid_dtype(vector_dtype) - else: - vector_bin_path = os.path.join(index_path, f"{index_prefix}_vectors.bin") - # if Path(vector_bin_path).exists(): - # raise ValueError( - # f"The path {vector_bin_path} already exists. Remove it and try again." - # ) - vector_dtype_actual = valid_dtype(data.dtype) - # vectors_to_file(vector_file=vector_bin_path, vectors=data) - - return vector_bin_path, vector_dtype_actual - - -def build_disk_index( - data: Union[str, VectorLikeBatch], - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - search_memory_maximum: float, - build_memory_maximum: float, - num_threads: int, - pq_disk_bytes: int = defaults.PQ_DISK_BYTES, - vector_dtype: Optional[VectorDType] = None, - index_prefix: str = "ann", - codebook_prefix: str = "", -) -> None: - """ - This function will construct a DiskANN disk index. Disk indices are ideal for very large datasets that - are too large to fit in memory. Memory is still used, but it is primarily used to provide precise disk - locations for fast retrieval of smaller subsets of the index without compromising much on recall. - - If you provide a numpy array, it will save this array to disk in a temp location - in the format DiskANN's PQ Flash Index builder requires. This temp folder is deleted upon index creation completion - or error. - - ## Distance Metric and Vector Datatype Restrictions - | Metric \ Datatype | np.float32 | np.uint8 | np.int8 | - |-------------------|------------|----------|---------| - | L2 | ✅ | ✅ | ✅ | - | MIPS | ✅ | ❌ | ❌ | - | Cosine [^bug-in-disk-cosine] | ❌ | ❌ | ❌ | - - [^bug-in-disk-cosine]: For StaticDiskIndex, Cosine distances are not currently supported. - - ### Parameters - - **data**: Either a `str` representing a path to a DiskANN vector bin file, or a numpy.ndarray, - of a supported dtype, in 2 dimensions. Note that `vector_dtype` must be provided if data is a `str` - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. - - **index_directory**: The index files will be saved to this **existing** directory path - - **complexity**: The size of the candidate nearest neighbor list to use when building the index. Values between 75 - and 200 are typical. Larger values will take more time to build but result in indices that provide higher recall - for the same search complexity. Use a value that is at least as large as `graph_degree` unless you are prepared - to compromise on quality - - **graph_degree**: The degree of the graph index, typically between 60 and 150. A larger maximum degree will - result in larger indices and longer indexing times, but better search quality. - - **search_memory_maximum**: Build index with the expectation that the search will use at most - `search_memory_maximum`, in gb. - - **build_memory_maximum**: Build index using at most `build_memory_maximum` in gb. Building processes typically - require more memory, while search memory can be reduced. - - **num_threads**: Number of threads to use when creating this index. `0` is used to indicate all available - logical processors should be used. - - **pq_disk_bytes**: Use `0` to store uncompressed data on SSD. This allows the index to asymptote to 100% - recall. If your vectors are too large to store in SSD, this parameter provides the option to compress the - vectors using PQ for storing on SSD. This will trade off recall. You would also want this to be greater - than the number of bytes used for the PQ compressed data stored in-memory. Default is `0`. - - **vector_dtype**: Required if the provided `data` is of type `str`, else we use the `data.dtype` if np array. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - """ - - _assert( - (isinstance(data, str) and vector_dtype is not None) - or isinstance(data, np.ndarray), - "vector_dtype is required if data is a str representing a path to the vector bin file", - ) - dap_metric = _valid_metric(distance_metric) - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_positive_uint32(graph_degree, "graph_degree") - _assert(search_memory_maximum > 0, "search_memory_maximum must be larger than 0") - _assert(build_memory_maximum > 0, "build_memory_maximum must be larger than 0") - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_nonnegative_uint32(pq_disk_bytes, "pq_disk_bytes") - _assert(index_prefix != "", "index_prefix cannot be an empty string") - - index_path = Path(index_directory) - _assert( - index_path.exists() and index_path.is_dir(), - "index_directory must both exist and be a directory", - ) - - vector_bin_path, vector_dtype_actual = _valid_path_and_dtype( - data, vector_dtype, index_directory, index_prefix - ) - _assert(dap_metric != _native_dap.COSINE, "Cosine is currently not supported in StaticDiskIndex") - if dap_metric == _native_dap.INNER_PRODUCT: - _assert( - vector_dtype_actual == np.float32, - "Integral vector dtypes (np.uint8, np.int8) are not supported with distance metric mips" - ) - - num_points, dimensions = vectors_metadata_from_file(vector_bin_path) - - if vector_dtype_actual == np.uint8: - _builder = _native_dap.build_disk_uint8_index - elif vector_dtype_actual == np.int8: - _builder = _native_dap.build_disk_int8_index - else: - _builder = _native_dap.build_disk_float_index - - index_prefix_path = os.path.join(index_directory, index_prefix) - - _builder( - distance_metric=dap_metric, - data_file_path=vector_bin_path, - index_prefix_path=index_prefix_path, - complexity=complexity, - graph_degree=graph_degree, - final_index_ram_limit=search_memory_maximum, - indexing_ram_budget=build_memory_maximum, - num_threads=num_threads, - pq_disk_bytes=pq_disk_bytes, - codebook_prefix=codebook_prefix, - ) - _write_index_metadata( - index_prefix_path, vector_dtype_actual, dap_metric, num_points, dimensions - ) - - -def build_memory_index( - data: Union[str, VectorLikeBatch], - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - num_threads: int, - alpha: float = defaults.ALPHA, - use_pq_build: bool = defaults.USE_PQ_BUILD, - num_pq_bytes: int = defaults.NUM_PQ_BYTES, - use_opq: bool = defaults.USE_OPQ, - vector_dtype: Optional[VectorDType] = None, - tags: Union[str, VectorIdentifierBatch] = "", - filter_labels: Optional[list[list[str]]] = None, - universal_label: str = "", - filter_complexity: int = defaults.FILTER_COMPLEXITY, - index_prefix: str = "ann", -) -> None: - """ - This function will construct a DiskANN memory index. Memory indices are ideal for smaller datasets whose - indices can fit into memory. Memory indices are faster than disk indices, but usually cannot scale to massive - sizes in an individual index on an individual machine. - - `diskannpy`'s memory indices take two forms: a `diskannpy.StaticMemoryIndex`, which will not be mutated, only - searched upon, and a `diskannpy.DynamicMemoryIndex`, which can be mutated AND searched upon in the same process. - - ## Important Note: - You **must** determine the type of index you are building for. If you are building for a - `diskannpy.DynamicMemoryIndex`, you **must** supply a valid value for the `tags` parameter. **Do not supply - tags if the index is intended to be `diskannpy.StaticMemoryIndex`**! - - ## Distance Metric and Vector Datatype Restrictions - - | Metric \ Datatype | np.float32 | np.uint8 | np.int8 | - |-------------------|------------|----------|---------| - | L2 | ✅ | ✅ | ✅ | - | MIPS | ✅ | ❌ | ❌ | - | Cosine | ✅ | ✅ | ✅ | - - ### Parameters - - - **data**: Either a `str` representing a path to an existing DiskANN vector bin file, or a numpy.ndarray of a - supported dtype in 2 dimensions. Note that `vector_dtype` must be provided if `data` is a `str`. - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. - - **index_directory**: The index files will be saved to this **existing** directory path - - **complexity**: The size of the candidate nearest neighbor list to use when building the index. Values between 75 - and 200 are typical. Larger values will take more time to build but result in indices that provide higher recall - for the same search complexity. Use a value that is at least as large as `graph_degree` unless you are prepared - to compromise on quality - - **graph_degree**: The degree of the graph index, typically between 60 and 150. A larger maximum degree will - result in larger indices and longer indexing times, but better search quality. - - **num_threads**: Number of threads to use when creating this index. `0` is used to indicate all available - logical processors should be used. - - **alpha**: The alpha parameter (>=1) is used to control the nature and number of points that are added to the - graph. A higher alpha value (e.g., 1.4) will result in fewer hops (and IOs) to convergence, but probably more - distance comparisons compared to a lower alpha value. - - **use_pq_build**: Use product quantization during build. Product quantization is a lossy compression technique - that can reduce the size of the index on disk. This will trade off recall. Default is `True`. - - **num_pq_bytes**: The number of bytes used to store the PQ compressed data in memory. This will trade off recall. - Default is `0`. - - **use_opq**: Use optimized product quantization during build. - - **vector_dtype**: Required if the provided `data` is of type `str`, else we use the `data.dtype` if np array. - - **tags**: Tags can be defined either as a path on disk to an existing .tags file, or provided as a np.array of - the same length as the number of vectors. Tags are used to identify vectors in the index via your *own* - numbering conventions, and is absolutely required for loading DynamicMemoryIndex indices `from_file`. - - **filter_labels**: An optional, but exhaustive list of categories for each vector. This is used to filter - search results by category. If provided, this must be a list of lists, where each inner list is a list of - categories for the corresponding vector. For example, if you have 3 vectors, and the first vector belongs to - categories "a" and "b", the second vector belongs to category "b", and the third vector belongs to no categories, - you would provide `filter_labels=[["a", "b"], ["b"], []]`. If you do not want to provide categories for a - particular vector, you can provide an empty list. If you do not want to provide categories for any vectors, - you can provide `None` for this parameter (which is the default) - - **universal_label**: An optional label that indicates that this vector should be included in *every* search - in which it also meets the knn search criteria. - - **filter_complexity**: Complexity to use when using filters. Default is 0. 0 is strictly invalid if you are - using filters. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - """ - _assert( - (isinstance(data, str) and vector_dtype is not None) - or isinstance(data, np.ndarray), - "vector_dtype is required if data is a str representing a path to the vector bin file", - ) - dap_metric = _valid_metric(distance_metric) - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_positive_uint32(graph_degree, "graph_degree") - _assert( - alpha >= 1, - "alpha must be >= 1, and realistically should be kept between [1.0, 2.0)", - ) - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_nonnegative_uint32(num_pq_bytes, "num_pq_bytes") - _assert_is_nonnegative_uint32(filter_complexity, "filter_complexity") - _assert(index_prefix != "", "index_prefix cannot be an empty string") - _assert( - filter_labels is None or filter_complexity > 0, - "if filter_labels is provided, filter_complexity must not be 0" - ) - - index_path = Path(index_directory) - _assert( - index_path.exists() and index_path.is_dir(), - "index_directory must both exist and be a directory", - ) - - vector_bin_path, vector_dtype_actual = _valid_path_and_dtype( - data, vector_dtype, index_directory, index_prefix - ) - if dap_metric == _native_dap.INNER_PRODUCT: - _assert( - vector_dtype_actual == np.float32, - "Integral vector dtypes (np.uint8, np.int8) are not supported with distance metric mips" - ) - - num_points, dimensions = vectors_metadata_from_file(vector_bin_path) - if filter_labels is not None: - _assert( - len(filter_labels) == num_points, - "filter_labels must be the same length as the number of points" - ) - - if vector_dtype_actual == np.uint8: - _builder = _native_dap.build_memory_uint8_index - elif vector_dtype_actual == np.int8: - _builder = _native_dap.build_memory_int8_index - else: - _builder = _native_dap.build_memory_float_index - - index_prefix_path = os.path.join(index_directory, index_prefix) - - filter_labels_file = "" - if filter_labels is not None: - label_counts = {} - filter_labels_file = f"{index_prefix_path}_pylabels.txt" - with open(filter_labels_file, "w") as labels_file: - for labels in filter_labels: - for label in labels: - label_counts[label] = 1 if label not in label_counts else label_counts[label] + 1 - if len(labels) == 0: - print("default", file=labels_file) - else: - print(",".join(labels), file=labels_file) - with open(f"{index_prefix_path}_label_metadata.json", "w") as label_metadata_file: - json.dump(label_counts, label_metadata_file, indent=True) - - if isinstance(tags, str) and tags != "": - use_tags = True - shutil.copy(tags, index_prefix_path + ".tags") - elif not isinstance(tags, str): - use_tags = True - tags_as_array = _castable_dtype_or_raise(tags, expected=np.uint32) - _assert(len(tags_as_array.shape) == 1, "Provided tags must be 1 dimensional") - _assert( - tags_as_array.shape[0] == num_points, - "Provided tags must contain an identical population to the number of points, " - f"{tags_as_array.shape[0]=}, {num_points=}", - ) - tags_to_file(index_prefix_path + ".tags", tags_as_array) - else: - use_tags = False - - _builder( - distance_metric=dap_metric, - data_file_path=vector_bin_path, - index_output_path=index_prefix_path, - complexity=complexity, - graph_degree=graph_degree, - alpha=alpha, - num_threads=num_threads, - use_pq_build=use_pq_build, - num_pq_bytes=num_pq_bytes, - use_opq=use_opq, - use_tags=use_tags, - filter_labels_file=filter_labels_file, - universal_label=universal_label, - filter_complexity=filter_complexity, - ) - - _write_index_metadata( - index_prefix_path, vector_dtype_actual, dap_metric, num_points, dimensions - ) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.pyi b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.pyi deleted file mode 100644 index 223e6c9..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_builder.pyi +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -from typing import BinaryIO, Optional, overload - -import numpy as np - -from . import DistanceMetric, VectorDType, VectorIdentifierBatch, VectorLikeBatch - -def numpy_to_diskann_file(vectors: np.ndarray, file_handler: BinaryIO): ... -@overload -def build_disk_index( - data: str, - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - search_memory_maximum: float, - build_memory_maximum: float, - num_threads: int, - pq_disk_bytes: int, - vector_dtype: VectorDType, - index_prefix: str, -) -> None: ... -@overload -def build_disk_index( - data: VectorLikeBatch, - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - search_memory_maximum: float, - build_memory_maximum: float, - num_threads: int, - pq_disk_bytes: int, - index_prefix: str, -) -> None: ... -@overload -def build_memory_index( - data: VectorLikeBatch, - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - alpha: float, - num_threads: int, - use_pq_build: bool, - num_pq_bytes: int, - use_opq: bool, - tags: Union[str, VectorIdentifierBatch], - filter_labels: Optional[list[list[str]]], - universal_label: str, - filter_complexity: int, - index_prefix: str -) -> None: ... -@overload -def build_memory_index( - data: str, - distance_metric: DistanceMetric, - index_directory: str, - complexity: int, - graph_degree: int, - alpha: float, - num_threads: int, - use_pq_build: bool, - num_pq_bytes: int, - use_opq: bool, - vector_dtype: VectorDType, - tags: Union[str, VectorIdentifierBatch], - filter_labels_file: Optional[list[list[str]]], - universal_label: str, - filter_complexity: int, - index_prefix: str -) -> None: ... diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_common.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_common.py deleted file mode 100644 index 2b28802..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_common.py +++ /dev/null @@ -1,251 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import os -import warnings -from enum import Enum -from pathlib import Path -from typing import Literal, NamedTuple, Optional, Tuple, Type, Union - -import numpy as np - -from . import ( - DistanceMetric, - VectorDType, - VectorIdentifierBatch, - VectorLike, - VectorLikeBatch, -) -from . import _diskannpy as _native_dap - -__ALL__ = ["valid_dtype"] - -_VALID_DTYPES = [np.float32, np.int8, np.uint8] - - -def valid_dtype(dtype: Type) -> VectorDType: - """ - Utility method to determine whether the provided dtype is supported by `diskannpy`, and if so, the canonical - dtype we will use internally (e.g. np.single -> np.float32) - """ - _assert_dtype(dtype) - if dtype == np.uint8: - return np.uint8 - if dtype == np.int8: - return np.int8 - if dtype == np.float32: - return np.float32 - - -def _assert(statement_eval: bool, message: str): - if not statement_eval: - raise ValueError(message) - - -def _valid_metric(metric: str) -> _native_dap.Metric: - if not isinstance(metric, str): - raise ValueError("distance_metric must be a string") - if metric.lower() == "l2": - return _native_dap.L2 - elif metric.lower() == "mips": - return _native_dap.INNER_PRODUCT - elif metric.lower() == "cosine": - return _native_dap.COSINE - else: - raise ValueError("distance_metric must be one of 'l2', 'mips', or 'cosine'") - - -def _assert_dtype(dtype: Type): - _assert( - any(np.can_cast(dtype, _dtype) for _dtype in _VALID_DTYPES), - f"Vector dtype must be of one of type {{(np.single, np.float32), (np.byte, np.int8), (np.ubyte, np.uint8)}}", - ) - - -def _castable_dtype_or_raise( - data: Union[VectorLike, VectorLikeBatch, VectorIdentifierBatch], expected: np.dtype -) -> np.ndarray: - if isinstance(data, np.ndarray) and np.can_cast(data.dtype, expected): - return data.astype(expected, casting="safe") - else: - raise TypeError( - f"expecting a numpy ndarray of dtype {expected}, not a {type(data)}" - ) - - -def _assert_2d(vectors: np.ndarray, name: str): - _assert(len(vectors.shape) == 2, f"{name} must be 2d numpy array") - - -__MAX_UINT32_VAL = 4_294_967_295 - - -def _assert_is_positive_uint32(test_value: int, parameter: str): - _assert( - test_value is not None and 0 < test_value < __MAX_UINT32_VAL, - f"{parameter} must be a positive integer in the uint32 range", - ) - - -def _assert_is_nonnegative_uint32(test_value: int, parameter: str): - _assert( - test_value is not None and -1 < test_value < __MAX_UINT32_VAL, - f"{parameter} must be a non-negative integer in the uint32 range", - ) - - -def _assert_is_nonnegative_uint64(test_value: int, parameter: str): - _assert( - -1 < test_value, - f"{parameter} must be a non-negative integer in the uint64 range", - ) - - -def _assert_existing_directory(path: str, parameter: str): - _path = Path(path) - _assert( - _path.exists() and _path.is_dir(), f"{parameter} must be an existing directory" - ) - - -def _assert_existing_file(path: str, parameter: str): - _path = Path(path) - _assert(_path.exists() and _path.is_file(), f"{parameter} must be an existing file") - - -class _DataType(Enum): - FLOAT32 = 0 - INT8 = 1 - UINT8 = 2 - - @classmethod - def from_type(cls, vector_dtype: VectorDType) -> "DataType": - if vector_dtype == np.float32: - return cls.FLOAT32 - if vector_dtype == np.int8: - return cls.INT8 - if vector_dtype == np.uint8: - return cls.UINT8 - - def to_type(self) -> VectorDType: - if self is _DataType.FLOAT32: - return np.float32 - if self is _DataType.INT8: - return np.int8 - if self is _DataType.UINT8: - return np.uint8 - - -class _Metric(Enum): - L2 = 0 - MIPS = 1 - COSINE = 2 - - @classmethod - def from_native(cls, metric: _native_dap.Metric) -> "_Metric": - if metric == _native_dap.L2: - return cls.L2 - if metric == _native_dap.INNER_PRODUCT: - return cls.MIPS - if metric == _native_dap.COSINE: - return cls.COSINE - - def to_native(self) -> _native_dap.Metric: - if self is _Metric.L2: - return _native_dap.L2 - if self is _Metric.MIPS: - return _native_dap.INNER_PRODUCT - if self is _Metric.COSINE: - return _native_dap.COSINE - - def to_str(self) -> _native_dap.Metric: - if self is _Metric.L2: - return "l2" - if self is _Metric.MIPS: - return "mips" - if self is _Metric.COSINE: - return "cosine" - - -def _build_metadata_path(index_path_and_prefix: str) -> str: - return index_path_and_prefix + "_metadata.bin" - - -def _write_index_metadata( - index_path_and_prefix: str, - dtype: VectorDType, - metric: _native_dap.Metric, - num_points: int, - dimensions: int, -): - np.array( - [ - _DataType.from_type(dtype).value, - _Metric.from_native(metric).value, - num_points, - dimensions, - ], - dtype=np.uint64, - ).tofile(_build_metadata_path(index_path_and_prefix)) - - -def _read_index_metadata( - index_path_and_prefix: str, -) -> Optional[Tuple[VectorDType, str, np.uint64, np.uint64]]: - path = _build_metadata_path(index_path_and_prefix) - if not Path(path).exists(): - return None - else: - metadata = np.fromfile(path, dtype=np.uint64, count=-1) - return ( - _DataType(int(metadata[0])).to_type(), - _Metric(int(metadata[1])).to_str(), - metadata[2], - metadata[3], - ) - - -def _ensure_index_metadata( - index_path_and_prefix: str, - vector_dtype: Optional[VectorDType], - distance_metric: Optional[DistanceMetric], - max_vectors: int, - dimensions: Optional[int], - warn_size_exceeded: bool = False, -) -> Tuple[VectorDType, str, np.uint64, np.uint64]: - possible_metadata = _read_index_metadata(index_path_and_prefix) - if possible_metadata is None: - _assert( - all([vector_dtype, distance_metric, dimensions]), - "distance_metric, vector_dtype, and dimensions must provided if a corresponding metadata file has not " - "been built for this index, such as when an index was built via the CLI tools or prior to the addition " - "of a metadata file", - ) - _assert_dtype(vector_dtype) - _assert_is_positive_uint32(max_vectors, "max_vectors") - _assert_is_positive_uint32(dimensions, "dimensions") - return vector_dtype, distance_metric, max_vectors, dimensions # type: ignore - else: - vector_dtype, distance_metric, num_vectors, dimensions = possible_metadata - if warn_size_exceeded: - if max_vectors is not None and num_vectors > max_vectors: - warnings.warn( - "The number of vectors in the saved index exceeds the max_vectors parameter. " - "max_vectors is being adjusted to accommodate the dataset, but any insertions will fail." - ) - max_vectors = num_vectors - if num_vectors == max_vectors: - warnings.warn( - "The number of vectors in the saved index equals max_vectors parameter. Any insertions will fail." - ) - return possible_metadata - - -def _valid_index_prefix(index_directory: str, index_prefix: str) -> str: - _assert( - index_directory is not None and index_directory != "", - "index_directory cannot be None or empty", - ) - _assert_existing_directory(index_directory, "index_directory") - _assert(index_prefix != "", "index_prefix cannot be an empty string") - return os.path.join(index_directory, index_prefix) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_diskannpy.cpython-310-x86_64-linux-gnu.so.bak b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_diskannpy.cpython-310-x86_64-linux-gnu.so.bak deleted file mode 100755 index 5741ecd..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_diskannpy.cpython-310-x86_64-linux-gnu.so.bak and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_dynamic_memory_index.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_dynamic_memory_index.py deleted file mode 100644 index cdf6432..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_dynamic_memory_index.py +++ /dev/null @@ -1,511 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import os -import warnings -from pathlib import Path -from typing import Optional - -import numpy as np - -from . import ( - DistanceMetric, - QueryResponse, - QueryResponseBatch, - VectorDType, - VectorIdentifier, - VectorIdentifierBatch, - VectorLike, - VectorLikeBatch, -) -from . import _diskannpy as _native_dap -from ._common import ( - _assert, - _assert_2d, - _assert_dtype, - _assert_existing_directory, - _assert_is_nonnegative_uint32, - _assert_is_positive_uint32, - _castable_dtype_or_raise, - _ensure_index_metadata, - _valid_index_prefix, - _valid_metric, - _write_index_metadata, -) -from ._diskannpy import defaults - -__ALL__ = ["DynamicMemoryIndex"] - - -class DynamicMemoryIndex: - """ - A DynamicMemoryIndex instance is used to both search and mutate a `diskannpy` memory index. This index is unlike - either `diskannpy.StaticMemoryIndex` or `diskannpy.StaticDiskIndex` in the following ways: - - - It requires an explicit vector identifier for each vector added to it. - - Insert and (lazy) deletion operations are provided for a flexible, living index - - The mutable aspect of this index will absolutely impact search time performance as new vectors are added and - old deleted. `DynamicMemoryIndex.consolidate_deletes()` should be called periodically to restructure the index - to remove deleted vectors and improve per-search performance, at the cost of an expensive index consolidation to - occur. - """ - - @classmethod - def from_file( - cls, - index_directory: str, - max_vectors: int, - complexity: int, - graph_degree: int, - saturate_graph: bool = defaults.SATURATE_GRAPH, - max_occlusion_size: int = defaults.MAX_OCCLUSION_SIZE, - alpha: float = defaults.ALPHA, - num_threads: int = defaults.NUM_THREADS, - filter_complexity: int = defaults.FILTER_COMPLEXITY, - num_frozen_points: int = defaults.NUM_FROZEN_POINTS_DYNAMIC, - initial_search_complexity: int = 0, - search_threads: int = 0, - concurrent_consolidation: bool = True, - index_prefix: str = "ann", - distance_metric: Optional[DistanceMetric] = None, - vector_dtype: Optional[VectorDType] = None, - dimensions: Optional[int] = None, - ) -> "DynamicMemoryIndex": - """ - The `from_file` classmethod is used to load a previously saved index from disk. This index *must* have been - created with a valid `tags` file or `tags` np.ndarray of `diskannpy.VectorIdentifier`s. It is *strongly* - recommended that you use the same parameters as the `diskannpy.build_memory_index()` function that created - the index. - - ### Parameters - - **index_directory**: The directory containing the index files. This directory must contain the following - files: - - `{index_prefix}.data` - - `{index_prefix}.tags` - - `{index_prefix}` - - It may also include the following optional files: - - `{index_prefix}_vectors.bin`: Optional. `diskannpy` builder functions may create this file in the - `index_directory` if the index was created from a numpy array - - `{index_prefix}_metadata.bin`: Optional. `diskannpy` builder functions create this file to store metadata - about the index, such as vector dtype, distance metric, number of vectors and vector dimensionality. - If an index is built from the `diskann` cli tools, this file will not exist. - - **max_vectors**: Capacity of the memory index including space for future insertions. - - **complexity**: Complexity (a.k.a `L`) references the size of the list we store candidate approximate - neighbors in. It's used during save (which is an index rebuild), and it's used as an initial search size to - warm up our index and lower the latency for initial real searches. - - **graph_degree**: Graph degree (a.k.a. `R`) is the maximum degree allowed for a node in the index's graph - structure. This degree will be pruned throughout the course of the index build, but it will never grow beyond - this value. Higher R values require longer index build times, but may result in an index showing excellent - recall and latency characteristics. - - **saturate_graph**: If True, the adjacency list of each node will be saturated with neighbors to have exactly - `graph_degree` neighbors. If False, each node will have between 1 and `graph_degree` neighbors. - - **max_occlusion_size**: The maximum number of points that can be considered by occlude_list function. - - **alpha**: The alpha parameter (>=1) is used to control the nature and number of points that are added to the - graph. A higher alpha value (e.g., 1.4) will result in fewer hops (and IOs) to convergence, but probably - more distance comparisons compared to a lower alpha value. - - **num_threads**: Number of threads to use when creating this index. `0` indicates we should use all available - logical processors. - - **filter_complexity**: Complexity to use when using filters. Default is 0. - - **num_frozen_points**: Number of points to freeze. Default is 1. - - **initial_search_complexity**: Should be set to the most common `complexity` expected to be used during the - life of this `diskannpy.DynamicMemoryIndex` object. The working scratch memory allocated is based off of - `initial_search_complexity` * `search_threads`. Note that it may be resized if a `search` or `batch_search` - operation requests a space larger than can be accommodated by these values. - - **search_threads**: Should be set to the most common `num_threads` expected to be used during the - life of this `diskannpy.DynamicMemoryIndex` object. The working scratch memory allocated is based off of - `initial_search_complexity` * `search_threads`. Note that it may be resized if a `batch_search` - operation requests a space larger than can be accommodated by these values. - - **concurrent_consolidation**: This flag dictates whether consolidation can be run alongside inserts and - deletes, or whether the index is locked down to changes while consolidation is ongoing. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. Default is `None`. **This - value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, - you are required to provide it. - - **vector_dtype**: The vector dtype this index has been built with. **This value is only used if a - `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, you are required to provide it. - - **dimensions**: The vector dimensionality of this index. All new vectors inserted must be the same - dimensionality. **This value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it - does not exist, you are required to provide it. - - ### Returns - A `diskannpy.DynamicMemoryIndex` object, with the index loaded from disk and ready to use for insertions, - deletions, and searches. - - """ - index_prefix_path = _valid_index_prefix(index_directory, index_prefix) - - # do tags exist? - tags_file = index_prefix_path + ".tags" - _assert( - Path(tags_file).exists(), - f"The file {tags_file} does not exist in {index_directory}", - ) - vector_dtype, dap_metric, num_vectors, dimensions = _ensure_index_metadata( - index_prefix_path, vector_dtype, distance_metric, max_vectors, dimensions, warn_size_exceeded=True - ) - - index = cls( - distance_metric=dap_metric, # type: ignore - vector_dtype=vector_dtype, - dimensions=dimensions, - max_vectors=max_vectors, - complexity=complexity, - graph_degree=graph_degree, - saturate_graph=saturate_graph, - max_occlusion_size=max_occlusion_size, - alpha=alpha, - num_threads=num_threads, - filter_complexity=filter_complexity, - num_frozen_points=num_frozen_points, - initial_search_complexity=initial_search_complexity, - search_threads=search_threads, - concurrent_consolidation=concurrent_consolidation, - ) - index._index.load(index_prefix_path) - index._num_vectors = num_vectors # current number of vectors loaded - return index - - def __init__( - self, - distance_metric: DistanceMetric, - vector_dtype: VectorDType, - dimensions: int, - max_vectors: int, - complexity: int, - graph_degree: int, - saturate_graph: bool = defaults.SATURATE_GRAPH, - max_occlusion_size: int = defaults.MAX_OCCLUSION_SIZE, - alpha: float = defaults.ALPHA, - num_threads: int = defaults.NUM_THREADS, - filter_complexity: int = defaults.FILTER_COMPLEXITY, - num_frozen_points: int = defaults.NUM_FROZEN_POINTS_DYNAMIC, - initial_search_complexity: int = 0, - search_threads: int = 0, - concurrent_consolidation: bool = True, - ): - """ - The `diskannpy.DynamicMemoryIndex` represents our python API into a mutable DiskANN memory index. - - This constructor is used to create a new, empty index. If you wish to load a previously saved index from disk, - please use the `diskannpy.DynamicMemoryIndex.from_file` classmethod instead. - - ### Parameters - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. - - **vector_dtype**: One of {`np.float32`, `np.int8`, `np.uint8`}. The dtype of the vectors this index will - be storing. - - **dimensions**: The vector dimensionality of this index. All new vectors inserted must be the same - dimensionality. - - **max_vectors**: Capacity of the data store including space for future insertions - - **graph_degree**: Graph degree (a.k.a. `R`) is the maximum degree allowed for a node in the index's graph - structure. This degree will be pruned throughout the course of the index build, but it will never grow beyond - this value. Higher `graph_degree` values require longer index build times, but may result in an index showing - excellent recall and latency characteristics. - - **saturate_graph**: If True, the adjacency list of each node will be saturated with neighbors to have exactly - `graph_degree` neighbors. If False, each node will have between 1 and `graph_degree` neighbors. - - **max_occlusion_size**: The maximum number of points that can be considered by occlude_list function. - - **alpha**: The alpha parameter (>=1) is used to control the nature and number of points that are added to the - graph. A higher alpha value (e.g., 1.4) will result in fewer hops (and IOs) to convergence, but probably - more distance comparisons compared to a lower alpha value. - - **num_threads**: Number of threads to use when creating this index. `0` indicates we should use all available - logical processors. - - **filter_complexity**: Complexity to use when using filters. Default is 0. - - **num_frozen_points**: Number of points to freeze. Default is 1. - - **initial_search_complexity**: Should be set to the most common `complexity` expected to be used during the - life of this `diskannpy.DynamicMemoryIndex` object. The working scratch memory allocated is based off of - `initial_search_complexity` * `search_threads`. Note that it may be resized if a `search` or `batch_search` - operation requests a space larger than can be accommodated by these values. - - **search_threads**: Should be set to the most common `num_threads` expected to be used during the - life of this `diskannpy.DynamicMemoryIndex` object. The working scratch memory allocated is based off of - `initial_search_complexity` * `search_threads`. Note that it may be resized if a `batch_search` - operation requests a space larger than can be accommodated by these values. - - **concurrent_consolidation**: This flag dictates whether consolidation can be run alongside inserts and - deletes, or whether the index is locked down to changes while consolidation is ongoing. - - """ - self._num_vectors = 0 - self._removed_num_vectors = 0 - dap_metric = _valid_metric(distance_metric) - self._dap_metric = dap_metric - _assert_dtype(vector_dtype) - _assert_is_positive_uint32(dimensions, "dimensions") - - self._vector_dtype = vector_dtype - self._dimensions = dimensions - - _assert_is_positive_uint32(max_vectors, "max_vectors") - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_positive_uint32(graph_degree, "graph_degree") - _assert( - alpha >= 1, - "alpha must be >= 1, and realistically should be kept between [1.0, 2.0)", - ) - _assert_is_nonnegative_uint32(max_occlusion_size, "max_occlusion_size") - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_nonnegative_uint32(filter_complexity, "filter_complexity") - _assert_is_nonnegative_uint32(num_frozen_points, "num_frozen_points") - _assert_is_nonnegative_uint32( - initial_search_complexity, "initial_search_complexity" - ) - _assert_is_nonnegative_uint32(search_threads, "search_threads") - - self._max_vectors = max_vectors - self._complexity = complexity - self._graph_degree = graph_degree - - if vector_dtype == np.uint8: - _index = _native_dap.DynamicMemoryUInt8Index - elif vector_dtype == np.int8: - _index = _native_dap.DynamicMemoryInt8Index - else: - _index = _native_dap.DynamicMemoryFloatIndex - - self._index = _index( - distance_metric=dap_metric, - dimensions=dimensions, - max_vectors=max_vectors, - complexity=complexity, - graph_degree=graph_degree, - saturate_graph=saturate_graph, - max_occlusion_size=max_occlusion_size, - alpha=alpha, - num_threads=num_threads, - filter_complexity=filter_complexity, - num_frozen_points=num_frozen_points, - initial_search_complexity=initial_search_complexity, - search_threads=search_threads, - concurrent_consolidation=concurrent_consolidation, - ) - self._points_deleted = False - - def search( - self, query: VectorLike, k_neighbors: int, complexity: int - ) -> QueryResponse: - """ - Searches the index by a single query vector. - - ### Parameters - - **query**: 1d numpy array of the same dimensionality and dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - """ - _query = _castable_dtype_or_raise(query, expected=self._vector_dtype) - _assert(len(_query.shape) == 1, "query vector must be 1-d") - _assert( - _query.shape[0] == self._dimensions, - f"query vector must have the same dimensionality as the index; index dimensionality: {self._dimensions}, " - f"query dimensionality: {_query.shape[0]}", - ) - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_nonnegative_uint32(complexity, "complexity") - - if k_neighbors > complexity: - warnings.warn( - f"k_neighbors={k_neighbors} asked for, but list_size={complexity} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - neighbors, distances = self._index.search(query=_query, knn=k_neighbors, complexity=complexity) - return QueryResponse(identifiers=neighbors, distances=distances) - - def batch_search( - self, - queries: VectorLikeBatch, - k_neighbors: int, - complexity: int, - num_threads: int, - ) -> QueryResponseBatch: - """ - Searches the index by a batch of query vectors. - - This search is parallelized and far more efficient than searching for each vector individually. - - ### Parameters - - **queries**: 2d numpy array, with column dimensionality matching the index and row dimensionality being the - number of queries intended to search for in parallel. Dtype must match dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - - **num_threads**: Number of threads to use when searching this index. (>= 0), 0 = num_threads in system - """ - _queries = _castable_dtype_or_raise(queries, expected=self._vector_dtype) - _assert_2d(_queries, "queries") - _assert( - _queries.shape[1] == self._dimensions, - f"query vectors must have the same dimensionality as the index; index dimensionality: {self._dimensions}, " - f"query dimensionality: {_queries.shape[1]}", - ) - - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_nonnegative_uint32(num_threads, "num_threads") - - if k_neighbors > complexity: - warnings.warn( - f"k_neighbors={k_neighbors} asked for, but list_size={complexity} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - - num_queries, dim = queries.shape - neighbors, distances = self._index.batch_search( - queries=_queries, - num_queries=num_queries, - knn=k_neighbors, - complexity=complexity, - num_threads=num_threads, - ) - return QueryResponseBatch(identifiers=neighbors, distances=distances) - - def save(self, save_path: str, index_prefix: str = "ann"): - """ - Saves this index to file. - - ### Parameters - - **save_path**: The path to save these index files to. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - """ - if save_path == "": - raise ValueError("save_path cannot be empty") - if index_prefix == "": - raise ValueError("index_prefix cannot be empty") - - index_prefix = index_prefix.format(complexity=self._complexity, graph_degree=self._graph_degree) - _assert_existing_directory(save_path, "save_path") - save_path = os.path.join(save_path, index_prefix) - if self._points_deleted is True: - warnings.warn( - "DynamicMemoryIndex.save() currently requires DynamicMemoryIndex.consolidate_delete() to be called " - "prior to save when items have been marked for deletion. This is being done automatically now, though" - "it will increase the time it takes to save; on large sets of data it can take a substantial amount of " - "time. In the future, we will implement a faster save with unconsolidated deletes, but for now this is " - "required." - ) - self._index.consolidate_delete() - self._index.save( - save_path=save_path, compact_before_save=True - ) # we do not yet support uncompacted saves - _write_index_metadata( - save_path, - self._vector_dtype, - self._dap_metric, - self._index.num_points(), - self._dimensions, - ) - - def insert(self, vector: VectorLike, vector_id: VectorIdentifier): - """ - Inserts a single vector into the index with the provided vector_id. - - If this insertion will overrun the `max_vectors` count boundaries of this index, `consolidate_delete()` will - be executed automatically. - - ### Parameters - - **vector**: The vector to insert. Note that dtype must match. - - **vector_id**: The vector_id to use for this vector. - """ - _vector = _castable_dtype_or_raise(vector, expected=self._vector_dtype) - _assert(len(vector.shape) == 1, "insert vector must be 1-d") - _assert_is_positive_uint32(vector_id, "vector_id") - if self._num_vectors + 1 > self._max_vectors: - if self._removed_num_vectors > 0: - warnings.warn(f"Inserting this vector would overrun the max_vectors={self._max_vectors} specified at index " - f"construction. We are attempting to consolidate_delete() to make space.") - self.consolidate_delete() - else: - raise RuntimeError(f"Inserting this vector would overrun the max_vectors={self._max_vectors} specified " - f"at index construction. Unable to make space by consolidating deletions. The insert" - f"operation has failed.") - status = self._index.insert(_vector, np.uint32(vector_id)) - if status == 0: - self._num_vectors += 1 - else: - raise RuntimeError( - f"Insert was unable to complete successfully; error code returned from diskann C++ lib: {status}" - ) - - - def batch_insert( - self, - vectors: VectorLikeBatch, - vector_ids: VectorIdentifierBatch, - num_threads: int = 0, - ): - """ - Inserts a batch of vectors into the index with the provided vector_ids. - - If this batch insertion will overrun the `max_vectors` count boundaries of this index, `consolidate_delete()` - will be executed automatically. - - ### Parameters - - **vectors**: The 2d numpy array of vectors to insert. - - **vector_ids**: The 1d array of vector ids to use. This array must have the same number of elements as - the vectors array has rows. The dtype of vector_ids must be `np.uint32` - - **num_threads**: Number of threads to use when inserting into this index. (>= 0), 0 = num_threads in system - """ - _query = _castable_dtype_or_raise(vectors, expected=self._vector_dtype) - _assert(len(vectors.shape) == 2, "vectors must be a 2-d array") - _assert( - vectors.shape[0] == vector_ids.shape[0], - "Number of vectors must be equal to number of ids", - ) - _vectors = vectors.astype(dtype=self._vector_dtype, casting="safe", copy=False) - _vector_ids = vector_ids.astype(dtype=np.uint32, casting="safe", copy=False) - - if self._num_vectors + _vector_ids.shape[0] > self._max_vectors: - if self._max_vectors + self._removed_num_vectors >= _vector_ids.shape[0]: - warnings.warn(f"Inserting these vectors, count={_vector_ids.shape[0]} would overrun the " - f"max_vectors={self._max_vectors} specified at index construction. We are attempting to " - f"consolidate_delete() to make space.") - self.consolidate_delete() - else: - raise RuntimeError(f"Inserting these vectors count={_vector_ids.shape[0]} would overrun the " - f"max_vectors={self._max_vectors} specified at index construction. Unable to make " - f"space by consolidating deletions. The batch insert operation has failed.") - - statuses = self._index.batch_insert( - _vectors, _vector_ids, _vector_ids.shape[0], num_threads - ) - successes = [] - failures = [] - for i in range(0, len(statuses)): - if statuses[i] == 0: - successes.append(i) - else: - failures.append(i) - self._num_vectors += len(successes) - if len(failures) == 0: - return - failed_ids = vector_ids[failures] - raise RuntimeError( - f"During batch insert, the following vector_ids were unable to be inserted into the index: {failed_ids}. " - f"{len(successes)} were successfully inserted" - ) - - - def mark_deleted(self, vector_id: VectorIdentifier): - """ - Mark vector for deletion. This is a soft delete that won't return the vector id in any results, but does not - remove it from the underlying index files or memory structure. To execute a hard delete, call this method and - then call the much more expensive `consolidate_delete` method on this index. - ### Parameters - - **vector_id**: The vector id to delete. Must be a uint32. - """ - _assert_is_positive_uint32(vector_id, "vector_id") - self._points_deleted = True - self._removed_num_vectors += 1 - # we do not decrement self._num_vectors until consolidate_delete - self._index.mark_deleted(np.uint32(vector_id)) - - def consolidate_delete(self): - """ - This method actually restructures the DiskANN index to remove the items that have been marked for deletion. - """ - self._index.consolidate_delete() - self._points_deleted = False - self._num_vectors -= self._removed_num_vectors - self._removed_num_vectors = 0 diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_files.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_files.py deleted file mode 100644 index 7740c34..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_files.py +++ /dev/null @@ -1,122 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import warnings -from typing import BinaryIO, Literal, NamedTuple - -import numpy as np -import numpy.typing as npt - -from . import VectorDType, VectorIdentifierBatch, VectorLikeBatch -from ._common import _assert, _assert_2d, _assert_dtype, _assert_existing_file - - -class Metadata(NamedTuple): - """DiskANN binary vector files contain a small stanza containing some metadata about them.""" - - num_vectors: int - """ The number of vectors in the file. """ - dimensions: int - """ The dimensionality of the vectors in the file. """ - - -def vectors_metadata_from_file(vector_file: str) -> Metadata: - """ - Read the metadata from a DiskANN binary vector file. - ### Parameters - - **vector_file**: The path to the vector file to read the metadata from. - - ### Returns - `diskannpy.Metadata` - """ - _assert_existing_file(vector_file, "vector_file") - points, dims = np.fromfile(file=vector_file, dtype=np.int32, count=2) - return Metadata(points, dims) - - -def _write_bin(data: np.ndarray, file_handler: BinaryIO): - if len(data.shape) == 1: - _ = file_handler.write(np.array([data.shape[0], 1], dtype=np.int32).tobytes()) - else: - _ = file_handler.write(np.array(data.shape, dtype=np.int32).tobytes()) - _ = file_handler.write(data.tobytes()) - - -def vectors_to_file(vector_file: str, vectors: VectorLikeBatch) -> None: - """ - Utility function that writes a DiskANN binary vector formatted file to the location of your choosing. - - ### Parameters - - **vector_file**: The path to the vector file to write the vectors to. - - **vectors**: A 2d array of dtype `numpy.float32`, `numpy.uint8`, or `numpy.int8` - """ - _assert_dtype(vectors.dtype) - _assert_2d(vectors, "vectors") - with open(vector_file, "wb") as fh: - _write_bin(vectors, fh) - - -def vectors_from_file( - vector_file: str, - dtype: VectorDType, - use_memmap: bool = False, - mode: Literal["r", "r+"] = "r" -) -> npt.NDArray[VectorDType]: - """ - Read vectors from a DiskANN binary vector file. - - ### Parameters - - **vector_file**: The path to the vector file to read the vectors from. - - **dtype**: The data type of the vectors in the file. Ensure you match the data types exactly - - **use_memmap**: If True, return a np.memmap, else a standard np.ndarray will be returned - - **mode**: Read-only (r) or read-write (r+) (memmap only). Unlike np.memmap, default is read-only (r) - - ### Returns - `numpy.typing.NDArray[dtype] | numpy.memmap` - """ - assert mode in ["r", "r+"] - points, dims = vectors_metadata_from_file(vector_file) - if not use_memmap: - return np.fromfile(file=vector_file, dtype=dtype, offset=8).reshape(points, dims) - else: - return np.memmap(vector_file, dtype=dtype, mode=mode, offset=8, shape=(points, dims), order='C') - - -def tags_to_file(tags_file: str, tags: VectorIdentifierBatch) -> None: - """ - Write tags to a DiskANN binary tag file. - - ### Parameters - - **tags_file**: The path to the tag file to write the tags to. - - **tags**: A 1d array of dtype `numpy.uint32` containing the tags to write. If you have a 2d array of tags with - one column, you can pass it here and it will be reshaped and copied to a new array. It is more efficient for you - to reshape on your own without copying it first, as it should be a constant time operation vs. linear time - - """ - _assert(np.can_cast(tags.dtype, np.uint32), "valid tags must be uint32") - _assert( - len(tags.shape) == 1 or tags.shape[1] == 1, - "tags must be 1d or 2d with 1 column", - ) - if len(tags.shape) == 2: - warnings.warn( - "Tags in 2d with one column will be reshaped and copied to a new array. " - "It is more efficient for you to reshape without copying first." - ) - tags = tags.reshape(tags.shape[0], copy=True) - with open(tags_file, "wb") as fh: - _write_bin(tags.astype(np.uint32), fh) - - -def tags_from_file(tags_file: str) -> VectorIdentifierBatch: - """ - Read tags from a DiskANN binary tag file and return them as a 1d array of dtype `numpy.uint32`. - - ### Parameters - - **tags_file**: The path to the tag file to read the tags from. - """ - _assert_existing_file(tags_file, "tags_file") - points, dims = vectors_metadata_from_file( - tags_file - ) # tag files contain the same metadata stanza - return np.fromfile(file=tags_file, dtype=np.uint32, offset=8).reshape(points) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_disk_index.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_disk_index.py deleted file mode 100644 index 47af362..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_disk_index.py +++ /dev/null @@ -1,244 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import os -import warnings -from typing import Optional - -import numpy as np - -from . import ( - DistanceMetric, - QueryResponse, - QueryResponseBatch, - VectorDType, - VectorLike, - VectorLikeBatch, -) -from . import _diskannpy as _native_dap -from ._common import ( - _assert, - _assert_2d, - _assert_is_nonnegative_uint32, - _assert_is_positive_uint32, - _castable_dtype_or_raise, - _ensure_index_metadata, - _valid_index_prefix, - _valid_metric, -) - -__ALL__ = ["StaticDiskIndex"] - - -class StaticDiskIndex: - """ - A StaticDiskIndex is a disk-backed index that is not mutable. - """ - - def __init__( - self, - index_directory: str, - num_threads: int, - num_nodes_to_cache: int, - cache_mechanism: int = 1, - distance_metric: Optional[DistanceMetric] = None, - vector_dtype: Optional[VectorDType] = None, - dimensions: Optional[int] = None, - index_prefix: str = "ann", - pq_prefix: str = "", - partition_prefix: str = "", - ): - """ - ### Parameters - - **index_directory**: The directory containing the index files. This directory must contain the following - files: - - `{index_prefix}_sample_data.bin` - - `{index_prefix}_mem.index.data` - - `{index_prefix}_pq_compressed.bin` - - `{index_prefix}_pq_pivots.bin` - - `{index_prefix}_sample_ids.bin` - - `{index_prefix}_disk.index` - - It may also include the following optional files: - - `{index_prefix}_vectors.bin`: Optional. `diskannpy` builder functions may create this file in the - `index_directory` if the index was created from a numpy array - - `{index_prefix}_metadata.bin`: Optional. `diskannpy` builder functions create this file to store metadata - about the index, such as vector dtype, distance metric, number of vectors and vector dimensionality. - If an index is built from the `diskann` cli tools, this file will not exist. - - **num_threads**: Number of threads to use when searching this index. (>= 0), 0 = num_threads in system - - **num_nodes_to_cache**: Number of nodes to cache in memory (> -1) - - **cache_mechanism**: 1 -> use the generated sample_data.bin file for - the index to initialize a set of cached nodes, up to `num_nodes_to_cache`, 2 -> ready the cache for up to - `num_nodes_to_cache`, but do not initialize it with any nodes. Any other value disables node caching. - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. Default is `None`. **This - value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, - you are required to provide it. - - **vector_dtype**: The vector dtype this index has been built with. **This value is only used if a - `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, you are required to provide it. - - **dimensions**: The vector dimensionality of this index. All new vectors inserted must be the same - dimensionality. **This value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it - does not exist, you are required to provide it. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - """ - index_prefix_path = _valid_index_prefix(index_directory, index_prefix) - vector_dtype, metric, _, _ = _ensure_index_metadata( - index_prefix_path, - vector_dtype, - distance_metric, - 1, # it doesn't matter because we don't need it in this context anyway - dimensions, - ) - dap_metric = _valid_metric(metric) - - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_nonnegative_uint32(num_nodes_to_cache, "num_nodes_to_cache") - - self._vector_dtype = vector_dtype - if vector_dtype == np.uint8: - _index = _native_dap.StaticDiskUInt8Index - elif vector_dtype == np.int8: - _index = _native_dap.StaticDiskInt8Index - else: - _index = _native_dap.StaticDiskFloatIndex - self._index = _index( - distance_metric=dap_metric, - index_path_prefix=index_prefix_path, - num_threads=num_threads, - num_nodes_to_cache=num_nodes_to_cache, - cache_mechanism=cache_mechanism, - pq_prefix=pq_prefix, - partition_prefix=partition_prefix, - ) - print("After index init") - - def search( - self, - query: VectorLike, - k_neighbors: int, - complexity: int, - beam_width: int = 2, - USE_DEFERRED_FETCH: bool = False, - skip_search_reorder: bool = False, - recompute_beighbor_embeddings: bool = False, - dedup_node_dis: bool = False, - prune_ratio: float = 0, - batch_recompute: bool = False, - global_pruning: bool = False, - ) -> QueryResponse: - """ - Searches the index by a single query vector. - - ### Parameters - - **query**: 1d numpy array of the same dimensionality and dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - - **beam_width**: The beamwidth to be used for search. This is the maximum number of IO requests each query - will issue per iteration of search code. Larger beamwidth will result in fewer IO round-trips per query, - but might result in slightly higher total number of IO requests to SSD per query. For the highest query - throughput with a fixed SSD IOps rating, use W=1. For best latency, use W=4,8 or higher complexity search. - Specifying 0 will optimize the beamwidth depending on the number of threads performing search, but will - involve some tuning overhead. - - **skip_search_reorder**: Whether to skip search reorder for diskann search. - - **recompute_beighbor_embeddings**: Whether to recompute the neighbor embeddings. - - **dedup_node_dis**: Whether to dedup node distances. - - **batch_recompute**: Whether to batch recompute. - """ - _query = _castable_dtype_or_raise(query, expected=self._vector_dtype) - _assert(len(_query.shape) == 1, "query vector must be 1-d") - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_positive_uint32(beam_width, "beam_width") - - if k_neighbors > complexity: - warnings.warn( - f"{k_neighbors=} asked for, but {complexity=} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - - neighbors, distances = self._index.search( - query=_query, - knn=k_neighbors, - complexity=complexity, - beam_width=beam_width, - USE_DEFERRED_FETCH=USE_DEFERRED_FETCH, - skip_search_reorder=skip_search_reorder, - recompute_beighbor_embeddings=recompute_beighbor_embeddings, - dedup_node_dis=dedup_node_dis, - prune_ratio=prune_ratio, - batch_recompute=batch_recompute, - global_pruning=global_pruning, - ) - return QueryResponse(identifiers=neighbors, distances=distances) - - def batch_search( - self, - queries: VectorLikeBatch, - k_neighbors: int, - complexity: int, - num_threads: int, - beam_width: int = 2, - USE_DEFERRED_FETCH: bool = False, - skip_search_reorder: bool = False, - recompute_beighbor_embeddings: bool = False, - dedup_node_dis: bool = False, - prune_ratio: float = 0, - batch_recompute: bool = False, - global_pruning: bool = False, - ) -> QueryResponseBatch: - """ - Searches the index by a batch of query vectors. - - This search is parallelized and far more efficient than searching for each vector individually. - - ### Parameters - - **queries**: 2d numpy array, with column dimensionality matching the index and row dimensionality being the - number of queries intended to search for in parallel. Dtype must match dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - - **num_threads**: Number of threads to use when searching this index. (>= 0), 0 = num_threads in system - - **beam_width**: The beamwidth to be used for search. This is the maximum number of IO requests each query - will issue per iteration of search code. Larger beamwidth will result in fewer IO round-trips per query, - but might result in slightly higher total number of IO requests to SSD per query. For the highest query - throughput with a fixed SSD IOps rating, use W=1. For best latency, use W=4,8 or higher complexity search. - Specifying 0 will optimize the beamwidth depending on the number of threads performing search, but will - involve some tuning overhead. - - **skip_search_reorder**: Whether to skip search reorder for diskann search. - """ - _queries = _castable_dtype_or_raise(queries, expected=self._vector_dtype) - _assert_2d(_queries, "queries") - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_positive_uint32(beam_width, "beam_width") - - if k_neighbors > complexity: - warnings.warn( - f"{k_neighbors=} asked for, but {complexity=} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - - num_queries, dim = _queries.shape - print( - f"USE_DEFERRED_FETCH={USE_DEFERRED_FETCH} skip_search_reorder={skip_search_reorder} recompute_beighbor_embeddings={recompute_beighbor_embeddings}, dedup_node_dis={dedup_node_dis}" - ) - neighbors, distances = self._index.batch_search( - queries=_queries, - num_queries=num_queries, - knn=k_neighbors, - complexity=complexity, - beam_width=beam_width, - num_threads=num_threads, - USE_DEFERRED_FETCH=USE_DEFERRED_FETCH, - skip_search_reorder=skip_search_reorder, - recompute_beighbor_embeddings=recompute_beighbor_embeddings, - dedup_node_dis=dedup_node_dis, - prune_ratio=prune_ratio, - batch_recompute=batch_recompute, - global_pruning=global_pruning, - ) - return QueryResponseBatch(identifiers=neighbors, distances=distances) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_memory_index.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_memory_index.py deleted file mode 100644 index 1380360..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/_static_memory_index.py +++ /dev/null @@ -1,262 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import json -import os -import warnings -from typing import Optional - -import numpy as np - -from . import ( - DistanceMetric, - QueryResponse, - QueryResponseBatch, - VectorDType, - VectorLike, - VectorLikeBatch, -) -from . import _diskannpy as _native_dap -from ._common import ( - _assert, - _assert_is_nonnegative_uint32, - _assert_is_positive_uint32, - _castable_dtype_or_raise, - _ensure_index_metadata, - _valid_index_prefix, - _valid_metric, -) - -__ALL__ = ["StaticMemoryIndex"] - - -class StaticMemoryIndex: - """ - A StaticMemoryIndex is an immutable in-memory DiskANN index. - """ - - def __init__( - self, - index_directory: str, - num_threads: int, - initial_search_complexity: int, - index_prefix: str = "ann", - distance_metric: Optional[DistanceMetric] = None, - vector_dtype: Optional[VectorDType] = None, - dimensions: Optional[int] = None, - enable_filters: bool = False, - ): - """ - ### Parameters - - **index_directory**: The directory containing the index files. This directory must contain the following - files: - - `{index_prefix}.data` - - `{index_prefix}` - - - It may also include the following optional files: - - `{index_prefix}_vectors.bin`: Optional. `diskannpy` builder functions may create this file in the - `index_directory` if the index was created from a numpy array - - `{index_prefix}_metadata.bin`: Optional. `diskannpy` builder functions create this file to store metadata - about the index, such as vector dtype, distance metric, number of vectors and vector dimensionality. - If an index is built from the `diskann` cli tools, this file will not exist. - - **num_threads**: Number of threads to use when searching this index. (>= 0), 0 = num_threads in system - - **initial_search_complexity**: Should be set to the most common `complexity` expected to be used during the - life of this `diskannpy.DynamicMemoryIndex` object. The working scratch memory allocated is based off of - `initial_search_complexity` * `search_threads`. Note that it may be resized if a `search` or `batch_search` - operation requests a space larger than can be accommodated by these values. - - **index_prefix**: The prefix of the index files. Defaults to "ann". - - **distance_metric**: A `str`, strictly one of {"l2", "mips", "cosine"}. `l2` and `cosine` are supported for all 3 - vector dtypes, but `mips` is only available for single precision floats. Default is `None`. **This - value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, - you are required to provide it. - - **vector_dtype**: The vector dtype this index has been built with. **This value is only used if a - `{index_prefix}_metadata.bin` file does not exist.** If it does not exist, you are required to provide it. - - **dimensions**: The vector dimensionality of this index. All new vectors inserted must be the same - dimensionality. **This value is only used if a `{index_prefix}_metadata.bin` file does not exist.** If it - does not exist, you are required to provide it. - - **enable_filters**: Indexes built with filters can also be used for filtered search. - """ - index_prefix_path = _valid_index_prefix(index_directory, index_prefix) - self._labels_map = {} - self._labels_metadata = {} - if enable_filters: - try: - with open(f"{index_prefix_path}_labels_map.txt", "r") as labels_map_if: - for line in labels_map_if: - (key, val) = line.split("\t") - self._labels_map[key] = int(val) - with open( - f"{index_prefix_path}_label_metadata.json", "r" - ) as labels_metadata_if: - self._labels_metadata = json.load(labels_metadata_if) - except: # noqa: E722 - # exceptions are basically presumed to be either file not found or file not formatted correctly - raise RuntimeException("Filter labels file was unable to be processed.") - vector_dtype, metric, num_points, dims = _ensure_index_metadata( - index_prefix_path, - vector_dtype, - distance_metric, - 1, # it doesn't matter because we don't need it in this context anyway - dimensions, - ) - dap_metric = _valid_metric(metric) - - _assert_is_nonnegative_uint32(num_threads, "num_threads") - _assert_is_positive_uint32( - initial_search_complexity, "initial_search_complexity" - ) - - self._vector_dtype = vector_dtype - self._dimensions = dims - - if vector_dtype == np.uint8: - _index = _native_dap.StaticMemoryUInt8Index - elif vector_dtype == np.int8: - _index = _native_dap.StaticMemoryInt8Index - else: - _index = _native_dap.StaticMemoryFloatIndex - - self._index = _index( - distance_metric=dap_metric, - num_points=num_points, - dimensions=dims, - index_path=index_prefix_path, - num_threads=num_threads, - initial_search_complexity=initial_search_complexity, - ) - - def search( - self, - query: VectorLike, - k_neighbors: int, - complexity: int, - filter_label: str = "", - USE_DEFERRED_FETCH: bool = False, - skip_search_reorder: bool = False, - recompute_beighbor_embeddings: bool = False, - dedup_node_dis: bool = False, - prune_ratio: float = 0, - batch_recompute: bool = False, - global_pruning: bool = False, - ) -> QueryResponse: - """ - Searches the index by a single query vector. - - ### Parameters - - **query**: 1d numpy array of the same dimensionality and dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - """ - if filter_label != "": - if len(self._labels_map) == 0: - raise ValueError( - f"A filter label of {filter_label} was provided, but this class was not initialized with filters " - "enabled, e.g. StaticDiskMemory(..., enable_filters=True)" - ) - if filter_label not in self._labels_map: - raise ValueError( - f"A filter label of {filter_label} was provided, but the external(str)->internal(np.uint32) labels map " - f"does not include that label." - ) - k_neighbors = min(k_neighbors, self._labels_metadata[filter_label]) - _query = _castable_dtype_or_raise(query, expected=self._vector_dtype) - _assert(len(_query.shape) == 1, "query vector must be 1-d") - _assert( - _query.shape[0] == self._dimensions, - f"query vector must have the same dimensionality as the index; index dimensionality: {self._dimensions}, " - f"query dimensionality: {_query.shape[0]}", - ) - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_nonnegative_uint32(complexity, "complexity") - - if k_neighbors > complexity: - warnings.warn( - f"k_neighbors={k_neighbors} asked for, but list_size={complexity} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - - if filter_label == "": - neighbors, distances = self._index.search( - query=_query, - knn=k_neighbors, - complexity=complexity, - USE_DEFERRED_FETCH=USE_DEFERRED_FETCH, - skip_search_reorder=skip_search_reorder, - recompute_beighbor_embeddings=recompute_beighbor_embeddings, - dedup_node_dis=dedup_node_dis, - prune_ratio=prune_ratio, - batch_recompute=batch_recompute, - global_pruning=global_pruning, - ) - else: - filter = self._labels_map[filter_label] - neighbors, distances = self._index.search_with_filter( - query=query, knn=k_neighbors, complexity=complexity, filter=filter - ) - return QueryResponse(identifiers=neighbors, distances=distances) - - def batch_search( - self, - queries: VectorLikeBatch, - k_neighbors: int, - complexity: int, - num_threads: int, - USE_DEFERRED_FETCH: bool = False, - skip_search_reorder: bool = False, - recompute_beighbor_embeddings: bool = False, - dedup_node_dis: bool = False, - prune_ratio: float = 0, - batch_recompute: bool = False, - global_pruning: bool = False, - ) -> QueryResponseBatch: - """ - Searches the index by a batch of query vectors. - - This search is parallelized and far more efficient than searching for each vector individually. - - ### Parameters - - **queries**: 2d numpy array, with column dimensionality matching the index and row dimensionality being the - number of queries intended to search for in parallel. Dtype must match dtype of the index. - - **k_neighbors**: Number of neighbors to be returned. If query vector exists in index, it almost definitely - will be returned as well, so adjust your ``k_neighbors`` as appropriate. Must be > 0. - - **complexity**: Size of distance ordered list of candidate neighbors to use while searching. List size - increases accuracy at the cost of latency. Must be at least k_neighbors in size. - - **num_threads**: Number of threads to use when searching this index. (>= 0), 0 = num_threads in system - """ - - _queries = _castable_dtype_or_raise(queries, expected=self._vector_dtype) - _assert(len(_queries.shape) == 2, "queries must must be 2-d np array") - _assert( - _queries.shape[1] == self._dimensions, - f"query vectors must have the same dimensionality as the index; index dimensionality: {self._dimensions}, " - f"query dimensionality: {_queries.shape[1]}", - ) - _assert_is_positive_uint32(k_neighbors, "k_neighbors") - _assert_is_positive_uint32(complexity, "complexity") - _assert_is_nonnegative_uint32(num_threads, "num_threads") - - if k_neighbors > complexity: - warnings.warn( - f"k_neighbors={k_neighbors} asked for, but list_size={complexity} was smaller. Increasing {complexity} to {k_neighbors}" - ) - complexity = k_neighbors - - num_queries, dim = _queries.shape - neighbors, distances = self._index.batch_search( - queries=_queries, - num_queries=num_queries, - knn=k_neighbors, - complexity=complexity, - num_threads=num_threads, - USE_DEFERRED_FETCH=USE_DEFERRED_FETCH, - skip_search_reorder=skip_search_reorder, - recompute_beighbor_embeddings=recompute_beighbor_embeddings, - dedup_node_dis=dedup_node_dis, - prune_ratio=prune_ratio, - batch_recompute=batch_recompute, - global_pruning=global_pruning, - ) - return QueryResponseBatch(identifiers=neighbors, distances=distances) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/builder.cpp b/packages/leann-backend-diskann/third_party/DiskANN/python/src/builder.cpp deleted file mode 100644 index 2b91eac..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/builder.cpp +++ /dev/null @@ -1,136 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "builder.h" -#include "common.h" -#include "disk_utils.h" -#include "index.h" -#include "parameters.h" - -namespace diskannpy -{ -template -void build_disk_index(const diskann::Metric metric, const std::string &data_file_path, - const std::string &index_prefix_path, const uint32_t complexity, const uint32_t graph_degree, - const double final_index_ram_limit, const double indexing_ram_budget, const uint32_t num_threads, - const uint32_t pq_disk_bytes, const std::string &codebook_prefix) -{ - std::string params = std::to_string(graph_degree) + " " + std::to_string(complexity) + " " + - std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + " " + - std::to_string(num_threads); - if (pq_disk_bytes > 0) - params = params + " " + std::to_string(pq_disk_bytes); - if (!codebook_prefix.empty()) - params = params + " " + codebook_prefix; - diskann::build_disk_index
(data_file_path.c_str(), index_prefix_path.c_str(), params.c_str(), metric, false, - codebook_prefix); -} - -template void build_disk_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - double, double, uint32_t, uint32_t, const std::string &); - -template void build_disk_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - double, double, uint32_t, uint32_t, const std::string &); -template void build_disk_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - double, double, uint32_t, uint32_t, const std::string &); - -template -std::string prepare_filtered_label_map(diskann::Index &index, const std::string &index_output_path, - const std::string &filter_labels_file, const std::string &universal_label) -{ - std::string labels_file_to_use = index_output_path + "_label_formatted.txt"; - std::string mem_labels_int_map_file = index_output_path + "_labels_map.txt"; - convert_labels_string_to_int(filter_labels_file, labels_file_to_use, mem_labels_int_map_file, universal_label); - if (!universal_label.empty()) - { - uint32_t unv_label_as_num = 0; - index.set_universal_label(unv_label_as_num); - } - return labels_file_to_use; -} - -template std::string prepare_filtered_label_map(diskann::Index &, const std::string &, - const std::string &, const std::string &); - -template std::string prepare_filtered_label_map(diskann::Index &, - const std::string &, const std::string &, const std::string &); - -template std::string prepare_filtered_label_map(diskann::Index &, - const std::string &, const std::string &, const std::string &); - -template -void build_memory_index(const diskann::Metric metric, const std::string &vector_bin_path, - const std::string &index_output_path, const uint32_t graph_degree, const uint32_t complexity, - const float alpha, const uint32_t num_threads, const bool use_pq_build, - const size_t num_pq_bytes, const bool use_opq, const bool use_tags, - const std::string &filter_labels_file, const std::string &universal_label, - const uint32_t filter_complexity) -{ - diskann::IndexWriteParameters index_build_params = diskann::IndexWriteParametersBuilder(complexity, graph_degree) - .with_filter_list_size(filter_complexity) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - diskann::IndexSearchParams index_search_params = - diskann::IndexSearchParams(index_build_params.search_list_size, num_threads); - size_t data_num, data_dim; - diskann::get_bin_metadata(vector_bin_path, data_num, data_dim); - - diskann::Index index(metric, data_dim, data_num, - std::make_shared(index_build_params), - std::make_shared(index_search_params), 0, - use_tags, use_tags, false, use_pq_build, num_pq_bytes, use_opq); - - if (use_tags) - { - const std::string tags_file = index_output_path + ".tags"; - if (!file_exists(tags_file)) - { - throw std::runtime_error("tags file not found at expected path: " + tags_file); - } - TagT *tags_data; - size_t tag_dims = 1; - diskann::load_bin(tags_file, tags_data, data_num, tag_dims); - std::vector tags(tags_data, tags_data + data_num); - if (filter_labels_file.empty()) - { - index.build(vector_bin_path.c_str(), data_num, tags); - } - else - { - auto labels_file = prepare_filtered_label_map(index, index_output_path, filter_labels_file, - universal_label); - index.build_filtered_index(vector_bin_path.c_str(), labels_file, data_num, tags); - } - } - else - { - if (filter_labels_file.empty()) - { - index.build(vector_bin_path.c_str(), data_num); - } - else - { - auto labels_file = prepare_filtered_label_map(index, index_output_path, filter_labels_file, - universal_label); - index.build_filtered_index(vector_bin_path.c_str(), labels_file, data_num); - } - } - - index.save(index_output_path.c_str()); -} - -template void build_memory_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - float, uint32_t, bool, size_t, bool, bool, const std::string &, - const std::string &, uint32_t); - -template void build_memory_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - float, uint32_t, bool, size_t, bool, bool, const std::string &, - const std::string &, uint32_t); - -template void build_memory_index(diskann::Metric, const std::string &, const std::string &, uint32_t, uint32_t, - float, uint32_t, bool, size_t, bool, bool, const std::string &, - const std::string &, uint32_t); - -} // namespace diskannpy diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/defaults.py b/packages/leann-backend-diskann/third_party/DiskANN/python/src/defaults.py deleted file mode 100644 index 4e22983..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/defaults.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -""" -# Parameter Defaults -These parameter defaults are re-exported from the C++ extension module, and used to keep the pythonic wrapper in sync with the C++. -""" -from ._diskannpy import defaults as _defaults - -ALPHA = _defaults.ALPHA -""" -Note that, as ALPHA is a `float32` (single precision float) in C++, when converted into Python it becomes a -`float64` (double precision float). The actual value is 1.2f. The alpha parameter (>=1) is used to control the nature -and number of points that are added to the graph. A higher alpha value (e.g., 1.4) will result in fewer hops (and IOs) -to convergence, but probably more distance comparisons compared to a lower alpha value. -""" -NUM_THREADS = _defaults.NUM_THREADS -""" Number of threads to use. `0` will use all available detected logical processors """ -MAX_OCCLUSION_SIZE = _defaults.MAX_OCCLUSION_SIZE -""" -The maximum number of points that can be occluded by a single point. This is used to prevent a single point from -dominating the graph structure. If a point has more than `max_occlusion_size` neighbors closer to it than the current -point, it will not be added to the graph. This is a tradeoff between index build time and search quality. -""" -FILTER_COMPLEXITY = _defaults.FILTER_COMPLEXITY -""" -Complexity (a.k.a. `L`) references the size of the list we store candidate approximate neighbors in while doing a -filtered search. This value must be larger than `k_neighbors`, and larger values tend toward higher recall in the -resultant ANN search at the cost of more time. -""" -NUM_FROZEN_POINTS_STATIC = _defaults.NUM_FROZEN_POINTS_STATIC -""" Number of points frozen by default in a StaticMemoryIndex """ -NUM_FROZEN_POINTS_DYNAMIC = _defaults.NUM_FROZEN_POINTS_DYNAMIC -""" Number of points frozen by default in a DynamicMemoryIndex """ -SATURATE_GRAPH = _defaults.SATURATE_GRAPH -""" Whether to saturate the graph or not. Default is `True` """ -GRAPH_DEGREE = _defaults.GRAPH_DEGREE -""" -Graph degree (a.k.a. `R`) is the maximum degree allowed for a node in the index's graph structure. This degree will be -pruned throughout the course of the index build, but it will never grow beyond this value. Higher R values require -longer index build times, but may result in an index showing excellent recall and latency characteristics. -""" -COMPLEXITY = _defaults.COMPLEXITY -""" -Complexity (a.k.a `L`) references the size of the list we store candidate approximate neighbors in while doing build -or search tasks. It's used during index build as part of the index optimization processes. It's used in index search -classes both to help mitigate poor latencies during cold start, as well as on subsequent queries to conduct the search. -Large values will likely increase latency but also may improve recall, and tuning these values for your particular -index is certainly a reasonable choice. -""" -PQ_DISK_BYTES = _defaults.PQ_DISK_BYTES -""" -Use `0` to store uncompressed data on SSD. This allows the index to asymptote to 100% recall. If your vectors are -too large to store in SSD, this parameter provides the option to compress the vectors using PQ for storing on SSD. -This will trade off recall. You would also want this to be greater than the number of bytes used for the PQ -compressed data stored in-memory. Default is `0`. -""" -USE_PQ_BUILD = _defaults.USE_PQ_BUILD -""" - Whether to use product quantization in the index building process. Product quantization is an approximation -technique that can vastly speed up vector computations and comparisons in a spatial neighborhood, but it is still an -approximation technique. It should be preferred when index creation times take longer than you can afford for your -use case. -""" -NUM_PQ_BYTES = _defaults.NUM_PQ_BYTES -""" -The number of product quantization bytes to use. More bytes requires more resources in both memory and time, but is -like to result in better approximations. -""" -USE_OPQ = _defaults.USE_OPQ -""" Whether to use Optimized Product Quantization or not. """ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/dynamic_memory_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/python/src/dynamic_memory_index.cpp deleted file mode 100644 index d05e54d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/dynamic_memory_index.cpp +++ /dev/null @@ -1,167 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "parameters.h" -#include "dynamic_memory_index.h" - -#include "pybind11/numpy.h" - -namespace diskannpy -{ - -diskann::IndexWriteParameters dynamic_index_write_parameters(const uint32_t complexity, const uint32_t graph_degree, - const bool saturate_graph, - const uint32_t max_occlusion_size, const float alpha, - const uint32_t num_threads, - const uint32_t filter_complexity) -{ - return diskann::IndexWriteParametersBuilder(complexity, graph_degree) - .with_saturate_graph(saturate_graph) - .with_max_occlusion_size(max_occlusion_size) - .with_alpha(alpha) - .with_num_threads(num_threads) - .with_filter_list_size(filter_complexity) - .build(); -} - -template -diskann::Index dynamic_index_builder( - const diskann::Metric m, const diskann::IndexWriteParameters &write_params, const size_t dimensions, - const size_t max_vectors, const uint32_t initial_search_complexity, const uint32_t initial_search_threads, - const bool concurrent_consolidation, const uint32_t num_frozen_points) -{ - const uint32_t _initial_search_threads = initial_search_threads != 0 ? initial_search_threads : omp_get_num_procs(); - - auto index_search_params = diskann::IndexSearchParams(initial_search_complexity, _initial_search_threads); - return diskann::Index( - m, dimensions, max_vectors, - std::make_shared(write_params), // index write params - std::make_shared(index_search_params), // index_search_params - num_frozen_points, // frozen_points - true, // dynamic_index - true, // enable_tags - concurrent_consolidation, - false, // pq_dist_build - 0, // num_pq_chunks - false); // use_opq = false -} - -template -DynamicMemoryIndex
::DynamicMemoryIndex(const diskann::Metric m, const size_t dimensions, const size_t max_vectors, - const uint32_t complexity, const uint32_t graph_degree, - const bool saturate_graph, const uint32_t max_occlusion_size, - const float alpha, const uint32_t num_threads, - const uint32_t filter_complexity, const uint32_t num_frozen_points, - const uint32_t initial_search_complexity, - const uint32_t initial_search_threads, const bool concurrent_consolidation) - : _initial_search_complexity(initial_search_complexity != 0 ? initial_search_complexity : complexity), - _write_parameters(dynamic_index_write_parameters(complexity, graph_degree, saturate_graph, max_occlusion_size, - alpha, num_threads, filter_complexity)), - _index(dynamic_index_builder
(m, _write_parameters, dimensions, max_vectors, _initial_search_complexity, - initial_search_threads, concurrent_consolidation, num_frozen_points)) -{ -} - -template void DynamicMemoryIndex
::load(const std::string &index_path) -{ - const std::string tags_file = index_path + ".tags"; - if (!file_exists(tags_file)) - { - throw std::runtime_error("tags file not found at expected path: " + tags_file); - } - _index.load(index_path.c_str(), _write_parameters.num_threads, _initial_search_complexity); -} - -template -int DynamicMemoryIndex
::insert(const py::array_t &vector, - const DynamicIdType id) -{ - return _index.insert_point(vector.data(), id); -} - -template -py::array_t DynamicMemoryIndex
::batch_insert( - py::array_t &vectors, - py::array_t &ids, const int32_t num_inserts, - const int num_threads) -{ - if (num_threads == 0) - omp_set_num_threads(omp_get_num_procs()); - else - omp_set_num_threads(num_threads); - py::array_t insert_retvals(num_inserts); - -#pragma omp parallel for schedule(dynamic, 1) default(none) shared(num_inserts, insert_retvals, vectors, ids) - for (int32_t i = 0; i < num_inserts; i++) - { - insert_retvals.mutable_data()[i] = _index.insert_point(vectors.data(i), *(ids.data(i))); - } - - return insert_retvals; -} - -template int DynamicMemoryIndex
::mark_deleted(const DynamicIdType id) -{ - return this->_index.lazy_delete(id); -} - -template void DynamicMemoryIndex
::save(const std::string &save_path, const bool compact_before_save) -{ - if (save_path.empty()) - { - throw std::runtime_error("A save_path must be provided"); - } - _index.save(save_path.c_str(), compact_before_save); -} - -template -NeighborsAndDistances DynamicMemoryIndex
::search( - py::array_t &query, const uint64_t knn, const uint64_t complexity) -{ - py::array_t ids(knn); - py::array_t dists(knn); - std::vector
empty_vector; - _index.search_with_tags(query.data(), knn, complexity, ids.mutable_data(), dists.mutable_data(), empty_vector); - return std::make_pair(ids, dists); -} - -template -NeighborsAndDistances DynamicMemoryIndex
::batch_search( - py::array_t &queries, const uint64_t num_queries, const uint64_t knn, - const uint64_t complexity, const uint32_t num_threads) -{ - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector
empty_vector; - - if (num_threads == 0) - omp_set_num_threads(omp_get_num_procs()); - else - omp_set_num_threads(static_cast(num_threads)); - -#pragma omp parallel for schedule(dynamic, 1) default(none) \ - shared(num_queries, queries, knn, complexity, ids, dists, empty_vector) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - _index.search_with_tags(queries.data(i), knn, complexity, ids.mutable_data(i), dists.mutable_data(i), - empty_vector); - } - - return std::make_pair(ids, dists); -} - -template void DynamicMemoryIndex
::consolidate_delete() -{ - _index.consolidate_deletes(_write_parameters); -} - -template size_t DynamicMemoryIndex
::num_points() -{ - return _index.get_num_points(); -} - -template class DynamicMemoryIndex; -template class DynamicMemoryIndex; -template class DynamicMemoryIndex; - -}; // namespace diskannpy diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/module.cpp b/packages/leann-backend-diskann/third_party/DiskANN/python/src/module.cpp deleted file mode 100644 index 0f295e1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/module.cpp +++ /dev/null @@ -1,142 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include - -#include -#include - -#include "defaults.h" -#include "distance.h" - -#include "builder.h" -#include "dynamic_memory_index.h" -#include "static_disk_index.h" -#include "static_memory_index.h" - -PYBIND11_MAKE_OPAQUE(std::vector); -PYBIND11_MAKE_OPAQUE(std::vector); -PYBIND11_MAKE_OPAQUE(std::vector); -PYBIND11_MAKE_OPAQUE(std::vector); - -namespace py = pybind11; -using namespace pybind11::literals; - -struct Variant -{ - std::string disk_builder_name; - std::string memory_builder_name; - std::string dynamic_memory_index_name; - std::string static_memory_index_name; - std::string static_disk_index_name; -}; - -const Variant FloatVariant{"build_disk_float_index", "build_memory_float_index", "DynamicMemoryFloatIndex", - "StaticMemoryFloatIndex", "StaticDiskFloatIndex"}; - -const Variant UInt8Variant{"build_disk_uint8_index", "build_memory_uint8_index", "DynamicMemoryUInt8Index", - "StaticMemoryUInt8Index", "StaticDiskUInt8Index"}; - -const Variant Int8Variant{"build_disk_int8_index", "build_memory_int8_index", "DynamicMemoryInt8Index", - "StaticMemoryInt8Index", "StaticDiskInt8Index"}; - -template inline void add_variant(py::module_ &m, const Variant &variant) -{ - m.def(variant.disk_builder_name.c_str(), &diskannpy::build_disk_index, "distance_metric"_a, "data_file_path"_a, - "index_prefix_path"_a, "complexity"_a, "graph_degree"_a, "final_index_ram_limit"_a, "indexing_ram_budget"_a, - "num_threads"_a, "pq_disk_bytes"_a, "codebook_prefix"_a = ""); - - m.def(variant.memory_builder_name.c_str(), &diskannpy::build_memory_index, "distance_metric"_a, - "data_file_path"_a, "index_output_path"_a, "graph_degree"_a, "complexity"_a, "alpha"_a, "num_threads"_a, - "use_pq_build"_a, "num_pq_bytes"_a, "use_opq"_a, "use_tags"_a = false, "filter_labels_file"_a = "", - "universal_label"_a = "", "filter_complexity"_a = 0); - - py::class_>(m, variant.static_memory_index_name.c_str()) - .def(py::init(), - "distance_metric"_a, "index_path"_a, "num_points"_a, "dimensions"_a, "num_threads"_a, - "initial_search_complexity"_a) - .def("search", &diskannpy::StaticMemoryIndex::search, "query"_a, "knn"_a, "complexity"_a) - .def("search_with_filter", &diskannpy::StaticMemoryIndex::search_with_filter, "query"_a, "knn"_a, - "complexity"_a, "filter"_a) - .def("batch_search", &diskannpy::StaticMemoryIndex::batch_search, "queries"_a, "num_queries"_a, "knn"_a, - "complexity"_a, "num_threads"_a); - - py::class_>(m, variant.dynamic_memory_index_name.c_str()) - .def(py::init(), - "distance_metric"_a, "dimensions"_a, "max_vectors"_a, "complexity"_a, "graph_degree"_a, - "saturate_graph"_a = diskann::defaults::SATURATE_GRAPH, - "max_occlusion_size"_a = diskann::defaults::MAX_OCCLUSION_SIZE, "alpha"_a = diskann::defaults::ALPHA, - "num_threads"_a = diskann::defaults::NUM_THREADS, - "filter_complexity"_a = diskann::defaults::FILTER_LIST_SIZE, - "num_frozen_points"_a = diskann::defaults::NUM_FROZEN_POINTS_DYNAMIC, "initial_search_complexity"_a = 0, - "search_threads"_a = 0, "concurrent_consolidation"_a = true) - .def("search", &diskannpy::DynamicMemoryIndex::search, "query"_a, "knn"_a, "complexity"_a) - .def("load", &diskannpy::DynamicMemoryIndex::load, "index_path"_a) - .def("batch_search", &diskannpy::DynamicMemoryIndex::batch_search, "queries"_a, "num_queries"_a, "knn"_a, - "complexity"_a, "num_threads"_a) - .def("batch_insert", &diskannpy::DynamicMemoryIndex::batch_insert, "vectors"_a, "ids"_a, "num_inserts"_a, - "num_threads"_a) - .def("save", &diskannpy::DynamicMemoryIndex::save, "save_path"_a = "", "compact_before_save"_a = false) - .def("insert", &diskannpy::DynamicMemoryIndex::insert, "vector"_a, "id"_a) - .def("mark_deleted", &diskannpy::DynamicMemoryIndex::mark_deleted, "id"_a) - .def("consolidate_delete", &diskannpy::DynamicMemoryIndex::consolidate_delete) - .def("num_points", &diskannpy::DynamicMemoryIndex::num_points); - - py::class_>(m, variant.static_disk_index_name.c_str()) - .def(py::init(), - "distance_metric"_a, "index_path_prefix"_a, "num_threads"_a, "num_nodes_to_cache"_a, - "cache_mechanism"_a = 1, "pq_prefix"_a = "", "partition_prefix"_a) - .def("cache_bfs_levels", &diskannpy::StaticDiskIndex::cache_bfs_levels, "num_nodes_to_cache"_a) - .def("search", &diskannpy::StaticDiskIndex::search, "query"_a, "knn"_a, "complexity"_a, "beam_width"_a, - "USE_DEFERRED_FETCH"_a = false, "skip_search_reorder"_a = false, "recompute_beighbor_embeddings"_a = false, - "dedup_node_dis"_a = false, "prune_ratio"_a = 0, "batch_recompute"_a = false, "global_pruning"_a = false) - .def("batch_search", &diskannpy::StaticDiskIndex::batch_search, "queries"_a, "num_queries"_a, "knn"_a, - "complexity"_a, "beam_width"_a, "num_threads"_a, "USE_DEFERRED_FETCH"_a = false, - "skip_search_reorder"_a = false, "recompute_beighbor_embeddings"_a = false, "dedup_node_dis"_a = false, - "prune_ratio"_a = 0, "batch_recompute"_a = false, "global_pruning"_a = false); -} - -PYBIND11_MODULE(_diskannpy, m) -{ - m.doc() = "DiskANN Python Bindings"; -#ifdef VERSION_INFO - m.attr("__version__") = VERSION_INFO; -#else - m.attr("__version__") = "dev"; -#endif - - // let's re-export our defaults - py::module_ default_values = m.def_submodule( - "defaults", - "A collection of the default values used for common diskann operations. `GRAPH_DEGREE` and `COMPLEXITY` are not" - " set as defaults, but some semi-reasonable default values are selected for your convenience. We urge you to " - "investigate their meaning and adjust them for your use cases."); - - default_values.attr("ALPHA") = diskann::defaults::ALPHA; - default_values.attr("NUM_THREADS") = diskann::defaults::NUM_THREADS; - default_values.attr("MAX_OCCLUSION_SIZE") = diskann::defaults::MAX_OCCLUSION_SIZE; - default_values.attr("FILTER_COMPLEXITY") = diskann::defaults::FILTER_LIST_SIZE; - default_values.attr("NUM_FROZEN_POINTS_STATIC") = diskann::defaults::NUM_FROZEN_POINTS_STATIC; - default_values.attr("NUM_FROZEN_POINTS_DYNAMIC") = diskann::defaults::NUM_FROZEN_POINTS_DYNAMIC; - default_values.attr("SATURATE_GRAPH") = diskann::defaults::SATURATE_GRAPH; - default_values.attr("GRAPH_DEGREE") = diskann::defaults::MAX_DEGREE; - default_values.attr("COMPLEXITY") = diskann::defaults::BUILD_LIST_SIZE; - default_values.attr("PQ_DISK_BYTES") = (uint32_t)0; - default_values.attr("USE_PQ_BUILD") = false; - default_values.attr("NUM_PQ_BYTES") = (uint32_t)0; - default_values.attr("USE_OPQ") = false; - - add_variant(m, FloatVariant); - add_variant(m, UInt8Variant); - add_variant(m, Int8Variant); - - py::enum_(m, "Metric") - .value("L2", diskann::Metric::L2) - .value("INNER_PRODUCT", diskann::Metric::INNER_PRODUCT) - .value("COSINE", diskann::Metric::COSINE) - .export_values(); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/py.typed b/packages/leann-backend-diskann/third_party/DiskANN/python/src/py.typed deleted file mode 100644 index e69de29..0000000 diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_disk_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_disk_index.cpp deleted file mode 100644 index 47dc09b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_disk_index.cpp +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "static_disk_index.h" - -#include "pybind11/numpy.h" - -namespace diskannpy -{ - -template -StaticDiskIndex
::StaticDiskIndex(const diskann::Metric metric, const std::string &index_path_prefix, - const uint32_t num_threads, const size_t num_nodes_to_cache, - const uint32_t cache_mechanism, const std::string &pq_prefix, - const std::string &partition_prefix) - : _reader(std::make_shared()), - _graph_reader(std::make_shared()), _index(_reader, _graph_reader, metric) -{ - std::cout << "Before index load" << std::endl; - - const uint32_t _num_threads = num_threads != 0 ? num_threads : omp_get_num_procs(); - int load_success = - _index.load(_num_threads, index_path_prefix.c_str(), pq_prefix.c_str(), partition_prefix.c_str()); - if (load_success != 0) - { - throw std::runtime_error("index load failed, " + index_path_prefix); - } - if (cache_mechanism == 1) - { - std::string sample_file = index_path_prefix + std::string("_sample_data.bin"); - cache_sample_paths(num_nodes_to_cache, sample_file, _num_threads); - } - else if (cache_mechanism == 2) - { - cache_bfs_levels(num_nodes_to_cache); - } - std::cout << "After index load" << std::endl; -} - -template void StaticDiskIndex
::cache_bfs_levels(const size_t num_nodes_to_cache) -{ - std::vector node_list; - _index.cache_bfs_levels(num_nodes_to_cache, node_list); - _index.load_cache_list(node_list); -} - -template -void StaticDiskIndex
::cache_sample_paths(const size_t num_nodes_to_cache, const std::string &warmup_query_file, - const uint32_t num_threads) -{ - if (!file_exists(warmup_query_file)) - { - return; - } - - std::vector node_list; - _index.generate_cache_list_from_sample_queries(warmup_query_file, 15, 4, num_nodes_to_cache, num_threads, - node_list); - _index.load_cache_list(node_list); -} - -template -NeighborsAndDistances StaticDiskIndex
::search( - py::array_t &query, const uint64_t knn, const uint64_t complexity, - const uint64_t beam_width, const bool USE_DEFERRED_FETCH, const bool skip_search_reorder, - const bool recompute_beighbor_embeddings, const bool dedup_node_dis, const float prune_ratio, - const bool batch_recompute, const bool global_pruning) -{ - py::array_t ids(knn); - py::array_t dists(knn); - - std::vector u32_ids(knn); - std::vector u64_ids(knn); - diskann::QueryStats stats; - - _index.cached_beam_search(query.data(), knn, complexity, u64_ids.data(), dists.mutable_data(), beam_width, false, - &stats, USE_DEFERRED_FETCH, skip_search_reorder, recompute_beighbor_embeddings, - dedup_node_dis, prune_ratio, batch_recompute, global_pruning); - - auto r = ids.mutable_unchecked<1>(); - for (uint64_t i = 0; i < knn; ++i) - r(i) = (unsigned)u64_ids[i]; - - return std::make_pair(ids, dists); -} - -template -NeighborsAndDistances StaticDiskIndex
::batch_search( - py::array_t &queries, const uint64_t num_queries, const uint64_t knn, - const uint64_t complexity, const uint64_t beam_width, const uint32_t num_threads, const bool USE_DEFERRED_FETCH, - const bool skip_search_reorder, const bool recompute_beighbor_embeddings, const bool dedup_node_dis, - const float prune_ratio, const bool batch_recompute, const bool global_pruning) -{ - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - - omp_set_num_threads(num_threads); - - std::vector u64_ids(knn * num_queries); - -#pragma omp parallel for schedule(dynamic, 1) default(none) \ - shared(num_queries, queries, knn, complexity, u64_ids, dists, beam_width, USE_DEFERRED_FETCH, skip_search_reorder, \ - recompute_beighbor_embeddings, dedup_node_dis, prune_ratio, batch_recompute, global_pruning) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - _index.cached_beam_search(queries.data(i), knn, complexity, u64_ids.data() + i * knn, dists.mutable_data(i), - beam_width, false, nullptr, USE_DEFERRED_FETCH, skip_search_reorder, - recompute_beighbor_embeddings, dedup_node_dis, prune_ratio, batch_recompute, - global_pruning); - } - - auto r = ids.mutable_unchecked(); - for (uint64_t i = 0; i < num_queries; ++i) - for (uint64_t j = 0; j < knn; ++j) - r(i, j) = (uint32_t)u64_ids[i * knn + j]; - - return std::make_pair(ids, dists); -} - -template class StaticDiskIndex; -template class StaticDiskIndex; -template class StaticDiskIndex; -} // namespace diskannpy \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_memory_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_memory_index.cpp deleted file mode 100644 index d3ac079..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/src/static_memory_index.cpp +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "static_memory_index.h" - -#include "pybind11/numpy.h" - -namespace diskannpy -{ - -template -diskann::Index static_index_builder(const diskann::Metric m, const size_t num_points, - const size_t dimensions, - const uint32_t initial_search_complexity) -{ - if (initial_search_complexity == 0) - { - throw std::runtime_error("initial_search_complexity must be a positive uint32_t"); - } - auto index_search_params = diskann::IndexSearchParams(initial_search_complexity, omp_get_num_procs()); - return diskann::Index
(m, dimensions, num_points, - nullptr, // index write params - std::make_shared(index_search_params), // index search params - 0, // num frozen points - false, // not a dynamic_index - false, // no enable_tags/ids - false, // no concurrent_consolidate, - false, // pq_dist_build - 0, // num_pq_chunks - false); // use_opq = false -} - -template -StaticMemoryIndex
::StaticMemoryIndex(const diskann::Metric m, const std::string &index_prefix, - const size_t num_points, const size_t dimensions, const uint32_t num_threads, - const uint32_t initial_search_complexity) - : _index(static_index_builder
(m, num_points, dimensions, initial_search_complexity)) -{ - const uint32_t _num_threads = num_threads != 0 ? num_threads : omp_get_num_procs(); - _index.load(index_prefix.c_str(), _num_threads, initial_search_complexity); -} - -template -NeighborsAndDistances StaticMemoryIndex
::search( - py::array_t &query, const uint64_t knn, const uint64_t complexity) -{ - py::array_t ids(knn); - py::array_t dists(knn); - std::vector
empty_vector; - _index.search(query.data(), knn, complexity, ids.mutable_data(), dists.mutable_data()); - return std::make_pair(ids, dists); -} - -template -NeighborsAndDistances StaticMemoryIndex
::search_with_filter( - py::array_t &query, const uint64_t knn, const uint64_t complexity, - const filterT filter) -{ - py::array_t ids(knn); - py::array_t dists(knn); - std::vector
empty_vector; - _index.search_with_filters(query.data(), filter, knn, complexity, ids.mutable_data(), dists.mutable_data()); - return std::make_pair(ids, dists); -} - -template -NeighborsAndDistances StaticMemoryIndex
::batch_search( - py::array_t &queries, const uint64_t num_queries, const uint64_t knn, - const uint64_t complexity, const uint32_t num_threads) -{ - const uint32_t _num_threads = num_threads != 0 ? num_threads : omp_get_num_procs(); - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector
empty_vector; - - omp_set_num_threads(static_cast(_num_threads)); - -#pragma omp parallel for schedule(dynamic, 1) default(none) shared(num_queries, queries, knn, complexity, ids, dists) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - _index.search(queries.data(i), knn, complexity, ids.mutable_data(i), dists.mutable_data(i)); - } - - return std::make_pair(ids, dists); -} - -template class StaticMemoryIndex; -template class StaticMemoryIndex; -template class StaticMemoryIndex; - -} // namespace diskannpy \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/__init__.py b/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/__init__.py deleted file mode 100644 index 4aeb960..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -from .build_memory_index import build_random_vectors_and_memory_index -from .create_test_data import random_vectors, vectors_as_temp_file, write_vectors -from .recall import calculate_recall diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/build_memory_index.py b/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/build_memory_index.py deleted file mode 100644 index 3c30bed..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/build_memory_index.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import os -from tempfile import mkdtemp - -import diskannpy as dap -import numpy as np - -from .create_test_data import random_vectors - - -def build_random_vectors_and_memory_index( - dtype, metric, with_tags: bool = False, index_prefix: str = "ann", seed: int = 12345 -): - query_vectors: np.ndarray = random_vectors(1000, 10, dtype=dtype, seed=seed) - index_vectors: np.ndarray = random_vectors(10000, 10, dtype=dtype, seed=seed) - ann_dir = mkdtemp() - - if with_tags: - rng = np.random.default_rng(seed) - tags = np.arange(start=1, stop=10001, dtype=np.uint32) - rng.shuffle(tags) - else: - tags = "" - - dap.build_memory_index( - data=index_vectors, - distance_metric=metric, - index_directory=ann_dir, - graph_degree=16, - complexity=32, - alpha=1.2, - num_threads=0, - use_pq_build=False, - num_pq_bytes=8, - use_opq=False, - filter_complexity=32, - tags=tags, - index_prefix=index_prefix, - ) - - return ( - metric, - dtype, - query_vectors, - index_vectors, - ann_dir, - os.path.join(ann_dir, "vectors.bin"), - tags, - ) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/create_test_data.py b/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/create_test_data.py deleted file mode 100644 index 44e413e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/create_test_data.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -from contextlib import contextmanager -from pathlib import Path -from tempfile import NamedTemporaryFile -from typing import BinaryIO - -import numpy as np - - -def random_vectors(rows: int, dimensions: int, dtype, seed: int = 12345) -> np.ndarray: - rng = np.random.default_rng(seed) - if dtype == np.float32: - vectors = rng.random((rows, dimensions), dtype=dtype) - elif dtype == np.uint8: - vectors = rng.integers( - low=0, high=256, size=(rows, dimensions), dtype=dtype - ) # low is inclusive, high is exclusive - elif dtype == np.int8: - vectors = rng.integers( - low=-128, high=128, size=(rows, dimensions), dtype=dtype - ) # low is inclusive, high is exclusive - else: - raise RuntimeError("Only np.float32, np.int8, and np.uint8 are supported") - return vectors - - -def write_vectors(file_handler: BinaryIO, vectors: np.ndarray): - _ = file_handler.write(np.array(vectors.shape, dtype=np.int32).tobytes()) - _ = file_handler.write(vectors.tobytes()) - - -@contextmanager -def vectors_as_temp_file(vectors: np.ndarray) -> str: - temp = NamedTemporaryFile(mode="wb", delete=False) - write_vectors(temp, vectors) - temp.close() - yield temp.name - Path(temp.name).unlink() diff --git a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/recall.py b/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/recall.py deleted file mode 100644 index 03f38f3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/python/tests/fixtures/recall.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import numpy as np - - -def calculate_recall( - result_set_indices: np.ndarray, truth_set_indices: np.ndarray, recall_at: int = 5 -) -> float: - """ - result_set_indices and truth_set_indices correspond by row index. the columns in each row contain the indices of - the nearest neighbors, with result_set_indices being the approximate nearest neighbor results and truth_set_indices - being the brute force nearest neighbor calculation via sklearn's NearestNeighbor class. - :param result_set_indices: - :param truth_set_indices: - :param recall_at: - :return: - """ - found = 0 - for i in range(0, result_set_indices.shape[0]): - result_set_set = set(result_set_indices[i][0:recall_at]) - truth_set_set = set(truth_set_indices[i][0:recall_at]) - found += len(result_set_set.intersection(truth_set_set)) - return found / (result_set_indices.shape[0] * recall_at) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.lock b/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.lock deleted file mode 100644 index 3a8a252..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.lock +++ /dev/null @@ -1,1820 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 3 - -[[package]] -name = "adler" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" - -[[package]] -name = "ahash" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c99f64d1e06488f620f932677e24bc6e2897582980441ae90a671415bd7ec2f" -dependencies = [ - "cfg-if", - "once_cell", - "version_check", -] - -[[package]] -name = "anes" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" - -[[package]] -name = "anstream" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ca84f3628370c59db74ee214b3263d58f9aadd9b4fe7e711fd87dc452b7f163" -dependencies = [ - "anstyle", - "anstyle-parse", - "anstyle-query", - "anstyle-wincon", - "colorchoice", - "is-terminal", - "utf8parse", -] - -[[package]] -name = "anstyle" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a30da5c5f2d5e72842e00bcb57657162cdabef0931f40e2deb9b4140440cecd" - -[[package]] -name = "anstyle-parse" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "938874ff5980b03a87c5524b3ae5b59cf99b1d6bc836848df7bc5ada9643c333" -dependencies = [ - "utf8parse", -] - -[[package]] -name = "anstyle-query" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b" -dependencies = [ - "windows-sys 0.48.0", -] - -[[package]] -name = "anstyle-wincon" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "180abfa45703aebe0093f79badacc01b8fd4ea2e35118747e5811127f926e188" -dependencies = [ - "anstyle", - "windows-sys 0.48.0", -] - -[[package]] -name = "anyhow" -version = "1.0.71" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c7d0618f0e0b7e8ff11427422b64564d5fb0be1940354bfe2e0529b18a9d9b8" - -[[package]] -name = "approx" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" -dependencies = [ - "num-traits", -] - -[[package]] -name = "autocfg" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" - -[[package]] -name = "base64" -version = "0.21.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "604178f6c5c21f02dc555784810edfb88d34ac2c73b2eae109655649ee73ce3d" - -[[package]] -name = "bincode" -version = "1.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" -dependencies = [ - "serde", -] - -[[package]] -name = "bit-vec" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" - -[[package]] -name = "bitflags" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" - -[[package]] -name = "bitflags" -version = "2.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "327762f6e5a765692301e5bb513e0d9fef63be86bbc14528052b1cd3e6f03e07" - -[[package]] -name = "build_and_insert_delete_memory_index" -version = "0.1.0" -dependencies = [ - "diskann", - "logger", - "vector", -] - -[[package]] -name = "build_and_insert_memory_index" -version = "0.1.0" -dependencies = [ - "diskann", - "logger", - "vector", -] - -[[package]] -name = "build_disk_index" -version = "0.1.0" -dependencies = [ - "diskann", - "logger", - "openblas-src", - "vector", -] - -[[package]] -name = "build_memory_index" -version = "0.1.0" -dependencies = [ - "clap", - "diskann", - "logger", - "vector", -] - -[[package]] -name = "bumpalo" -version = "3.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1" - -[[package]] -name = "bytemuck" -version = "1.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17febce684fd15d89027105661fec94afb475cb995fbc59d2865198446ba2eea" - -[[package]] -name = "byteorder" -version = "1.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" - -[[package]] -name = "bytes" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" - -[[package]] -name = "cast" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" - -[[package]] -name = "cblas" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3de46dff748ed7e891bc46faae117f48d2a7911041c6630aed3c61a3fe12326f" -dependencies = [ - "cblas-sys", - "libc", - "num-complex", -] - -[[package]] -name = "cblas-sys" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6feecd82cce51b0204cf063f0041d69f24ce83f680d87514b004248e7b0fa65" -dependencies = [ - "libc", -] - -[[package]] -name = "cc" -version = "1.0.79" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f" - -[[package]] -name = "cfg-if" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" - -[[package]] -name = "ciborium" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "effd91f6c78e5a4ace8a5d3c0b6bfaec9e2baaef55f3efc00e45fb2e477ee926" -dependencies = [ - "ciborium-io", - "ciborium-ll", - "serde", -] - -[[package]] -name = "ciborium-io" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdf919175532b369853f5d5e20b26b43112613fd6fe7aee757e35f7a44642656" - -[[package]] -name = "ciborium-ll" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "defaa24ecc093c77630e6c15e17c51f5e187bf35ee514f4e2d67baaa96dae22b" -dependencies = [ - "ciborium-io", - "half 1.8.2", -] - -[[package]] -name = "clap" -version = "4.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9394150f5b4273a1763355bd1c2ec54cc5a2593f790587bcd6b2c947cfa9211" -dependencies = [ - "clap_builder", - "clap_derive", - "once_cell", -] - -[[package]] -name = "clap_builder" -version = "4.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a78fbdd3cc2914ddf37ba444114bc7765bbdcb55ec9cbe6fa054f0137400717" -dependencies = [ - "anstream", - "anstyle", - "bitflags 1.3.2", - "clap_lex", - "strsim", -] - -[[package]] -name = "clap_derive" -version = "4.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8cd2b2a819ad6eec39e8f1d6b53001af1e5469f8c177579cdaeb313115b825f" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.18", -] - -[[package]] -name = "clap_lex" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b" - -[[package]] -name = "colorchoice" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" - -[[package]] -name = "convert_f32_to_bf16" -version = "0.1.0" -dependencies = [ - "half 2.2.1", -] - -[[package]] -name = "core-foundation" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "194a7a9e6de53fa55116934067c844d9d749312f75c6f6d0980e8c252f8c2146" -dependencies = [ - "core-foundation-sys", - "libc", -] - -[[package]] -name = "core-foundation-sys" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e496a50fda8aacccc86d7529e2c1e0892dbd0f898a6b5645b5561b89c3210efa" - -[[package]] -name = "crc32fast" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "criterion" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" -dependencies = [ - "anes", - "cast", - "ciborium", - "clap", - "criterion-plot", - "is-terminal", - "itertools", - "num-traits", - "once_cell", - "oorandom", - "plotters", - "rayon", - "regex", - "serde", - "serde_derive", - "serde_json", - "tinytemplate", - "walkdir", -] - -[[package]] -name = "criterion-plot" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" -dependencies = [ - "cast", - "itertools", -] - -[[package]] -name = "crossbeam" -version = "0.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2801af0d36612ae591caa9568261fddce32ce6e08a7275ea334a06a4ad021a2c" -dependencies = [ - "cfg-if", - "crossbeam-channel", - "crossbeam-deque", - "crossbeam-epoch", - "crossbeam-queue", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-channel" -version = "0.5.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a33c2bf77f2df06183c3aa30d1e96c0695a313d4f9c453cc3762a6db39f99200" -dependencies = [ - "cfg-if", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef" -dependencies = [ - "cfg-if", - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7" -dependencies = [ - "autocfg", - "cfg-if", - "crossbeam-utils", - "memoffset", - "scopeguard", -] - -[[package]] -name = "crossbeam-queue" -version = "0.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1cfb3ea8a53f37c40dea2c7bedcbd88bdfae54f5e2175d6ecaff1c988353add" -dependencies = [ - "cfg-if", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-utils" -version = "0.8.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "crunchy" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" - -[[package]] -name = "dirs" -version = "3.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30baa043103c9d0c2a57cf537cc2f35623889dc0d405e6c3cccfadbc81c71309" -dependencies = [ - "dirs-sys", -] - -[[package]] -name = "dirs-sys" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b1d1d91c932ef41c0f2663aa8b0ca0342d444d842c06914aa0a7e352d0bada6" -dependencies = [ - "libc", - "redox_users", - "winapi", -] - -[[package]] -name = "diskann" -version = "0.1.0" -dependencies = [ - "approx", - "bincode", - "bit-vec", - "byteorder", - "cblas", - "cc", - "criterion", - "crossbeam", - "half 2.2.1", - "hashbrown 0.13.2", - "logger", - "num-traits", - "once_cell", - "openblas-src", - "platform", - "rand", - "rayon", - "serde", - "thiserror", - "vector", - "winapi", -] - -[[package]] -name = "either" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91" - -[[package]] -name = "errno" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bcfec3a70f97c962c307b2d2c56e358cf1d00b558d74262b5f929ee8cc7e73a" -dependencies = [ - "errno-dragonfly", - "libc", - "windows-sys 0.48.0", -] - -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", -] - -[[package]] -name = "fastrand" -version = "1.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e51093e27b0797c359783294ca4f0a911c270184cb10f85783b118614a1501be" -dependencies = [ - "instant", -] - -[[package]] -name = "filetime" -version = "0.2.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cbc844cecaee9d4443931972e1289c8ff485cb4cc2767cb03ca139ed6885153" -dependencies = [ - "cfg-if", - "libc", - "redox_syscall 0.2.16", - "windows-sys 0.48.0", -] - -[[package]] -name = "fixedbitset" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" - -[[package]] -name = "flate2" -version = "1.0.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743" -dependencies = [ - "crc32fast", - "miniz_oxide", -] - -[[package]] -name = "foreign-types" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" -dependencies = [ - "foreign-types-shared", -] - -[[package]] -name = "foreign-types-shared" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" - -[[package]] -name = "form_urlencoded" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652" -dependencies = [ - "percent-encoding", -] - -[[package]] -name = "getrandom" -version = "0.2.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" -dependencies = [ - "cfg-if", - "libc", - "wasi", -] - -[[package]] -name = "half" -version = "1.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" - -[[package]] -name = "half" -version = "2.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02b4af3693f1b705df946e9fe5631932443781d0aabb423b62fcd4d73f6d2fd0" -dependencies = [ - "crunchy", -] - -[[package]] -name = "hashbrown" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" - -[[package]] -name = "hashbrown" -version = "0.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e" -dependencies = [ - "ahash", -] - -[[package]] -name = "heck" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" - -[[package]] -name = "hermit-abi" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee512640fe35acbfb4bb779db6f0d80704c2cacfa2e39b601ef3e3f47d1ae4c7" -dependencies = [ - "libc", -] - -[[package]] -name = "hermit-abi" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286" - -[[package]] -name = "idna" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d20d6b07bfbc108882d88ed8e37d39636dcc260e15e30c45e6ba089610b917c" -dependencies = [ - "unicode-bidi", - "unicode-normalization", -] - -[[package]] -name = "indexmap" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" -dependencies = [ - "autocfg", - "hashbrown 0.12.3", -] - -[[package]] -name = "instant" -version = "0.1.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "io-lifetimes" -version = "1.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eae7b9aee968036d54dce06cebaefd919e4472e753296daccd6d344e3e2df0c2" -dependencies = [ - "hermit-abi 0.3.1", - "libc", - "windows-sys 0.48.0", -] - -[[package]] -name = "is-terminal" -version = "0.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adcf93614601c8129ddf72e2d5633df827ba6551541c6d8c59520a371475be1f" -dependencies = [ - "hermit-abi 0.3.1", - "io-lifetimes", - "rustix", - "windows-sys 0.48.0", -] - -[[package]] -name = "itertools" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" -dependencies = [ - "either", -] - -[[package]] -name = "itoa" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "453ad9f582a441959e5f0d088b02ce04cfe8d51a8eaf077f12ac6d3e94164ca6" - -[[package]] -name = "js-sys" -version = "0.3.64" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5f195fe497f702db0f318b07fdd68edb16955aed830df8363d837542f8f935a" -dependencies = [ - "wasm-bindgen", -] - -[[package]] -name = "lazy_static" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" - -[[package]] -name = "libc" -version = "0.2.146" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f92be4933c13fd498862a9e02a3055f8a8d9c039ce33db97306fd5a6caa7f29b" - -[[package]] -name = "linux-raw-sys" -version = "0.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef53942eb7bf7ff43a617b3e2c1c4a5ecf5944a7c1bc12d7ee39bbb15e5c1519" - -[[package]] -name = "load_and_insert_memory_index" -version = "0.1.0" -dependencies = [ - "diskann", - "logger", - "vector", -] - -[[package]] -name = "log" -version = "0.4.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b06a4cde4c0f271a446782e3eff8de789548ce57dbc8eca9292c27f4a42004b4" - -[[package]] -name = "logger" -version = "0.1.0" -dependencies = [ - "lazy_static", - "log", - "once_cell", - "prost", - "prost-build", - "prost-types", - "thiserror", - "vcpkg", - "win_etw_macros", - "win_etw_provider", -] - -[[package]] -name = "memoffset" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c" -dependencies = [ - "autocfg", -] - -[[package]] -name = "miniz_oxide" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" -dependencies = [ - "adler", -] - -[[package]] -name = "multimap" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a" - -[[package]] -name = "native-tls" -version = "0.2.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07226173c32f2926027b63cce4bcd8076c3552846cbe7925f3aaffeac0a3b92e" -dependencies = [ - "lazy_static", - "libc", - "log", - "openssl", - "openssl-probe", - "openssl-sys", - "schannel", - "security-framework", - "security-framework-sys", - "tempfile", -] - -[[package]] -name = "num-complex" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02e0d21255c828d6f128a1e41534206671e8c3ea0c62f32291e808dc82cff17d" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-traits" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" -dependencies = [ - "autocfg", -] - -[[package]] -name = "num_cpus" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b" -dependencies = [ - "hermit-abi 0.2.6", - "libc", -] - -[[package]] -name = "once_cell" -version = "1.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" - -[[package]] -name = "oorandom" -version = "11.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575" - -[[package]] -name = "openblas-build" -version = "0.10.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eba42c395477605f400a8d79ee0b756cfb82abe3eb5618e35fa70d3a36010a7f" -dependencies = [ - "anyhow", - "flate2", - "native-tls", - "tar", - "thiserror", - "ureq", - "walkdir", -] - -[[package]] -name = "openblas-src" -version = "0.10.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38e5d8af0b707ac2fe1574daa88b4157da73b0de3dc7c39fe3e2c0bb64070501" -dependencies = [ - "dirs", - "openblas-build", - "vcpkg", -] - -[[package]] -name = "openssl" -version = "0.10.60" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79a4c6c3a2b158f7f8f2a2fc5a969fa3a068df6fc9dbb4a43845436e3af7c800" -dependencies = [ - "bitflags 2.4.1", - "cfg-if", - "foreign-types", - "libc", - "once_cell", - "openssl-macros", - "openssl-sys", -] - -[[package]] -name = "openssl-macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.18", -] - -[[package]] -name = "openssl-probe" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" - -[[package]] -name = "openssl-sys" -version = "0.9.96" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3812c071ba60da8b5677cc12bcb1d42989a65553772897a7e0355545a819838f" -dependencies = [ - "cc", - "libc", - "pkg-config", - "vcpkg", -] - -[[package]] -name = "percent-encoding" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" - -[[package]] -name = "petgraph" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dd7d28ee937e54fe3080c91faa1c3a46c06de6252988a7f4592ba2310ef22a4" -dependencies = [ - "fixedbitset", - "indexmap", -] - -[[package]] -name = "pkg-config" -version = "0.3.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" - -[[package]] -name = "platform" -version = "0.1.0" -dependencies = [ - "log", - "winapi", -] - -[[package]] -name = "plotters" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2c224ba00d7cadd4d5c660deaf2098e5e80e07846537c51f9cfa4be50c1fd45" -dependencies = [ - "num-traits", - "plotters-backend", - "plotters-svg", - "wasm-bindgen", - "web-sys", -] - -[[package]] -name = "plotters-backend" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e76628b4d3a7581389a35d5b6e2139607ad7c75b17aed325f210aa91f4a9609" - -[[package]] -name = "plotters-svg" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38f6d39893cca0701371e3c27294f09797214b86f1fb951b89ade8ec04e2abab" -dependencies = [ - "plotters-backend", -] - -[[package]] -name = "ppv-lite86" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" - -[[package]] -name = "prettyplease" -version = "0.1.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c8646e95016a7a6c4adea95bafa8a16baab64b583356217f2c85db4a39d9a86" -dependencies = [ - "proc-macro2", - "syn 1.0.109", -] - -[[package]] -name = "proc-macro2" -version = "1.0.60" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dec2b086b7a862cf4de201096214fa870344cf922b2b30c167badb3af3195406" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "prost" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd" -dependencies = [ - "bytes", - "prost-derive", -] - -[[package]] -name = "prost-build" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "119533552c9a7ffacc21e099c24a0ac8bb19c2a2a3f363de84cd9b844feab270" -dependencies = [ - "bytes", - "heck", - "itertools", - "lazy_static", - "log", - "multimap", - "petgraph", - "prettyplease", - "prost", - "prost-types", - "regex", - "syn 1.0.109", - "tempfile", - "which", -] - -[[package]] -name = "prost-derive" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5d2d8d10f3c6ded6da8b05b5fb3b8a5082514344d56c9f871412d29b4e075b4" -dependencies = [ - "anyhow", - "itertools", - "proc-macro2", - "quote", - "syn 1.0.109", -] - -[[package]] -name = "prost-types" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "213622a1460818959ac1181aaeb2dc9c7f63df720db7d788b3e24eacd1983e13" -dependencies = [ - "prost", -] - -[[package]] -name = "quote" -version = "1.0.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9ab9c7eadfd8df19006f1cf1a4aed13540ed5cbc047010ece5826e10825488" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "rand" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" -dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", -] - -[[package]] -name = "rand_core" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] - -[[package]] -name = "rayon" -version = "1.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d2df5196e37bcc87abebc0053e20787d73847bb33134a69841207dd0a47f03b" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b8f95bd6966f5c87776639160a66bd8ab9895d9d4ab01ddba9fc60661aebe8d" -dependencies = [ - "crossbeam-channel", - "crossbeam-deque", - "crossbeam-utils", - "num_cpus", -] - -[[package]] -name = "redox_syscall" -version = "0.2.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" -dependencies = [ - "bitflags 1.3.2", -] - -[[package]] -name = "redox_syscall" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "567664f262709473930a4bf9e51bf2ebf3348f2e748ccc50dea20646858f8f29" -dependencies = [ - "bitflags 1.3.2", -] - -[[package]] -name = "redox_users" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b033d837a7cf162d7993aded9304e30a83213c648b6e389db233191f891e5c2b" -dependencies = [ - "getrandom", - "redox_syscall 0.2.16", - "thiserror", -] - -[[package]] -name = "regex" -version = "1.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0ab3ca65655bb1e41f2a8c8cd662eb4fb035e67c3f78da1d61dffe89d07300f" -dependencies = [ - "regex-syntax", -] - -[[package]] -name = "regex-syntax" -version = "0.7.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "436b050e76ed2903236f032a59761c1eb99e1b0aead2c257922771dab1fc8c78" - -[[package]] -name = "rustix" -version = "0.37.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4eb579851244c2c03e7c24f501c3432bed80b8f720af1d6e5b0e0f01555a035" -dependencies = [ - "bitflags 1.3.2", - "errno", - "io-lifetimes", - "libc", - "linux-raw-sys", - "windows-sys 0.48.0", -] - -[[package]] -name = "rustls-native-certs" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9aace74cb666635c918e9c12bc0d348266037aa8eb599b5cba565709a8dff00" -dependencies = [ - "openssl-probe", - "rustls-pemfile", - "schannel", - "security-framework", -] - -[[package]] -name = "rustls-pemfile" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d194b56d58803a43635bdc398cd17e383d6f71f9182b9a192c127ca42494a59b" -dependencies = [ - "base64", -] - -[[package]] -name = "ryu" -version = "1.0.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041" - -[[package]] -name = "same-file" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" -dependencies = [ - "winapi-util", -] - -[[package]] -name = "schannel" -version = "0.1.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "713cfb06c7059f3588fb8044c0fad1d09e3c01d225e25b9220dbfdcf16dbb1b3" -dependencies = [ - "windows-sys 0.42.0", -] - -[[package]] -name = "scopeguard" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" - -[[package]] -name = "search_memory_index" -version = "0.1.0" -dependencies = [ - "bytemuck", - "diskann", - "num_cpus", - "rayon", - "vector", -] - -[[package]] -name = "security-framework" -version = "2.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fc758eb7bffce5b308734e9b0c1468893cae9ff70ebf13e7090be8dcbcc83a8" -dependencies = [ - "bitflags 1.3.2", - "core-foundation", - "core-foundation-sys", - "libc", - "security-framework-sys", -] - -[[package]] -name = "security-framework-sys" -version = "2.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f51d0c0d83bec45f16480d0ce0058397a69e48fcdc52d1dc8855fb68acbd31a7" -dependencies = [ - "core-foundation-sys", - "libc", -] - -[[package]] -name = "serde" -version = "1.0.164" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e8c8cf938e98f769bc164923b06dce91cea1751522f46f8466461af04c9027d" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.164" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9735b638ccc51c28bf6914d90a2e9725b377144fc612c49a611fddd1b631d68" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.18", -] - -[[package]] -name = "serde_json" -version = "1.0.97" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdf3bf93142acad5821c99197022e170842cdbc1c30482b98750c688c640842a" -dependencies = [ - "itoa", - "ryu", - "serde", -] - -[[package]] -name = "sha1_smol" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae1a47186c03a32177042e55dbc5fd5aee900b8e0069a8d70fba96a9375cd012" - -[[package]] -name = "strsim" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" - -[[package]] -name = "syn" -version = "1.0.109" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "syn" -version = "2.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32d41677bcbe24c20c52e7c70b0d8db04134c5d1066bf98662e2871ad200ea3e" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "tar" -version = "0.4.38" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b55807c0344e1e6c04d7c965f5289c39a8d94ae23ed5c0b57aabac549f871c6" -dependencies = [ - "filetime", - "libc", - "xattr", -] - -[[package]] -name = "tempfile" -version = "3.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31c0432476357e58790aaa47a8efb0c5138f137343f3b5f23bd36a27e3b0a6d6" -dependencies = [ - "autocfg", - "cfg-if", - "fastrand", - "redox_syscall 0.3.5", - "rustix", - "windows-sys 0.48.0", -] - -[[package]] -name = "thiserror" -version = "1.0.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "978c9a314bd8dc99be594bc3c175faaa9794be04a5a5e153caba6915336cebac" -dependencies = [ - "thiserror-impl", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.18", -] - -[[package]] -name = "tinytemplate" -version = "1.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" -dependencies = [ - "serde", - "serde_json", -] - -[[package]] -name = "tinyvec" -version = "1.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50" -dependencies = [ - "tinyvec_macros", -] - -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - -[[package]] -name = "unicode-bidi" -version = "0.3.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92888ba5573ff080736b3648696b70cafad7d250551175acbaa4e0385b3e1460" - -[[package]] -name = "unicode-ident" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b15811caf2415fb889178633e7724bad2509101cde276048e013b9def5e51fa0" - -[[package]] -name = "unicode-normalization" -version = "0.1.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921" -dependencies = [ - "tinyvec", -] - -[[package]] -name = "ureq" -version = "2.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b11c96ac7ee530603dcdf68ed1557050f374ce55a5a07193ebf8cbc9f8927e9" -dependencies = [ - "base64", - "flate2", - "log", - "native-tls", - "once_cell", - "rustls-native-certs", - "url", -] - -[[package]] -name = "url" -version = "2.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50bff7831e19200a85b17131d085c25d7811bc4e186efdaf54bbd132994a88cb" -dependencies = [ - "form_urlencoded", - "idna", - "percent-encoding", -] - -[[package]] -name = "utf8parse" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" - -[[package]] -name = "uuid" -version = "1.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fa2982af2eec27de306107c027578ff7f423d65f7250e40ce0fea8f45248b81" -dependencies = [ - "sha1_smol", -] - -[[package]] -name = "vcpkg" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" - -[[package]] -name = "vector" -version = "0.1.0" -dependencies = [ - "approx", - "base64", - "bincode", - "bytemuck", - "cc", - "half 2.2.1", - "rand", - "serde", - "thiserror", -] - -[[package]] -name = "vector_base64" -version = "0.1.0" -dependencies = [ - "base64", - "bincode", - "half 2.2.1", - "serde", -] - -[[package]] -name = "version_check" -version = "0.9.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" - -[[package]] -name = "w32-error" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7c61a6bd91e168c12fc170985725340f6b458eb6f971d1cf6c34f74ffafb43" -dependencies = [ - "winapi", -] - -[[package]] -name = "walkdir" -version = "2.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36df944cda56c7d8d8b7496af378e6b16de9284591917d307c9b4d313c44e698" -dependencies = [ - "same-file", - "winapi-util", -] - -[[package]] -name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" - -[[package]] -name = "wasm-bindgen" -version = "0.2.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7706a72ab36d8cb1f80ffbf0e071533974a60d0a308d01a5d0375bf60499a342" -dependencies = [ - "cfg-if", - "wasm-bindgen-macro", -] - -[[package]] -name = "wasm-bindgen-backend" -version = "0.2.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ef2b6d3c510e9625e5fe6f509ab07d66a760f0885d858736483c32ed7809abd" -dependencies = [ - "bumpalo", - "log", - "once_cell", - "proc-macro2", - "quote", - "syn 2.0.18", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-macro" -version = "0.2.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dee495e55982a3bd48105a7b947fd2a9b4a8ae3010041b9e0faab3f9cd028f1d" -dependencies = [ - "quote", - "wasm-bindgen-macro-support", -] - -[[package]] -name = "wasm-bindgen-macro-support" -version = "0.2.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.18", - "wasm-bindgen-backend", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-shared" -version = "0.2.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca6ad05a4870b2bf5fe995117d3728437bd27d7cd5f06f13c17443ef369775a1" - -[[package]] -name = "web-sys" -version = "0.3.64" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b85cbef8c220a6abc02aefd892dfc0fc23afb1c6a426316ec33253a3877249b" -dependencies = [ - "js-sys", - "wasm-bindgen", -] - -[[package]] -name = "which" -version = "4.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2441c784c52b289a054b7201fc93253e288f094e2f4be9058343127c4226a269" -dependencies = [ - "either", - "libc", - "once_cell", -] - -[[package]] -name = "widestring" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "653f141f39ec16bba3c5abe400a0c60da7468261cc2cbf36805022876bc721a8" - -[[package]] -name = "win_etw_macros" -version = "0.1.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bc4c591edb4858e3445f7a60c7e0a50915aedadfa044f28f17c98c145ef54d" -dependencies = [ - "proc-macro2", - "quote", - "sha1_smol", - "syn 1.0.109", - "uuid", - "win_etw_metadata", -] - -[[package]] -name = "win_etw_metadata" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e50d0fa665033a19ecefd281b4fb5481eba2972dedbb5ec129c9392a206d652f" -dependencies = [ - "bitflags 1.3.2", -] - -[[package]] -name = "win_etw_provider" -version = "0.1.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dffcc196e0e180e73a275a91f6914f173227fd627cabac3efdd8d6adec113892" -dependencies = [ - "w32-error", - "widestring", - "win_etw_metadata", - "winapi", - "zerocopy", -] - -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-util" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" -dependencies = [ - "winapi", -] - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - -[[package]] -name = "windows-sys" -version = "0.42.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" -dependencies = [ - "windows_aarch64_gnullvm 0.42.2", - "windows_aarch64_msvc 0.42.2", - "windows_i686_gnu 0.42.2", - "windows_i686_msvc 0.42.2", - "windows_x86_64_gnu 0.42.2", - "windows_x86_64_gnullvm 0.42.2", - "windows_x86_64_msvc 0.42.2", -] - -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets", -] - -[[package]] -name = "windows-targets" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b1eb6f0cd7c80c79759c929114ef071b87354ce476d9d94271031c0497adfd5" -dependencies = [ - "windows_aarch64_gnullvm 0.48.0", - "windows_aarch64_msvc 0.48.0", - "windows_i686_gnu 0.48.0", - "windows_i686_msvc 0.48.0", - "windows_x86_64_gnu 0.48.0", - "windows_x86_64_gnullvm 0.48.0", - "windows_x86_64_msvc 0.48.0", -] - -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" - -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91ae572e1b79dba883e0d315474df7305d12f569b400fcf90581b06062f7e1bc" - -[[package]] -name = "windows_aarch64_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" - -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2ef27e0d7bdfcfc7b868b317c1d32c641a6fe4629c171b8928c7b08d98d7cf3" - -[[package]] -name = "windows_i686_gnu" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" - -[[package]] -name = "windows_i686_gnu" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "622a1962a7db830d6fd0a69683c80a18fda201879f0f447f065a3b7467daa241" - -[[package]] -name = "windows_i686_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" - -[[package]] -name = "windows_i686_msvc" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4542c6e364ce21bf45d69fdd2a8e455fa38d316158cfd43b3ac1c5b1b19f8e00" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca2b8a661f7628cbd23440e50b05d705db3686f894fc9580820623656af974b1" - -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" - -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7896dbc1f41e08872e9d5e8f8baa8fdd2677f29468c4e156210174edc7f7b953" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a515f5799fe4961cb532f983ce2b23082366b898e52ffbce459c86f67c8378a" - -[[package]] -name = "xattr" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d1526bbe5aaeb5eb06885f4d987bcdfa5e23187055de9b83fe00156a821fabc" -dependencies = [ - "libc", -] - -[[package]] -name = "zerocopy" -version = "0.6.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "854e949ac82d619ee9a14c66a1b674ac730422372ccb759ce0c39cabcf2bf8e6" -dependencies = [ - "byteorder", - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.6.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "125139de3f6b9d625c39e2efdd73d41bdac468ccd556556440e322be0e1bbd91" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.18", -] diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.toml deleted file mode 100644 index 5236f96..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -[workspace] -members = [ - "cmd_drivers/build_memory_index", - "cmd_drivers/build_and_insert_memory_index", - "cmd_drivers/load_and_insert_memory_index", - "cmd_drivers/convert_f32_to_bf16", - "cmd_drivers/search_memory_index", - "cmd_drivers/build_disk_index", - "cmd_drivers/build_and_insert_delete_memory_index", - "vector", - "diskann", - "platform", - "logger", - "vector_base64" -] -resolver = "2" - -[profile.release] -opt-level = 3 -codegen-units=1 diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/Cargo.toml deleted file mode 100644 index 42aa185..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/Cargo.toml +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "build_and_insert_delete_memory_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -diskann = { path = "../../diskann" } -logger = { path = "../../logger" } -vector = { path = "../../vector" } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/src/main.rs deleted file mode 100644 index 4593a9e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_delete_memory_index/src/main.rs +++ /dev/null @@ -1,420 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::env; - -use diskann::{ - common::{ANNError, ANNResult}, - index::create_inmem_index, - model::{ - configuration::index_write_parameters::IndexWriteParametersBuilder, - vertex::{DIM_104, DIM_128, DIM_256}, - IndexConfiguration, - }, - utils::round_up, - utils::{file_exists, load_ids_to_delete_from_file, load_metadata_from_file, Timer}, -}; - -use vector::{FullPrecisionDistance, Half, Metric}; - -// The main function to build an in-memory index -#[allow(clippy::too_many_arguments)] -fn build_and_insert_delete_in_memory_index( - metric: Metric, - data_path: &str, - delta_path: &str, - r: u32, - l: u32, - alpha: f32, - save_path: &str, - num_threads: u32, - _use_pq_build: bool, - _num_pq_bytes: usize, - use_opq: bool, - delete_path: &str, -) -> ANNResult<()> -where - T: Default + Copy + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - let index_write_parameters = IndexWriteParametersBuilder::new(l, r) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - let (data_num, data_dim) = load_metadata_from_file(data_path)?; - - let config = IndexConfiguration::new( - metric, - data_dim, - round_up(data_dim as u64, 8_u64) as usize, - data_num, - false, - 0, - use_opq, - 0, - 2.0f32, - index_write_parameters, - ); - let mut index = create_inmem_index::(config)?; - - let timer = Timer::new(); - - index.build(data_path, data_num)?; - - let diff = timer.elapsed(); - - println!("Initial indexing time: {}", diff.as_secs_f64()); - - let (delta_data_num, _) = load_metadata_from_file(delta_path)?; - - index.insert(delta_path, delta_data_num)?; - - if !delete_path.is_empty() { - if !file_exists(delete_path) { - return Err(ANNError::log_index_error(format!( - "ERROR: Data file for delete {} does not exist.", - delete_path - ))); - } - - let (num_points_to_delete, vertex_ids_to_delete) = - load_ids_to_delete_from_file(delete_path)?; - index.soft_delete(vertex_ids_to_delete, num_points_to_delete)?; - } - - index.save(save_path)?; - - Ok(()) -} - -fn main() -> ANNResult<()> { - let mut data_type = String::new(); - let mut dist_fn = String::new(); - let mut data_path = String::new(); - let mut insert_path = String::new(); - let mut index_path_prefix = String::new(); - let mut delete_path = String::new(); - - let mut num_threads = 0u32; - let mut r = 64u32; - let mut l = 100u32; - - let mut alpha = 1.2f32; - let mut build_pq_bytes = 0u32; - let mut _use_pq_build = false; - let mut use_opq = false; - - let args: Vec = env::args().collect(); - let mut iter = args.iter().skip(1).peekable(); - - while let Some(arg) = iter.next() { - match arg.as_str() { - "--help" | "-h" => { - print_help(); - return Ok(()); - } - "--data_type" => { - data_type = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_type".to_string(), - "Missing data type".to_string(), - ) - })? - .to_owned(); - } - "--dist_fn" => { - dist_fn = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "dist_fn".to_string(), - "Missing distance function".to_string(), - ) - })? - .to_owned(); - } - "--data_path" => { - data_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_path".to_string(), - "Missing data path".to_string(), - ) - })? - .to_owned(); - } - "--insert_path" => { - insert_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "insert_path".to_string(), - "Missing insert path".to_string(), - ) - })? - .to_owned(); - } - "--index_path_prefix" => { - index_path_prefix = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "index_path_prefix".to_string(), - "Missing index path prefix".to_string(), - ) - })? - .to_owned(); - } - "--max_degree" | "-R" => { - r = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "max_degree".to_string(), - "Missing max degree".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "max_degree".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--Lbuild" | "-L" => { - l = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - "Missing build complexity".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--alpha" => { - alpha = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "alpha".to_string(), - "Missing alpha".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "alpha".to_string(), - format!("ParseFloatError: {}", err), - ) - })?; - } - "--num_threads" | "-T" => { - num_threads = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "num_threads".to_string(), - "Missing number of threads".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "num_threads".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--build_PQ_bytes" => { - build_pq_bytes = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - "Missing PQ bytes".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--use_opq" => { - use_opq = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "use_opq".to_string(), - "Missing use_opq flag".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "use_opq".to_string(), - format!("ParseBoolError: {}", err), - ) - })?; - } - "--delete_path" => { - delete_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "delete_path".to_string(), - "Missing delete_path".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "delete_set_path".to_string(), - format!("ParseStringError: {}", err), - ) - })?; - } - _ => { - return Err(ANNError::log_index_config_error( - String::from(""), - format!("Unknown argument: {}", arg), - )); - } - } - } - - if data_type.is_empty() - || dist_fn.is_empty() - || data_path.is_empty() - || index_path_prefix.is_empty() - { - return Err(ANNError::log_index_config_error( - String::from(""), - "Missing required arguments".to_string(), - )); - } - - _use_pq_build = build_pq_bytes > 0; - - let metric = dist_fn - .parse::() - .map_err(|err| ANNError::log_index_config_error("dist_fn".to_string(), err.to_string()))?; - - println!( - "Starting index build with R: {} Lbuild: {} alpha: {} #threads: {}", - r, l, alpha, num_threads - ); - - match data_type.as_str() { - "int8" => { - build_and_insert_delete_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - &delete_path, - )?; - } - "uint8" => { - build_and_insert_delete_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - &delete_path, - )?; - } - "float" => { - build_and_insert_delete_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - &delete_path, - )?; - } - "f16" => { - build_and_insert_delete_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - &delete_path, - )?; - } - _ => { - println!("Unsupported type. Use one of int8, uint8 or float."); - return Err(ANNError::log_index_config_error( - "data_type".to_string(), - "Invalid data type".to_string(), - )); - } - } - - Ok(()) -} - -fn print_help() { - println!("Arguments"); - println!("--help, -h Print information on arguments"); - println!("--data_type data type (required)"); - println!("--dist_fn distance function (required)"); - println!( - "--data_path Input data file in bin format for initial build (required)" - ); - println!("--insert_path Input data file in bin format for insert (required)"); - println!("--index_path_prefix Path prefix for saving index file components (required)"); - println!("--max_degree, -R Maximum graph degree (default: 64)"); - println!("--Lbuild, -L Build complexity, higher value results in better graphs (default: 100)"); - println!("--alpha alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for denser graphs with lower diameter (default: 1.2)"); - println!("--num_threads, -T Number of threads used for building index (defaults to num of CPU logic cores)"); - println!("--build_PQ_bytes Number of PQ bytes to build the index; 0 for full precision build (default: 0)"); - println!("--use_opq Set true for OPQ compression while using PQ distance comparisons for building the index, and false for PQ compression (default: false)"); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/Cargo.toml deleted file mode 100644 index d9811fc..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/Cargo.toml +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "build_and_insert_memory_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -diskann = { path = "../../diskann" } -logger = { path = "../../logger" } -vector = { path = "../../vector" } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/src/main.rs deleted file mode 100644 index 46e4ba4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_and_insert_memory_index/src/main.rs +++ /dev/null @@ -1,382 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::env; - -use diskann::{ - common::{ANNResult, ANNError}, - index::create_inmem_index, - utils::round_up, - model::{ - IndexWriteParametersBuilder, - IndexConfiguration, - vertex::{DIM_128, DIM_256, DIM_104} - }, - utils::{load_metadata_from_file, Timer}, -}; - -use vector::{Metric, FullPrecisionDistance, Half}; - -// The main function to build an in-memory index -#[allow(clippy::too_many_arguments)] -fn build_and_insert_in_memory_index ( - metric: Metric, - data_path: &str, - delta_path: &str, - r: u32, - l: u32, - alpha: f32, - save_path: &str, - num_threads: u32, - _use_pq_build: bool, - _num_pq_bytes: usize, - use_opq: bool -) -> ANNResult<()> -where - T: Default + Copy + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance -{ - let index_write_parameters = IndexWriteParametersBuilder::new(l, r) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - let (data_num, data_dim) = load_metadata_from_file(data_path)?; - - let config = IndexConfiguration::new( - metric, - data_dim, - round_up(data_dim as u64, 8_u64) as usize, - data_num, - false, - 0, - use_opq, - 0, - 2.0f32, - index_write_parameters, - ); - let mut index = create_inmem_index::(config)?; - - let timer = Timer::new(); - - index.build(data_path, data_num)?; - - let diff = timer.elapsed(); - - println!("Initial indexing time: {}", diff.as_secs_f64()); - - let (delta_data_num, _) = load_metadata_from_file(delta_path)?; - - index.insert(delta_path, delta_data_num)?; - - index.save(save_path)?; - - Ok(()) -} - -fn main() -> ANNResult<()> { - let mut data_type = String::new(); - let mut dist_fn = String::new(); - let mut data_path = String::new(); - let mut insert_path = String::new(); - let mut index_path_prefix = String::new(); - - let mut num_threads = 0u32; - let mut r = 64u32; - let mut l = 100u32; - - let mut alpha = 1.2f32; - let mut build_pq_bytes = 0u32; - let mut _use_pq_build = false; - let mut use_opq = false; - - let args: Vec = env::args().collect(); - let mut iter = args.iter().skip(1).peekable(); - - while let Some(arg) = iter.next() { - match arg.as_str() { - "--help" | "-h" => { - print_help(); - return Ok(()); - } - "--data_type" => { - data_type = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_type".to_string(), - "Missing data type".to_string(), - ) - })? - .to_owned(); - } - "--dist_fn" => { - dist_fn = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "dist_fn".to_string(), - "Missing distance function".to_string(), - ) - })? - .to_owned(); - } - "--data_path" => { - data_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_path".to_string(), - "Missing data path".to_string(), - ) - })? - .to_owned(); - } - "--insert_path" => { - insert_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "insert_path".to_string(), - "Missing insert path".to_string(), - ) - })? - .to_owned(); - } - "--index_path_prefix" => { - index_path_prefix = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "index_path_prefix".to_string(), - "Missing index path prefix".to_string(), - ) - })? - .to_owned(); - } - "--max_degree" | "-R" => { - r = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "max_degree".to_string(), - "Missing max degree".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "max_degree".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--Lbuild" | "-L" => { - l = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - "Missing build complexity".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--alpha" => { - alpha = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "alpha".to_string(), - "Missing alpha".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "alpha".to_string(), - format!("ParseFloatError: {}", err), - ) - })?; - } - "--num_threads" | "-T" => { - num_threads = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "num_threads".to_string(), - "Missing number of threads".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "num_threads".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--build_PQ_bytes" => { - build_pq_bytes = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - "Missing PQ bytes".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--use_opq" => { - use_opq = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "use_opq".to_string(), - "Missing use_opq flag".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "use_opq".to_string(), - format!("ParseBoolError: {}", err), - ) - })?; - } - _ => { - return Err(ANNError::log_index_config_error( - String::from(""), - format!("Unknown argument: {}", arg), - )); - } - } - } - - if data_type.is_empty() - || dist_fn.is_empty() - || data_path.is_empty() - || index_path_prefix.is_empty() - { - return Err(ANNError::log_index_config_error( - String::from(""), - "Missing required arguments".to_string(), - )); - } - - _use_pq_build = build_pq_bytes > 0; - - let metric = dist_fn - .parse::() - .map_err(|err| ANNError::log_index_config_error( - "dist_fn".to_string(), - err.to_string(), - ))?; - - println!( - "Starting index build with R: {} Lbuild: {} alpha: {} #threads: {}", - r, l, alpha, num_threads - ); - - match data_type.as_str() { - "int8" => { - build_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "uint8" => { - build_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "float" => { - build_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "f16" => { - build_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - _ => { - println!("Unsupported type. Use one of int8, uint8 or float."); - return Err(ANNError::log_index_config_error("data_type".to_string(), "Invalid data type".to_string())); - } - } - - Ok(()) -} - -fn print_help() { - println!("Arguments"); - println!("--help, -h Print information on arguments"); - println!("--data_type data type (required)"); - println!("--dist_fn distance function (required)"); - println!("--data_path Input data file in bin format for initial build (required)"); - println!("--insert_path Input data file in bin format for insert (required)"); - println!("--index_path_prefix Path prefix for saving index file components (required)"); - println!("--max_degree, -R Maximum graph degree (default: 64)"); - println!("--Lbuild, -L Build complexity, higher value results in better graphs (default: 100)"); - println!("--alpha alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for denser graphs with lower diameter (default: 1.2)"); - println!("--num_threads, -T Number of threads used for building index (defaults to num of CPU logic cores)"); - println!("--build_PQ_bytes Number of PQ bytes to build the index; 0 for full precision build (default: 0)"); - println!("--use_opq Set true for OPQ compression while using PQ distance comparisons for building the index, and false for PQ compression (default: false)"); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/Cargo.toml deleted file mode 100644 index afe5e5b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/Cargo.toml +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "build_disk_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -diskann = { path = "../../diskann" } -logger = { path = "../../logger" } -vector = { path = "../../vector" } -openblas-src = { version = "0.10.8", features = ["system", "static"] } diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/src/main.rs deleted file mode 100644 index e0b6dbe..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_disk_index/src/main.rs +++ /dev/null @@ -1,377 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::env; - -use diskann::{ - common::{ANNError, ANNResult}, - index::ann_disk_index::create_disk_index, - model::{ - default_param_vals::ALPHA, - vertex::{DIM_104, DIM_128, DIM_256}, - DiskIndexBuildParameters, IndexConfiguration, IndexWriteParametersBuilder, - }, - storage::DiskIndexStorage, - utils::round_up, - utils::{load_metadata_from_file, Timer}, -}; - -use vector::{FullPrecisionDistance, Half, Metric}; - -/// The main function to build a disk index -#[allow(clippy::too_many_arguments)] -fn build_disk_index( - metric: Metric, - data_path: &str, - r: u32, - l: u32, - index_path_prefix: &str, - num_threads: u32, - search_ram_limit_gb: f64, - index_build_ram_limit_gb: f64, - num_pq_chunks: usize, - use_opq: bool, -) -> ANNResult<()> -where - T: Default + Copy + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - let disk_index_build_parameters = - DiskIndexBuildParameters::new(search_ram_limit_gb, index_build_ram_limit_gb)?; - - let index_write_parameters = IndexWriteParametersBuilder::new(l, r) - .with_saturate_graph(true) - .with_num_threads(num_threads) - .build(); - - let (data_num, data_dim) = load_metadata_from_file(data_path)?; - - let config = IndexConfiguration::new( - metric, - data_dim, - round_up(data_dim as u64, 8_u64) as usize, - data_num, - num_pq_chunks > 0, - num_pq_chunks, - use_opq, - 0, - 1f32, - index_write_parameters, - ); - let storage = DiskIndexStorage::new(data_path.to_string(), index_path_prefix.to_string())?; - let mut index = create_disk_index::(Some(disk_index_build_parameters), config, storage)?; - - let timer = Timer::new(); - - index.build("")?; - - let diff = timer.elapsed(); - println!("Indexing time: {}", diff.as_secs_f64()); - - Ok(()) -} - -fn main() -> ANNResult<()> { - let mut data_type = String::new(); - let mut dist_fn = String::new(); - let mut data_path = String::new(); - let mut index_path_prefix = String::new(); - - let mut num_threads = 0u32; - let mut r = 64u32; - let mut l = 100u32; - let mut search_ram_limit_gb = 0f64; - let mut index_build_ram_limit_gb = 0f64; - - let mut build_pq_bytes = 0u32; - let mut use_opq = false; - - let args: Vec = env::args().collect(); - let mut iter = args.iter().skip(1).peekable(); - - while let Some(arg) = iter.next() { - match arg.as_str() { - "--help" | "-h" => { - print_help(); - return Ok(()); - } - "--data_type" => { - data_type = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_type".to_string(), - "Missing data type".to_string(), - ) - })? - .to_owned(); - } - "--dist_fn" => { - dist_fn = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "dist_fn".to_string(), - "Missing distance function".to_string(), - ) - })? - .to_owned(); - } - "--data_path" => { - data_path = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "data_path".to_string(), - "Missing data path".to_string(), - ) - })? - .to_owned(); - } - "--index_path_prefix" => { - index_path_prefix = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "index_path_prefix".to_string(), - "Missing index path prefix".to_string(), - ) - })? - .to_owned(); - } - "--max_degree" | "-R" => { - r = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "max_degree".to_string(), - "Missing max degree".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "max_degree".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--Lbuild" | "-L" => { - l = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - "Missing build complexity".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "Lbuild".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--num_threads" | "-T" => { - num_threads = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "num_threads".to_string(), - "Missing number of threads".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "num_threads".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--build_PQ_bytes" => { - build_pq_bytes = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - "Missing PQ bytes".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - format!("ParseIntError: {}", err), - ) - })?; - } - "--use_opq" => { - use_opq = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "use_opq".to_string(), - "Missing use_opq flag".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "use_opq".to_string(), - format!("ParseBoolError: {}", err), - ) - })?; - } - "--search_DRAM_budget" | "-B" => { - search_ram_limit_gb = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "search_DRAM_budget".to_string(), - "Missing search_DRAM_budget flag".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "search_DRAM_budget".to_string(), - format!("ParseBoolError: {}", err), - ) - })?; - } - "--build_DRAM_budget" | "-M" => { - index_build_ram_limit_gb = iter - .next() - .ok_or_else(|| { - ANNError::log_index_config_error( - "build_DRAM_budget".to_string(), - "Missing build_DRAM_budget flag".to_string(), - ) - })? - .parse() - .map_err(|err| { - ANNError::log_index_config_error( - "build_DRAM_budget".to_string(), - format!("ParseBoolError: {}", err), - ) - })?; - } - _ => { - return Err(ANNError::log_index_config_error( - String::from(""), - format!("Unknown argument: {}", arg), - )); - } - } - } - - if data_type.is_empty() - || dist_fn.is_empty() - || data_path.is_empty() - || index_path_prefix.is_empty() - { - return Err(ANNError::log_index_config_error( - String::from(""), - "Missing required arguments".to_string(), - )); - } - - let metric = dist_fn - .parse::() - .map_err(|err| ANNError::log_index_config_error("dist_fn".to_string(), err.to_string()))?; - - println!( - "Starting index build with R: {} Lbuild: {} alpha: {} #threads: {} search_DRAM_budget: {} build_DRAM_budget: {}", - r, l, ALPHA, num_threads, search_ram_limit_gb, index_build_ram_limit_gb - ); - - let err = match data_type.as_str() { - "int8" => build_disk_index::( - metric, - &data_path, - r, - l, - &index_path_prefix, - num_threads, - search_ram_limit_gb, - index_build_ram_limit_gb, - build_pq_bytes as usize, - use_opq, - ), - "uint8" => build_disk_index::( - metric, - &data_path, - r, - l, - &index_path_prefix, - num_threads, - search_ram_limit_gb, - index_build_ram_limit_gb, - build_pq_bytes as usize, - use_opq, - ), - "float" => build_disk_index::( - metric, - &data_path, - r, - l, - &index_path_prefix, - num_threads, - search_ram_limit_gb, - index_build_ram_limit_gb, - build_pq_bytes as usize, - use_opq, - ), - "f16" => build_disk_index::( - metric, - &data_path, - r, - l, - &index_path_prefix, - num_threads, - search_ram_limit_gb, - index_build_ram_limit_gb, - build_pq_bytes as usize, - use_opq, - ), - _ => { - println!("Unsupported type. Use one of int8, uint8, float or f16."); - return Err(ANNError::log_index_config_error( - "data_type".to_string(), - "Invalid data type".to_string(), - )); - } - }; - - match err { - Ok(_) => { - println!("Index build completed successfully"); - Ok(()) - } - Err(err) => { - eprintln!("Error: {:?}", err); - Err(err) - } - } -} - -fn print_help() { - println!("Arguments"); - println!("--help, -h Print information on arguments"); - println!("--data_type data type (required)"); - println!("--dist_fn distance function (required)"); - println!("--data_path Input data file in bin format (required)"); - println!("--index_path_prefix Path prefix for saving index file components (required)"); - println!("--max_degree, -R Maximum graph degree (default: 64)"); - println!("--Lbuild, -L Build complexity, higher value results in better graphs (default: 100)"); - println!("--search_DRAM_budget Bound on the memory footprint of the index at search time in GB. Once built, the index will use up only the specified RAM limit, the rest will reside on disk"); - println!("--build_DRAM_budget Limit on the memory allowed for building the index in GB"); - println!("--num_threads, -T Number of threads used for building index (defaults to num of CPU logic cores)"); - println!("--build_PQ_bytes Number of PQ bytes to build the index; 0 for full precision build (default: 0)"); - println!("--use_opq Set true for OPQ compression while using PQ distance comparisons for building the index, and false for PQ compression (default: false)"); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/Cargo.toml deleted file mode 100644 index eb4708d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/Cargo.toml +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "build_memory_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -clap = { version = "4.3.8", features = ["derive"] } -diskann = { path = "../../diskann" } -logger = { path = "../../logger" } -vector = { path = "../../vector" } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/args.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/args.rs deleted file mode 100644 index ede31f2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/args.rs +++ /dev/null @@ -1,62 +0,0 @@ -use clap::{Args, Parser}; - -#[derive(Debug, Args)] -enum DataType { - /// Float data type. - Float, - - /// Half data type. - FP16, -} - -#[derive(Debug, Args)] -enum DistanceFunction { - /// Euclidean distance. - L2, - - /// Cosine distance. - Cosine, -} - -#[derive(Debug, Parser)] -struct BuildMemoryIndexArgs { - /// Data type of the vectors. - #[clap(long, default_value = "float")] - pub data_type: DataType, - - /// Distance function to use. - #[clap(long, default_value = "l2")] - pub dist_fn: Metric, - - /// Path to the data file. The file should be in the format specified by the `data_type` argument. - #[clap(long, short, required = true)] - pub data_path: String, - - /// Path to the index file. The index will be saved to this prefixed name. - #[clap(long, short, required = true)] - pub index_path_prefix: String, - - /// Number of max out degree from a vertex. - #[clap(long, default_value = "32")] - pub max_degree: usize, - - /// Number of candidates to consider when building out edges - #[clap(long, short default_value = "50")] - pub l_build: usize, - - /// Alpha to use to build diverse edges - #[clap(long, short default_value = "1.0")] - pub alpha: f32, - - /// Number of threads to use. - #[clap(long, short, default_value = "1")] - pub num_threads: u8, - - /// Number of PQ bytes to use. - #[clap(long, short, default_value = "8")] - pub build_pq_bytes: usize, - - /// Use opq? - #[clap(long, short, default_value = "false")] - pub use_opq: bool, -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/main.rs deleted file mode 100644 index cdccc00..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/build_memory_index/src/main.rs +++ /dev/null @@ -1,174 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use clap::{Parser, ValueEnum}; -use std::path::PathBuf; - -use diskann::{ - common::ANNResult, - index::create_inmem_index, - model::{ - vertex::{DIM_104, DIM_128, DIM_256}, - IndexConfiguration, IndexWriteParametersBuilder, - }, - utils::round_up, - utils::{load_metadata_from_file, Timer}, -}; - -use vector::{FullPrecisionDistance, Half, Metric}; - -/// The main function to build an in-memory index -#[allow(clippy::too_many_arguments)] -fn build_in_memory_index( - metric: Metric, - data_path: &str, - r: u32, - l: u32, - alpha: f32, - save_path: &str, - num_threads: u32, - _use_pq_build: bool, - _num_pq_bytes: usize, - use_opq: bool, -) -> ANNResult<()> -where - T: Default + Copy + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - let index_write_parameters = IndexWriteParametersBuilder::new(l, r) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - let (data_num, data_dim) = load_metadata_from_file(data_path)?; - - let config = IndexConfiguration::new( - metric, - data_dim, - round_up(data_dim as u64, 8_u64) as usize, - data_num, - false, - 0, - use_opq, - 0, - 1f32, - index_write_parameters, - ); - let mut index = create_inmem_index::(config)?; - - let timer = Timer::new(); - - index.build(data_path, data_num)?; - - let diff = timer.elapsed(); - - println!("Indexing time: {}", diff.as_secs_f64()); - index.save(save_path)?; - - Ok(()) -} - -fn main() -> ANNResult<()> { - let args = BuildMemoryIndexArgs::parse(); - - let _use_pq_build = args.build_pq_bytes > 0; - - println!( - "Starting index build with R: {} Lbuild: {} alpha: {} #threads: {}", - args.max_degree, args.l_build, args.alpha, args.num_threads - ); - - let err = match args.data_type { - DataType::Float => build_in_memory_index::( - args.dist_fn, - &args.data_path.to_string_lossy(), - args.max_degree, - args.l_build, - args.alpha, - &args.index_path_prefix, - args.num_threads, - _use_pq_build, - args.build_pq_bytes, - args.use_opq, - ), - DataType::FP16 => build_in_memory_index::( - args.dist_fn, - &args.data_path.to_string_lossy(), - args.max_degree, - args.l_build, - args.alpha, - &args.index_path_prefix, - args.num_threads, - _use_pq_build, - args.build_pq_bytes, - args.use_opq, - ), - }; - - match err { - Ok(_) => { - println!("Index build completed successfully"); - Ok(()) - } - Err(err) => { - eprintln!("Error: {:?}", err); - Err(err) - } - } -} - -#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)] -enum DataType { - /// Float data type. - Float, - - /// Half data type. - FP16, -} - -#[derive(Debug, Parser)] -struct BuildMemoryIndexArgs { - /// data type (required) - #[arg(long = "data_type", default_value = "float")] - pub data_type: DataType, - - /// Distance function to use. - #[arg(long = "dist_fn", default_value = "l2")] - pub dist_fn: Metric, - - /// Path to the data file. The file should be in the format specified by the `data_type` argument. - #[arg(long = "data_path", short, required = true)] - pub data_path: PathBuf, - - /// Path to the index file. The index will be saved to this prefixed name. - #[arg(long = "index_path_prefix", short, required = true)] - pub index_path_prefix: String, - - /// Number of max out degree from a vertex. - #[arg(long = "max_degree", short = 'R', default_value = "64")] - pub max_degree: u32, - - /// Number of candidates to consider when building out edges - #[arg(long = "l_build", short = 'L', default_value = "100")] - pub l_build: u32, - - /// alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for denser graphs with lower diameter - #[arg(long, short, default_value = "1.2")] - pub alpha: f32, - - /// Number of threads to use. - #[arg(long = "num_threads", short = 'T', default_value = "1")] - pub num_threads: u32, - - /// Number of PQ bytes to build the index; 0 for full precision build - #[arg(long = "build_pq_bytes", short, default_value = "0")] - pub build_pq_bytes: usize, - - /// Set true for OPQ compression while using PQ distance comparisons for building the index, and false for PQ compression - #[arg(long = "use_opq", short, default_value = "false")] - pub use_opq: bool, -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/Cargo.toml deleted file mode 100644 index 1993aab..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/Cargo.toml +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "convert_f32_to_bf16" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -half = "2.2.1" diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/src/main.rs deleted file mode 100644 index 87b4fba..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/convert_f32_to_bf16/src/main.rs +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use half::{bf16, f16}; -use std::env; -use std::fs::{File, OpenOptions}; -use std::io::{self, Read, Write, BufReader, BufWriter}; - -enum F16OrBF16 { - F16(f16), - BF16(bf16), -} - -fn main() -> io::Result<()> { - // Retrieve command-line arguments - let args: Vec = env::args().collect(); - - match args.len() { - 3|4|5|6=> {}, - _ => { - print_usage(); - std::process::exit(1); - } - } - - // Retrieve the input and output file paths from the arguments - let input_file_path = &args[1]; - let output_file_path = &args[2]; - let use_f16 = args.len() >= 4 && args[3] == "f16"; - let save_as_float = args.len() >= 5 && args[4] == "save_as_float"; - let batch_size = if args.len() >= 6 { args[5].parse::().unwrap() } else { 100000 }; - println!("use_f16: {}", use_f16); - println!("save_as_float: {}", save_as_float); - println!("batch_size: {}", batch_size); - - // Open the input file for reading - let mut input_file = BufReader::new(File::open(input_file_path)?); - - // Open the output file for writing - let mut output_file = BufWriter::new(OpenOptions::new().write(true).create(true).open(output_file_path)?); - - // Read the first 8 bytes as metadata - let mut metadata = [0; 8]; - input_file.read_exact(&mut metadata)?; - - // Write the metadata to the output file - output_file.write_all(&metadata)?; - - // Extract the number of points and dimension from the metadata - let num_points = i32::from_le_bytes(metadata[..4].try_into().unwrap()); - let dimension = i32::from_le_bytes(metadata[4..].try_into().unwrap()); - let num_batches = num_points / batch_size; - // Calculate the size of one data point in bytes - let data_point_size = (dimension * 4 * batch_size) as usize; - let mut batches_processed = 0; - let numbers_to_print = 2; - let mut numbers_printed = 0; - let mut num_fb16_wins = 0; - let mut num_f16_wins = 0; - let mut bf16_overflow = 0; - let mut f16_overflow = 0; - - // Process each data point - for _ in 0..num_batches { - // Read one data point from the input file - let mut buffer = vec![0; data_point_size]; - match input_file.read_exact(&mut buffer){ - Ok(()) => { - // Convert the float32 data to bf16 - let half_data: Vec = buffer - .chunks_exact(4) - .map(|chunk| { - let value = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); - let converted_bf16 = bf16::from_f32(value); - let converted_f16 = f16::from_f32(value); - let distance_f16 = (converted_f16.to_f32() - value).abs(); - let distance_bf16 = (converted_bf16.to_f32() - value).abs(); - - if distance_f16 < distance_bf16 { - num_f16_wins += 1; - } else { - num_fb16_wins += 1; - } - - if (converted_bf16 == bf16::INFINITY) || (converted_bf16 == bf16::NEG_INFINITY) { - bf16_overflow += 1; - } - - if (converted_f16 == f16::INFINITY) || (converted_f16 == f16::NEG_INFINITY) { - f16_overflow += 1; - } - - if numbers_printed < numbers_to_print { - numbers_printed += 1; - println!("f32 value: {} f16 value: {} | distance {}, bf16 value: {} | distance {},", - value, converted_f16, converted_f16.to_f32() - value, converted_bf16, converted_bf16.to_f32() - value); - } - - if use_f16 { - F16OrBF16::F16(converted_f16) - } else { - F16OrBF16::BF16(converted_bf16) - } - }) - .collect(); - - batches_processed += 1; - - match save_as_float { - true => { - for float_val in half_data { - match float_val { - F16OrBF16::F16(f16_val) => output_file.write_all(&f16_val.to_f32().to_le_bytes())?, - F16OrBF16::BF16(bf16_val) => output_file.write_all(&bf16_val.to_f32().to_le_bytes())?, - } - } - } - false => { - for float_val in half_data { - match float_val { - F16OrBF16::F16(f16_val) => output_file.write_all(&f16_val.to_le_bytes())?, - F16OrBF16::BF16(bf16_val) => output_file.write_all(&bf16_val.to_le_bytes())?, - } - } - } - } - - // Print the number of points processed - println!("Processed {} points out of {}", batches_processed * batch_size, num_points); - } - Err(ref e) if e.kind() == io::ErrorKind::UnexpectedEof => { - println!("Conversion completed! {} of times f16 wins | overflow count {}, {} of times bf16 wins | overflow count{}", - num_f16_wins, f16_overflow, num_fb16_wins, bf16_overflow); - break; - } - Err(err) => { - println!("Error: {}", err); - break; - } - }; - } - - Ok(()) -} - -/// Prints the usage information -fn print_usage() { - println!("Usage: program_name input_file output_file [f16] [save_as_float] [batch_size]]"); - println!("specify f16 to downscale to f16. otherwise, downscale to bf16."); - println!("specify save_as_float to downcast to f16 or bf16, and upcast to float before saving the output data. otherwise, the data will be saved as half type."); - println!("specify the batch_size as a int, the default value is 100000."); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/Cargo.toml deleted file mode 100644 index cbb4e1e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/Cargo.toml +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "load_and_insert_memory_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -diskann = { path = "../../diskann" } -logger = { path = "../../logger" } -vector = { path = "../../vector" } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/src/main.rs deleted file mode 100644 index 4168046..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/load_and_insert_memory_index/src/main.rs +++ /dev/null @@ -1,313 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::env; - -use diskann::{ - common::{ANNResult, ANNError}, - index::create_inmem_index, - utils::round_up, - model::{ - IndexWriteParametersBuilder, - IndexConfiguration, - vertex::{DIM_128, DIM_256, DIM_104} - }, - utils::{Timer, load_metadata_from_file}, -}; - -use vector::{Metric, FullPrecisionDistance, Half}; - -// The main function to build an in-memory index -#[allow(clippy::too_many_arguments)] -fn load_and_insert_in_memory_index ( - metric: Metric, - data_path: &str, - delta_path: &str, - r: u32, - l: u32, - alpha: f32, - save_path: &str, - num_threads: u32, - _use_pq_build: bool, - _num_pq_bytes: usize, - use_opq: bool -) -> ANNResult<()> -where - T: Default + Copy + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance -{ - let index_write_parameters = IndexWriteParametersBuilder::new(l, r) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - let (data_num, data_dim) = load_metadata_from_file(&format!("{}.data", data_path))?; - - let config = IndexConfiguration::new( - metric, - data_dim, - round_up(data_dim as u64, 8_u64) as usize, - data_num, - false, - 0, - use_opq, - 0, - 2.0f32, - index_write_parameters, - ); - let mut index = create_inmem_index::(config)?; - - let timer = Timer::new(); - - index.load(data_path, data_num)?; - - let diff = timer.elapsed(); - - println!("Initial indexing time: {}", diff.as_secs_f64()); - - let (delta_data_num, _) = load_metadata_from_file(delta_path)?; - - index.insert(delta_path, delta_data_num)?; - - index.save(save_path)?; - - Ok(()) -} - -fn main() -> ANNResult<()> { - let mut data_type = String::new(); - let mut dist_fn = String::new(); - let mut data_path = String::new(); - let mut insert_path = String::new(); - let mut index_path_prefix = String::new(); - - let mut num_threads = 0u32; - let mut r = 64u32; - let mut l = 100u32; - - let mut alpha = 1.2f32; - let mut build_pq_bytes = 0u32; - let mut _use_pq_build = false; - let mut use_opq = false; - - let args: Vec = env::args().collect(); - let mut iter = args.iter().skip(1).peekable(); - - while let Some(arg) = iter.next() { - match arg.as_str() { - "--help" | "-h" => { - print_help(); - return Ok(()); - } - "--data_type" => { - data_type = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "data_type".to_string(), - "Missing data type".to_string()) - )? - .to_owned(); - } - "--dist_fn" => { - dist_fn = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "dist_fn".to_string(), - "Missing distance function".to_string()) - )? - .to_owned(); - } - "--data_path" => { - data_path = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "data_path".to_string(), - "Missing data path".to_string()) - )? - .to_owned(); - } - "--insert_path" => { - insert_path = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "insert_path".to_string(), - "Missing insert path".to_string()) - )? - .to_owned(); - } - "--index_path_prefix" => { - index_path_prefix = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "index_path_prefix".to_string(), - "Missing index path prefix".to_string()))? - .to_owned(); - } - "--max_degree" | "-R" => { - r = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "max_degree".to_string(), - "Missing max degree".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "max_degree".to_string(), - format!("ParseIntError: {}", err)) - )?; - } - "--Lbuild" | "-L" => { - l = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "Lbuild".to_string(), - "Missing build complexity".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "Lbuild".to_string(), - format!("ParseIntError: {}", err)) - )?; - } - "--alpha" => { - alpha = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "alpha".to_string(), - "Missing alpha".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "alpha".to_string(), - format!("ParseFloatError: {}", err)) - )?; - } - "--num_threads" | "-T" => { - num_threads = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "num_threads".to_string(), - "Missing number of threads".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "num_threads".to_string(), - format!("ParseIntError: {}", err)) - )?; - } - "--build_PQ_bytes" => { - build_pq_bytes = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - "Missing PQ bytes".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "build_PQ_bytes".to_string(), - format!("ParseIntError: {}", err)) - )?; - } - "--use_opq" => { - use_opq = iter.next().ok_or_else(|| ANNError::log_index_config_error( - "use_opq".to_string(), - "Missing use_opq flag".to_string()))? - .parse() - .map_err(|err| ANNError::log_index_config_error( - "use_opq".to_string(), - format!("ParseBoolError: {}", err)) - )?; - } - _ => { - return Err(ANNError::log_index_config_error(String::from(""), format!("Unknown argument: {}", arg))); - } - } - } - - if data_type.is_empty() - || dist_fn.is_empty() - || data_path.is_empty() - || index_path_prefix.is_empty() - { - return Err(ANNError::log_index_config_error(String::from(""), "Missing required arguments".to_string())); - } - - _use_pq_build = build_pq_bytes > 0; - - let metric = dist_fn - .parse::() - .map_err(|err| ANNError::log_index_config_error( - "dist_fn".to_string(), - err.to_string(), - ))?; - - println!( - "Starting index build with R: {} Lbuild: {} alpha: {} #threads: {}", - r, l, alpha, num_threads - ); - - match data_type.as_str() { - "int8" => { - load_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "uint8" => { - load_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "float" => { - load_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )?; - } - "f16" => { - load_and_insert_in_memory_index::( - metric, - &data_path, - &insert_path, - r, - l, - alpha, - &index_path_prefix, - num_threads, - _use_pq_build, - build_pq_bytes as usize, - use_opq, - )? - } - _ => { - println!("Unsupported type. Use one of int8, uint8 or float."); - return Err(ANNError::log_index_config_error("data_type".to_string(), "Invalid data type".to_string())); - } - } - - Ok(()) -} - -fn print_help() { - println!("Arguments"); - println!("--help, -h Print information on arguments"); - println!("--data_type data type (required)"); - println!("--dist_fn distance function (required)"); - println!("--data_path Input data file in bin format for initial build (required)"); - println!("--insert_path Input data file in bin format for insert (required)"); - println!("--index_path_prefix Path prefix for saving index file components (required)"); - println!("--max_degree, -R Maximum graph degree (default: 64)"); - println!("--Lbuild, -L Build complexity, higher value results in better graphs (default: 100)"); - println!("--alpha alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for denser graphs with lower diameter (default: 1.2)"); - println!("--num_threads, -T Number of threads used for building index (defaults to num of CPU logic cores)"); - println!("--build_PQ_bytes Number of PQ bytes to build the index; 0 for full precision build (default: 0)"); - println!("--use_opq Set true for OPQ compression while using PQ distance comparisons for building the index, and false for PQ compression (default: false)"); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/Cargo.toml deleted file mode 100644 index cba3709..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/Cargo.toml +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "search_memory_index" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -bytemuck = "1.13.1" -diskann = { path = "../../diskann" } -num_cpus = "1.15.0" -rayon = "1.7.0" -vector = { path = "../../vector" } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/main.rs deleted file mode 100644 index ca4d4cd..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/main.rs +++ /dev/null @@ -1,430 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod search_index_utils; -use bytemuck::Pod; -use diskann::{ - common::{ANNError, ANNResult}, - index, - model::{ - configuration::index_write_parameters::{default_param_vals, IndexWriteParametersBuilder}, - vertex::{DIM_104, DIM_128, DIM_256}, - IndexConfiguration, - }, - utils::{load_metadata_from_file, save_bin_u32}, -}; -use std::{env, path::Path, process::exit, time::Instant}; -use vector::{FullPrecisionDistance, Half, Metric}; - -use rayon::prelude::*; - -#[allow(clippy::too_many_arguments)] -fn search_memory_index( - metric: Metric, - index_path: &str, - result_path_prefix: &str, - query_file: &str, - truthset_file: &str, - num_threads: u32, - recall_at: u32, - print_all_recalls: bool, - l_vec: &Vec, - show_qps_per_thread: bool, - fail_if_recall_below: f32, -) -> ANNResult -where - T: Default + Copy + Sized + Pod + Sync + Send + Into, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - // Load the query file - let (query, query_num, query_dim, query_aligned_dim) = - search_index_utils::load_aligned_bin::(query_file)?; - let mut gt_dim: usize = 0; - let mut gt_ids: Option> = None; - let mut gt_dists: Option> = None; - - // Check for ground truth - let mut calc_recall_flag = false; - if !truthset_file.is_empty() && Path::new(truthset_file).exists() { - let ret = search_index_utils::load_truthset(truthset_file)?; - gt_ids = Some(ret.0); - gt_dists = ret.1; - let gt_num = ret.2; - gt_dim = ret.3; - - if gt_num != query_num { - println!("Error. Mismatch in number of queries and ground truth data"); - } - - calc_recall_flag = true; - } else { - println!( - "Truthset file {} not found. Not computing recall", - truthset_file - ); - } - - let num_frozen_pts = search_index_utils::get_graph_num_frozen_points(index_path)?; - - // C++ uses the max given L value, so we do the same here. Max degree is never specified in C++ so use the rust default - let index_write_params = IndexWriteParametersBuilder::new( - *l_vec.iter().max().unwrap(), - default_param_vals::MAX_DEGREE, - ) - .with_num_threads(num_threads) - .build(); - - let (index_num_points, _) = load_metadata_from_file(&format!("{}.data", index_path))?; - - let index_config = IndexConfiguration::new( - metric, - query_dim, - query_aligned_dim, - index_num_points, - false, - 0, - false, - num_frozen_pts, - 1f32, - index_write_params, - ); - let mut index = index::create_inmem_index::(index_config)?; - - index.load(index_path, index_num_points)?; - - println!("Using {} threads to search", num_threads); - let qps_title = if show_qps_per_thread { - "QPS/thread" - } else { - "QPS" - }; - let mut table_width = 4 + 12 + 18 + 20 + 15; - let mut table_header_str = format!( - "{:>4}{:>12}{:>18}{:>20}{:>15}", - "Ls", qps_title, "Avg dist cmps", "Mean Latency (mus)", "99.9 Latency" - ); - - let first_recall: u32 = if print_all_recalls { 1 } else { recall_at }; - let mut recalls_to_print: usize = 0; - if calc_recall_flag { - for curr_recall in first_recall..=recall_at { - let recall_str = format!("Recall@{}", curr_recall); - table_header_str.push_str(&format!("{:>12}", recall_str)); - recalls_to_print = (recall_at + 1 - first_recall) as usize; - table_width += recalls_to_print * 12; - } - } - - println!("{}", table_header_str); - println!("{}", "=".repeat(table_width)); - - let mut query_result_ids: Vec> = - vec![vec![0; query_num * recall_at as usize]; l_vec.len()]; - let mut latency_stats: Vec = vec![0.0; query_num]; - let mut cmp_stats: Vec = vec![0; query_num]; - let mut best_recall = 0.0; - - std::env::set_var("RAYON_NUM_THREADS", num_threads.to_string()); - - for test_id in 0..l_vec.len() { - let l_value = l_vec[test_id]; - - if l_value < recall_at { - println!( - "Ignoring search with L:{} since it's smaller than K:{}", - l_value, recall_at - ); - continue; - } - - let zipped = cmp_stats - .par_iter_mut() - .zip(latency_stats.par_iter_mut()) - .zip(query_result_ids[test_id].par_chunks_mut(recall_at as usize)) - .zip(query.par_chunks(query_aligned_dim)); - - let start = Instant::now(); - zipped.for_each(|(((cmp, latency), query_result), query_chunk)| { - let query_start = Instant::now(); - *cmp = index - .search(query_chunk, recall_at as usize, l_value, query_result) - .unwrap(); - - let query_end = Instant::now(); - let diff = query_end.duration_since(query_start); - *latency = diff.as_micros() as f32; - }); - let diff = Instant::now().duration_since(start); - - let mut displayed_qps: f32 = query_num as f32 / diff.as_secs_f32(); - if show_qps_per_thread { - displayed_qps /= num_threads as f32; - } - - let mut recalls: Vec = Vec::new(); - if calc_recall_flag { - recalls.reserve(recalls_to_print); - for curr_recall in first_recall..=recall_at { - recalls.push(search_index_utils::calculate_recall( - query_num, - gt_ids.as_ref().unwrap(), - >_dists, - gt_dim, - &query_result_ids[test_id], - recall_at, - curr_recall, - )? as f32); - } - } - - latency_stats.sort_by(|a, b| a.partial_cmp(b).unwrap()); - let mean_latency = latency_stats.iter().sum::() / query_num as f32; - let avg_cmps = cmp_stats.iter().sum::() as f32 / query_num as f32; - - let mut stat_str = format!( - "{: >4}{: >12.2}{: >18.2}{: >20.2}{: >15.2}", - l_value, - displayed_qps, - avg_cmps, - mean_latency, - latency_stats[(0.999 * query_num as f32).round() as usize] - ); - - for recall in recalls.iter() { - stat_str.push_str(&format!("{: >12.2}", recall)); - best_recall = f32::max(best_recall, *recall); - } - - println!("{}", stat_str); - } - - println!("Done searching. Now saving results"); - for (test_id, l_value) in l_vec.iter().enumerate() { - if *l_value < recall_at { - println!( - "Ignoring all search with L: {} since it's smaller than K: {}", - l_value, recall_at - ); - } - - let cur_result_path = format!("{}_{}_idx_uint32.bin", result_path_prefix, l_value); - save_bin_u32( - &cur_result_path, - query_result_ids[test_id].as_slice(), - query_num, - recall_at as usize, - 0, - )?; - } - - if best_recall >= fail_if_recall_below { - Ok(0) - } else { - Ok(-1) - } -} - -fn main() -> ANNResult<()> { - let return_val: i32; - { - let mut data_type: String = String::new(); - let mut metric: Option = None; - let mut index_path: String = String::new(); - let mut result_path_prefix: String = String::new(); - let mut query_file: String = String::new(); - let mut truthset_file: String = String::new(); - let mut num_cpus: u32 = num_cpus::get() as u32; - let mut recall_at: Option = None; - let mut print_all_recalls: bool = false; - let mut l_vec: Vec = Vec::new(); - let mut show_qps_per_thread: bool = false; - let mut fail_if_recall_below: f32 = 0.0; - - let args: Vec = env::args().collect(); - let mut iter = args.iter().skip(1).peekable(); - while let Some(arg) = iter.next() { - let ann_error = - || ANNError::log_index_config_error(String::from(arg), format!("Missing {}", arg)); - match arg.as_str() { - "--help" | "-h" => { - print_help(); - return Ok(()); - } - "--data_type" => { - data_type = iter.next().ok_or_else(ann_error)?.to_owned(); - } - "--dist_fn" => { - metric = Some(iter.next().ok_or_else(ann_error)?.parse().map_err(|err| { - ANNError::log_index_config_error( - String::from(arg), - format!("ParseError: {}", err), - ) - })?); - } - "--index_path_prefix" => { - index_path = iter.next().ok_or_else(ann_error)?.to_owned(); - } - "--result_path" => { - result_path_prefix = iter.next().ok_or_else(ann_error)?.to_owned(); - } - "--query_file" => { - query_file = iter.next().ok_or_else(ann_error)?.to_owned(); - } - "--gt_file" => { - truthset_file = iter.next().ok_or_else(ann_error)?.to_owned(); - } - "--recall_at" | "-K" => { - recall_at = - Some(iter.next().ok_or_else(ann_error)?.parse().map_err(|err| { - ANNError::log_index_config_error( - String::from(arg), - format!("ParseError: {}", err), - ) - })?); - } - "--print_all_recalls" => { - print_all_recalls = true; - } - "--search_list" | "-L" => { - while iter.peek().is_some() && !iter.peek().unwrap().starts_with('-') { - l_vec.push(iter.next().ok_or_else(ann_error)?.parse().map_err(|err| { - ANNError::log_index_config_error( - String::from(arg), - format!("ParseError: {}", err), - ) - })?); - } - } - "--num_threads" => { - num_cpus = iter.next().ok_or_else(ann_error)?.parse().map_err(|err| { - ANNError::log_index_config_error( - String::from(arg), - format!("ParseError: {}", err), - ) - })?; - } - "--qps_per_thread" => { - show_qps_per_thread = true; - } - "--fail_if_recall_below" => { - fail_if_recall_below = - iter.next().ok_or_else(ann_error)?.parse().map_err(|err| { - ANNError::log_index_config_error( - String::from(arg), - format!("ParseError: {}", err), - ) - })?; - } - _ => { - return Err(ANNError::log_index_error(format!( - "Unknown argument: {}", - arg - ))); - } - } - } - - if metric.is_none() { - return Err(ANNError::log_index_error(String::from("No metric given!"))); - } else if recall_at.is_none() { - return Err(ANNError::log_index_error(String::from( - "No recall_at given!", - ))); - } - - // Seems like float is the only supported data type for FullPrecisionDistance right now, - // but keep the structure in place here for future data types - match data_type.as_str() { - "float" => { - return_val = search_memory_index::( - metric.unwrap(), - &index_path, - &result_path_prefix, - &query_file, - &truthset_file, - num_cpus, - recall_at.unwrap(), - print_all_recalls, - &l_vec, - show_qps_per_thread, - fail_if_recall_below, - )?; - } - "int8" => { - return_val = search_memory_index::( - metric.unwrap(), - &index_path, - &result_path_prefix, - &query_file, - &truthset_file, - num_cpus, - recall_at.unwrap(), - print_all_recalls, - &l_vec, - show_qps_per_thread, - fail_if_recall_below, - )?; - } - "uint8" => { - return_val = search_memory_index::( - metric.unwrap(), - &index_path, - &result_path_prefix, - &query_file, - &truthset_file, - num_cpus, - recall_at.unwrap(), - print_all_recalls, - &l_vec, - show_qps_per_thread, - fail_if_recall_below, - )?; - } - "f16" => { - return_val = search_memory_index::( - metric.unwrap(), - &index_path, - &result_path_prefix, - &query_file, - &truthset_file, - num_cpus, - recall_at.unwrap(), - print_all_recalls, - &l_vec, - show_qps_per_thread, - fail_if_recall_below, - )?; - } - _ => { - return Err(ANNError::log_index_error(format!( - "Unknown data type: {}!", - data_type - ))); - } - } - } - - // Rust only allows returning values with this method, but this will immediately terminate the program without running destructors on the - // stack. To get around this enclose main function logic in a block so that by the time we return here all destructors have been called. - exit(return_val); -} - -fn print_help() { - println!("Arguments"); - println!("--help, -h Print information on arguments"); - println!("--data_type data type (required)"); - println!("--dist_fn distance function (required)"); - println!("--index_path_prefix Path prefix to the index (required)"); - println!("--result_path Path prefix for saving results of the queries (required)"); - println!("--query_file Query file in binary format"); - println!("--gt_file Ground truth file for the queryset"); - println!("--recall_at, -K Number of neighbors to be returned"); - println!("--print_all_recalls Print recalls at all positions, from 1 up to specified recall_at value"); - println!("--search_list List of L values of search"); - println!("----num_threads, -T Number of threads used for building index (defaults to num_cpus::get())"); - println!("--qps_per_thread Print overall QPS divided by the number of threads in the output table"); - println!("--fail_if_recall_below If set to a value >0 and <100%, program returns -1 if best recall found is below this threshold"); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/search_index_utils.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/search_index_utils.rs deleted file mode 100644 index c7b04a4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/cmd_drivers/search_memory_index/src/search_index_utils.rs +++ /dev/null @@ -1,186 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use bytemuck::{cast_slice, Pod}; -use diskann::{ - common::{ANNError, ANNResult, AlignedBoxWithSlice}, - model::data_store::DatasetDto, - utils::{copy_aligned_data_from_file, is_aligned, round_up}, -}; -use std::collections::HashSet; -use std::fs::File; -use std::io::Read; -use std::mem::size_of; - -pub(crate) fn calculate_recall( - num_queries: usize, - gold_std: &[u32], - gs_dist: &Option>, - dim_gs: usize, - our_results: &[u32], - dim_or: u32, - recall_at: u32, -) -> ANNResult { - let mut total_recall: f64 = 0.0; - let (mut gt, mut res): (HashSet, HashSet) = (HashSet::new(), HashSet::new()); - - for i in 0..num_queries { - gt.clear(); - res.clear(); - - let gt_slice = &gold_std[dim_gs * i..]; - let res_slice = &our_results[dim_or as usize * i..]; - let mut tie_breaker = recall_at as usize; - - if gs_dist.is_some() { - tie_breaker = (recall_at - 1) as usize; - let gt_dist_vec = &gs_dist.as_ref().unwrap()[dim_gs * i..]; - while tie_breaker < dim_gs - && gt_dist_vec[tie_breaker] == gt_dist_vec[(recall_at - 1) as usize] - { - tie_breaker += 1; - } - } - - (0..tie_breaker).for_each(|idx| { - gt.insert(gt_slice[idx]); - }); - - (0..tie_breaker).for_each(|idx| { - res.insert(res_slice[idx]); - }); - - let mut cur_recall: u32 = 0; - for v in gt.iter() { - if res.contains(v) { - cur_recall += 1; - } - } - - total_recall += cur_recall as f64; - } - - Ok(total_recall / num_queries as f64 * (100.0 / recall_at as f64)) -} - -pub(crate) fn get_graph_num_frozen_points(graph_file: &str) -> ANNResult { - let mut file = File::open(graph_file)?; - let mut usize_buffer = [0; size_of::()]; - let mut u32_buffer = [0; size_of::()]; - - file.read_exact(&mut usize_buffer)?; - file.read_exact(&mut u32_buffer)?; - file.read_exact(&mut u32_buffer)?; - file.read_exact(&mut usize_buffer)?; - let file_frozen_pts = usize::from_le_bytes(usize_buffer); - - Ok(file_frozen_pts) -} - -#[inline] -pub(crate) fn load_truthset( - bin_file: &str, -) -> ANNResult<(Vec, Option>, usize, usize)> { - let mut file = File::open(bin_file)?; - let actual_file_size = file.metadata()?.len() as usize; - - let mut buffer = [0; size_of::()]; - file.read_exact(&mut buffer)?; - let npts = i32::from_le_bytes(buffer) as usize; - - file.read_exact(&mut buffer)?; - let dim = i32::from_le_bytes(buffer) as usize; - - println!("Metadata: #pts = {npts}, #dims = {dim}... "); - - let expected_file_size_with_dists: usize = - 2 * npts * dim * size_of::() + 2 * size_of::(); - let expected_file_size_just_ids: usize = npts * dim * size_of::() + 2 * size_of::(); - - let truthset_type : i32 = match actual_file_size - { - // This is in the C++ code, but nothing is done in this case. Keeping it here for future reference just in case. - // expected_file_size_just_ids => 2, - x if x == expected_file_size_with_dists => 1, - _ => return Err(ANNError::log_index_error(format!("Error. File size mismatch. File should have bin format, with npts followed by ngt - followed by npts*ngt ids and optionally followed by npts*ngt distance values; actual size: {}, expected: {} or {}", - actual_file_size, - expected_file_size_with_dists, - expected_file_size_just_ids))) - }; - - let mut ids: Vec = vec![0; npts * dim]; - let mut buffer = vec![0; npts * dim * size_of::()]; - file.read_exact(&mut buffer)?; - ids.clone_from_slice(cast_slice::(&buffer)); - - if truthset_type == 1 { - let mut dists: Vec = vec![0.0; npts * dim]; - let mut buffer = vec![0; npts * dim * size_of::()]; - file.read_exact(&mut buffer)?; - dists.clone_from_slice(cast_slice::(&buffer)); - - return Ok((ids, Some(dists), npts, dim)); - } - - Ok((ids, None, npts, dim)) -} - -#[inline] -pub(crate) fn load_aligned_bin( - bin_file: &str, -) -> ANNResult<(AlignedBoxWithSlice, usize, usize, usize)> { - let t_size = size_of::(); - let (npts, dim, file_size): (usize, usize, usize); - { - println!("Reading (with alignment) bin file: {bin_file}"); - let mut file = File::open(bin_file)?; - file_size = file.metadata()?.len() as usize; - - let mut buffer = [0; size_of::()]; - file.read_exact(&mut buffer)?; - npts = i32::from_le_bytes(buffer) as usize; - - file.read_exact(&mut buffer)?; - dim = i32::from_le_bytes(buffer) as usize; - } - - let rounded_dim = round_up(dim, 8); - let expected_actual_file_size = npts * dim * size_of::() + 2 * size_of::(); - - if file_size != expected_actual_file_size { - return Err(ANNError::log_index_error(format!( - "ERROR: File size mismatch. Actual size is {} while expected size is {} - npts = {}, #dims = {}, aligned_dim = {}", - file_size, expected_actual_file_size, npts, dim, rounded_dim - ))); - } - - println!("Metadata: #pts = {npts}, #dims = {dim}, aligned_dim = {rounded_dim}..."); - - let alloc_size = npts * rounded_dim; - let alignment = 8 * t_size; - println!( - "allocating aligned memory of {} bytes... ", - alloc_size * t_size - ); - if !is_aligned(alloc_size * t_size, alignment) { - return Err(ANNError::log_index_error(format!( - "Requested memory size is not a multiple of {}. Can not be allocated.", - alignment - ))); - } - - let mut data = AlignedBoxWithSlice::::new(alloc_size, alignment)?; - let dto = DatasetDto { - data: &mut data, - rounded_dim, - }; - - println!("done. Copying data to mem_aligned buffer..."); - - let (_, _) = copy_aligned_data_from_file(bin_file, dto, 0)?; - - Ok((data, npts, dim, rounded_dim)) -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/Cargo.toml deleted file mode 100644 index a5be547..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/Cargo.toml +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "diskann" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -bincode = "1.3.3" -bit-vec = "0.6.3" -byteorder = "1.4.3" -cblas = "0.4.0" -crossbeam = "0.8.2" -half = "2.2.1" -hashbrown = "0.13.2" -num-traits = "0.2.15" -once_cell = "1.17.1" -openblas-src = { version = "0.10.8", features = ["system"] } -rand = { version = "0.8.5", features = [ "small_rng" ] } -rayon = "1.7.0" -serde = { version = "1.0.130", features = ["derive"] } -thiserror = "1.0.40" -winapi = { version = "0.3.9", features = ["errhandlingapi", "fileapi", "ioapiset", "handleapi", "winnt", "minwindef", "basetsd", "winerror", "winbase"] } - -logger = { path = "../logger" } -platform = { path = "../platform" } -vector = { path = "../vector" } - -[build-dependencies] -cc = "1.0.79" - -[dev-dependencies] -approx = "0.5.1" -criterion = "0.5.1" - - -[[bench]] -name = "distance_bench" -harness = false - -[[bench]] -name = "neighbor_bench" -harness = false diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/distance_bench.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/distance_bench.rs deleted file mode 100644 index 885c95b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/distance_bench.rs +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use criterion::{black_box, criterion_group, criterion_main, Criterion}; - -use rand::{thread_rng, Rng}; -use vector::{FullPrecisionDistance, Metric}; - -// make sure the vector is 256-bit (32 bytes) aligned required by _mm256_load_ps -#[repr(C, align(32))] -struct Vector32ByteAligned { - v: [f32; 256], -} - -fn benchmark_l2_distance_float_rust(c: &mut Criterion) { - let (a, b) = prepare_random_aligned_vectors(); - let mut group = c.benchmark_group("avx-computation"); - group.sample_size(5000); - - group.bench_function("AVX Rust run", |f| { - f.iter(|| { - black_box(<[f32; 256]>::distance_compare( - black_box(&a.v), - black_box(&b.v), - Metric::L2, - )) - }) - }); -} - -// make sure the vector is 256-bit (32 bytes) aligned required by _mm256_load_ps -fn prepare_random_aligned_vectors() -> (Box, Box) { - let a = Box::new(Vector32ByteAligned { - v: [(); 256].map(|_| thread_rng().gen_range(0.0..100.0)), - }); - - let b = Box::new(Vector32ByteAligned { - v: [(); 256].map(|_| thread_rng().gen_range(0.0..100.0)), - }); - - (a, b) -} - -criterion_group!(benches, benchmark_l2_distance_float_rust,); -criterion_main!(benches); - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/kmeans_bench.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/kmeans_bench.rs deleted file mode 100644 index c69c16a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/kmeans_bench.rs +++ /dev/null @@ -1,70 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use criterion::{criterion_group, criterion_main, Criterion}; -use diskann::utils::k_means_clustering; -use rand::Rng; - -const NUM_POINTS: usize = 10000; -const DIM: usize = 100; -const NUM_CENTERS: usize = 256; -const MAX_KMEANS_REPS: usize = 12; - -fn benchmark_kmeans_rust(c: &mut Criterion) { - let mut rng = rand::thread_rng(); - let data: Vec = (0..NUM_POINTS * DIM) - .map(|_| rng.gen_range(-1.0..1.0)) - .collect(); - let centers: Vec = vec![0.0; NUM_CENTERS * DIM]; - - let mut group = c.benchmark_group("kmeans-computation"); - group.sample_size(500); - - group.bench_function("K-Means Rust run", |f| { - f.iter(|| { - // let mut centers_copy = centers.clone(); - let data_copy = data.clone(); - let mut centers_copy = centers.clone(); - k_means_clustering( - &data_copy, - NUM_POINTS, - DIM, - &mut centers_copy, - NUM_CENTERS, - MAX_KMEANS_REPS, - ) - }) - }); -} - -fn benchmark_kmeans_c(c: &mut Criterion) { - let mut rng = rand::thread_rng(); - let data: Vec = (0..NUM_POINTS * DIM) - .map(|_| rng.gen_range(-1.0..1.0)) - .collect(); - let centers: Vec = vec![0.0; NUM_CENTERS * DIM]; - - let mut group = c.benchmark_group("kmeans-computation"); - group.sample_size(500); - - group.bench_function("K-Means C++ Run", |f| { - f.iter(|| { - let data_copy = data.clone(); - let mut centers_copy = centers.clone(); - let _ = k_means_clustering( - data_copy.as_slice(), - NUM_POINTS, - DIM, - centers_copy.as_mut_slice(), - NUM_CENTERS, - MAX_KMEANS_REPS, - ); - }) - }); -} - -criterion_group!(benches, benchmark_kmeans_rust, benchmark_kmeans_c); - -criterion_main!(benches); - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/neighbor_bench.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/neighbor_bench.rs deleted file mode 100644 index 958acdc..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/neighbor_bench.rs +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::time::Duration; - -use criterion::{black_box, criterion_group, criterion_main, Criterion}; - -use diskann::model::{Neighbor, NeighborPriorityQueue}; -use rand::distributions::{Distribution, Uniform}; -use rand::rngs::StdRng; -use rand::SeedableRng; - -fn benchmark_priority_queue_insert(c: &mut Criterion) { - let vec = generate_random_floats(); - let mut group = c.benchmark_group("neighborqueue-insert"); - group.measurement_time(Duration::from_secs(3)).sample_size(500); - - let mut queue = NeighborPriorityQueue::with_capacity(64_usize); - group.bench_function("Neighbor Priority Queue Insert", |f| { - f.iter(|| { - queue.clear(); - for n in vec.iter() { - queue.insert(*n); - } - - black_box(&1) - }); - }); -} - -fn generate_random_floats() -> Vec { - let seed: [u8; 32] = [73; 32]; - let mut rng: StdRng = SeedableRng::from_seed(seed); - let range = Uniform::new(0.0, 1.0); - let mut random_floats = Vec::with_capacity(100); - - for i in 0..100 { - let random_float = range.sample(&mut rng) as f32; - let n = Neighbor::new(i, random_float); - random_floats.push(n); - } - - random_floats -} - -criterion_group!(benches, benchmark_priority_queue_insert); -criterion_main!(benches); - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/mod.rs deleted file mode 100644 index 87e377c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod search; - -pub mod prune; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/mod.rs deleted file mode 100644 index 4627eeb..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -pub mod prune; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/prune.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/prune.rs deleted file mode 100644 index 40fec4a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/prune/prune.rs +++ /dev/null @@ -1,288 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use hashbrown::HashSet; -use vector::{FullPrecisionDistance, Metric}; - -use crate::common::{ANNError, ANNResult}; -use crate::index::InmemIndex; -use crate::model::graph::AdjacencyList; -use crate::model::neighbor::SortedNeighborVector; -use crate::model::scratch::InMemQueryScratch; -use crate::model::Neighbor; - -impl InmemIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - /// A method that occludes a list of neighbors based on some criteria - #[allow(clippy::too_many_arguments)] - fn occlude_list( - &self, - location: u32, - pool: &mut SortedNeighborVector, - alpha: f32, - degree: u32, - max_candidate_size: usize, - result: &mut AdjacencyList, - scratch: &mut InMemQueryScratch, - delete_set_ptr: Option<&HashSet>, - ) -> ANNResult<()> { - if pool.is_empty() { - return Ok(()); - } - - if !result.is_empty() { - return Err(ANNError::log_index_error( - "result is not empty.".to_string(), - )); - } - - // Truncate pool at max_candidate_size and initialize scratch spaces - if pool.len() > max_candidate_size { - pool.truncate(max_candidate_size); - } - - let occlude_factor = &mut scratch.occlude_factor; - - // occlude_list can be called with the same scratch more than once by - // search_for_point_and_add_link through inter_insert. - occlude_factor.clear(); - - // Initialize occlude_factor to pool.len() many 0.0 values for correctness - occlude_factor.resize(pool.len(), 0.0); - - let mut cur_alpha = 1.0; - while cur_alpha <= alpha && result.len() < degree as usize { - for (i, neighbor) in pool.iter().enumerate() { - if result.len() >= degree as usize { - break; - } - if occlude_factor[i] > cur_alpha { - continue; - } - // Set the entry to f32::MAX so that is not considered again - occlude_factor[i] = f32::MAX; - - // Add the entry to the result if its not been deleted, and doesn't - // add a self loop - if delete_set_ptr.map_or(true, |delete_set| !delete_set.contains(&neighbor.id)) - && neighbor.id != location - { - result.push(neighbor.id); - } - - // Update occlude factor for points from i+1 to pool.len() - for (j, neighbor2) in pool.iter().enumerate().skip(i + 1) { - if occlude_factor[j] > alpha { - continue; - } - - // todo - self.filtered_index - let djk = self.get_distance(neighbor2.id, neighbor.id)?; - match self.configuration.dist_metric { - Metric::L2 | Metric::Cosine => { - occlude_factor[j] = if djk == 0.0 { - f32::MAX - } else { - occlude_factor[j].max(neighbor2.distance / djk) - }; - } - } - } - } - - cur_alpha *= 1.2; - } - - Ok(()) - } - - /// Prunes the neighbors of a given data point based on some criteria and returns a list of pruned ids. - /// - /// # Arguments - /// - /// * `location` - The id of the data point whose neighbors are to be pruned. - /// * `pool` - A vector of neighbors to be pruned, sorted by distance to the query point. - /// * `pruned_list` - A vector to store the ids of the pruned neighbors. - /// * `scratch` - A mutable reference to a scratch space for in-memory queries. - /// - /// # Panics - /// - /// Panics if `pruned_list` contains more than `range` elements after pruning. - pub fn prune_neighbors( - &self, - location: u32, - pool: &mut Vec, - pruned_list: &mut AdjacencyList, - scratch: &mut InMemQueryScratch, - ) -> ANNResult<()> { - self.robust_prune( - location, - pool, - self.configuration.index_write_parameter.max_degree, - self.configuration.index_write_parameter.max_occlusion_size, - self.configuration.index_write_parameter.alpha, - pruned_list, - scratch, - ) - } - - /// Prunes the neighbors of a given data point based on some criteria and returns a list of pruned ids. - /// - /// # Arguments - /// - /// * `location` - The id of the data point whose neighbors are to be pruned. - /// * `pool` - A vector of neighbors to be pruned, sorted by distance to the query point. - /// * `range` - The maximum number of neighbors to keep after pruning. - /// * `max_candidate_size` - The maximum number of candidates to consider for pruning. - /// * `alpha` - A parameter that controls the occlusion pruning strategy. - /// * `pruned_list` - A vector to store the ids of the pruned neighbors. - /// * `scratch` - A mutable reference to a scratch space for in-memory queries. - /// - /// # Error - /// - /// Return error if `pruned_list` contains more than `range` elements after pruning. - #[allow(clippy::too_many_arguments)] - fn robust_prune( - &self, - location: u32, - pool: &mut Vec, - range: u32, - max_candidate_size: u32, - alpha: f32, - pruned_list: &mut AdjacencyList, - scratch: &mut InMemQueryScratch, - ) -> ANNResult<()> { - if pool.is_empty() { - // if the pool is empty, behave like a noop - pruned_list.clear(); - return Ok(()); - } - - // If using _pq_build, over-write the PQ distances with actual distances - // todo : pq_dist - - // sort the pool based on distance to query and prune it with occlude_list - let mut pool = SortedNeighborVector::new(pool); - pruned_list.clear(); - - self.occlude_list( - location, - &mut pool, - alpha, - range, - max_candidate_size as usize, - pruned_list, - scratch, - Option::None, - )?; - - if pruned_list.len() > range as usize { - return Err(ANNError::log_index_error(format!( - "pruned_list's len {} is over range {}.", - pruned_list.len(), - range - ))); - } - - if self.configuration.index_write_parameter.saturate_graph && alpha > 1.0f32 { - for neighbor in pool.iter() { - if pruned_list.len() >= (range as usize) { - break; - } - if !pruned_list.contains(&neighbor.id) && neighbor.id != location { - pruned_list.push(neighbor.id); - } - } - } - - Ok(()) - } - - /// A method that inserts a point n into the graph of its neighbors and their neighbors, - /// pruning the graph if necessary to keep it within the specified range - /// * `n` - The index of the new point - /// * `pruned_list` is a vector of the neighbors of n that have been pruned by a previous step - /// * `range` is the target number of neighbors for each point - /// * `scratch` is a mutable reference to a scratch space that can be reused for intermediate computations - pub fn inter_insert( - &self, - n: u32, - pruned_list: &Vec, - range: u32, - scratch: &mut InMemQueryScratch, - ) -> ANNResult<()> { - // Borrow the pruned_list as a source pool of neighbors - let src_pool = pruned_list; - - if src_pool.is_empty() { - return Err(ANNError::log_index_error("src_pool is empty.".to_string())); - } - - for &vertex_id in src_pool { - // vertex is the index of a neighbor of n - // Assert that vertex is within the valid range of points - if (vertex_id as usize) - >= self.configuration.max_points + self.configuration.num_frozen_pts - { - return Err(ANNError::log_index_error(format!( - "vertex_id {} is out of valid range of points {}", - vertex_id, - self.configuration.max_points + self.configuration.num_frozen_pts, - ))); - } - - let neighbors = self.add_to_neighbors(vertex_id, n, range)?; - - if let Some(copy_of_neighbors) = neighbors { - // Pruning is needed, create a dummy set and a dummy vector to store the unique neighbors of vertex_id - let mut dummy_pool = self.get_unique_neighbors(©_of_neighbors, vertex_id)?; - - // Create a new vector to store the pruned neighbors of vertex_id - let mut new_out_neighbors = - AdjacencyList::for_range(self.configuration.write_range()); - // Prune the neighbors of vertex_id using a helper method - self.prune_neighbors(vertex_id, &mut dummy_pool, &mut new_out_neighbors, scratch)?; - - self.set_neighbors(vertex_id, new_out_neighbors)?; - } - } - - Ok(()) - } - - /// Adds a node to the list of neighbors for the given node. - /// - /// # Arguments - /// - /// * `vertex_id` - The ID of the node to add the neighbor to. - /// * `node_id` - The ID of the node to add. - /// * `range` - The range of the graph. - /// - /// # Return - /// - /// Returns `None` if the node is already in the list of neighbors, or a `Vec` containing the updated list of neighbors if the list of neighbors is full. - fn add_to_neighbors( - &self, - vertex_id: u32, - node_id: u32, - range: u32, - ) -> ANNResult>> { - // vertex contains a vector of the neighbors of vertex_id - let mut vertex_guard = self.final_graph.write_vertex_and_neighbors(vertex_id)?; - - Ok(vertex_guard.add_to_neighbors(node_id, range)) - } - - fn set_neighbors(&self, vertex_id: u32, new_out_neighbors: AdjacencyList) -> ANNResult<()> { - // vertex contains a vector of the neighbors of vertex_id - let mut vertex_guard = self.final_graph.write_vertex_and_neighbors(vertex_id)?; - - vertex_guard.set_neighbors(new_out_neighbors); - Ok(()) - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/mod.rs deleted file mode 100644 index 9f007ab..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -pub mod search; - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/search.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/search.rs deleted file mode 100644 index ab6d016..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/algorithm/search/search.rs +++ /dev/null @@ -1,359 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Search algorithm for index construction and query - -use crate::common::{ANNError, ANNResult}; -use crate::index::InmemIndex; -use crate::model::{scratch::InMemQueryScratch, Neighbor, Vertex}; -use hashbrown::hash_set::Entry::*; -use vector::FullPrecisionDistance; - -impl InmemIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - /// Search for query using given L value, for benchmarking purposes - /// # Arguments - /// * `query` - query vertex - /// * `scratch` - in-memory query scratch - /// * `search_list_size` - search list size to use for the benchmark - pub fn search_with_l_override( - &self, - query: &Vertex, - scratch: &mut InMemQueryScratch, - search_list_size: usize, - ) -> ANNResult { - let init_ids = self.get_init_ids()?; - self.init_graph_for_point(query, init_ids, scratch)?; - // Scratch is created using largest L val from search_memory_index, so we artifically make it smaller here - // This allows us to use the same scratch for all L values without having to rebuild the query scratch - scratch.best_candidates.set_capacity(search_list_size); - let (_, cmp) = self.greedy_search(query, scratch)?; - - Ok(cmp) - } - - /// search for point - /// # Arguments - /// * `query` - query vertex - /// * `scratch` - in-memory query scratch - /// TODO: use_filter, filteredLindex - pub fn search_for_point( - &self, - query: &Vertex, - scratch: &mut InMemQueryScratch, - ) -> ANNResult> { - let init_ids = self.get_init_ids()?; - self.init_graph_for_point(query, init_ids, scratch)?; - let (mut visited_nodes, _) = self.greedy_search(query, scratch)?; - - visited_nodes.retain(|&element| element.id != query.vertex_id()); - Ok(visited_nodes) - } - - /// Returns the locations of start point and frozen points suitable for use with iterate_to_fixed_point. - fn get_init_ids(&self) -> ANNResult> { - let mut init_ids = Vec::with_capacity(1 + self.configuration.num_frozen_pts); - init_ids.push(self.start); - - for frozen in self.configuration.max_points - ..(self.configuration.max_points + self.configuration.num_frozen_pts) - { - let frozen_u32 = frozen.try_into()?; - if frozen_u32 != self.start { - init_ids.push(frozen_u32); - } - } - - Ok(init_ids) - } - - /// Initialize graph for point - /// # Arguments - /// * `query` - query vertex - /// * `init_ids` - initial nodes from which search starts - /// * `scratch` - in-memory query scratch - /// * `search_list_size_override` - override for search list size in index config - fn init_graph_for_point( - &self, - query: &Vertex, - init_ids: Vec, - scratch: &mut InMemQueryScratch, - ) -> ANNResult<()> { - scratch - .best_candidates - .reserve(self.configuration.index_write_parameter.search_list_size as usize); - scratch.query.memcpy(query.vector())?; - - if !scratch.id_scratch.is_empty() { - return Err(ANNError::log_index_error( - "id_scratch is not empty.".to_string(), - )); - } - - let query_vertex = Vertex::::try_from((&scratch.query[..], query.vertex_id())) - .map_err(|err| { - ANNError::log_index_error(format!( - "TryFromSliceError: failed to get Vertex for query, err={}", - err - )) - })?; - - for id in init_ids { - if (id as usize) >= self.configuration.max_points + self.configuration.num_frozen_pts { - return Err(ANNError::log_index_error(format!( - "vertex_id {} is out of valid range of points {}", - id, - self.configuration.max_points + self.configuration.num_frozen_pts - ))); - } - - if let Vacant(entry) = scratch.node_visited_robinset.entry(id) { - entry.insert(); - - let vertex = self.dataset.get_vertex(id)?; - - let distance = vertex.compare(&query_vertex, self.configuration.dist_metric); - let neighbor = Neighbor::new(id, distance); - scratch.best_candidates.insert(neighbor); - } - } - - Ok(()) - } - - /// GreedySearch against query node - /// Returns visited nodes - /// # Arguments - /// * `query` - query vertex - /// * `scratch` - in-memory query scratch - /// TODO: use_filter, filter_label, search_invocation - fn greedy_search( - &self, - query: &Vertex, - scratch: &mut InMemQueryScratch, - ) -> ANNResult<(Vec, u32)> { - let mut visited_nodes = - Vec::with_capacity((3 * scratch.candidate_size + scratch.max_degree) as usize); - - // TODO: uncomment hops? - // let mut hops: u32 = 0; - let mut cmps: u32 = 0; - - let query_vertex = Vertex::::try_from((&scratch.query[..], query.vertex_id())) - .map_err(|err| { - ANNError::log_index_error(format!( - "TryFromSliceError: failed to get Vertex for query, err={}", - err - )) - })?; - - while scratch.best_candidates.has_notvisited_node() { - let closest_node = scratch.best_candidates.closest_notvisited(); - - // Add node to visited nodes to create pool for prune later - // TODO: search_invocation and use_filter - visited_nodes.push(closest_node); - - // Find which of the nodes in des have not been visited before - scratch.id_scratch.clear(); - - let max_vertex_id = self.configuration.max_points + self.configuration.num_frozen_pts; - - for id in self - .final_graph - .read_vertex_and_neighbors(closest_node.id)? - .get_neighbors() - { - let current_vertex_id = *id; - debug_assert!( - (current_vertex_id as usize) < max_vertex_id, - "current_vertex_id {} is out of valid range of points {}", - current_vertex_id, - max_vertex_id - ); - if current_vertex_id as usize >= max_vertex_id { - continue; - } - - // quickly de-dup. Remember, we are in a read lock - // we want to exit out of it quickly - if scratch.node_visited_robinset.insert(current_vertex_id) { - scratch.id_scratch.push(current_vertex_id); - } - } - - let len = scratch.id_scratch.len(); - for (m, &id) in scratch.id_scratch.iter().enumerate() { - if m + 1 < len { - let next_node = unsafe { *scratch.id_scratch.get_unchecked(m + 1) }; - self.dataset.prefetch_vector(next_node); - } - - let vertex = self.dataset.get_vertex(id)?; - let distance = query_vertex.compare(&vertex, self.configuration.dist_metric); - - // Insert pairs into the pool of candidates - scratch.best_candidates.insert(Neighbor::new(id, distance)); - } - - cmps += len as u32; - } - - Ok((visited_nodes, cmps)) - } -} - -#[cfg(test)] -mod search_test { - use vector::Metric; - - use crate::model::configuration::index_write_parameters::IndexWriteParametersBuilder; - use crate::model::graph::AdjacencyList; - use crate::model::IndexConfiguration; - use crate::test_utils::inmem_index_initialization::create_index_with_test_data; - - use super::*; - - #[test] - fn get_init_ids_no_forzen_pts() { - let index_write_parameters = IndexWriteParametersBuilder::new(50, 4) - .with_alpha(1.2) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - 256, - 256, - 256, - false, - 0, - false, - 0, - 1f32, - index_write_parameters, - ); - - let index = InmemIndex::::new(config).unwrap(); - let init_ids = index.get_init_ids().unwrap(); - assert_eq!(init_ids.len(), 1); - assert_eq!(init_ids[0], 256); - } - - #[test] - fn get_init_ids_with_forzen_pts() { - let index_write_parameters = IndexWriteParametersBuilder::new(50, 4) - .with_alpha(1.2) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - 256, - 256, - 256, - false, - 0, - false, - 2, - 1f32, - index_write_parameters, - ); - - let index = InmemIndex::::new(config).unwrap(); - let init_ids = index.get_init_ids().unwrap(); - assert_eq!(init_ids.len(), 2); - assert_eq!(init_ids[0], 256); - assert_eq!(init_ids[1], 257); - } - - #[test] - fn search_for_point_initial_call() { - let index = create_index_with_test_data(); - let query = index.dataset.get_vertex(0).unwrap(); - - let mut scratch = InMemQueryScratch::new( - index.configuration.index_write_parameter.search_list_size, - &index.configuration.index_write_parameter, - false, - ) - .unwrap(); - let visited_nodes = index.search_for_point(&query, &mut scratch).unwrap(); - assert_eq!(visited_nodes.len(), 1); - assert_eq!(scratch.best_candidates.size(), 1); - assert_eq!(scratch.best_candidates[0].id, 72); - assert_eq!(scratch.best_candidates[0].distance, 125678.0_f32); - assert!(scratch.best_candidates[0].visited); - } - - fn set_neighbors(index: &InmemIndex, vertex_id: u32, neighbors: Vec) { - index - .final_graph - .write_vertex_and_neighbors(vertex_id) - .unwrap() - .set_neighbors(AdjacencyList::from(neighbors)); - } - #[test] - fn search_for_point_works_with_edges() { - let index = create_index_with_test_data(); - let query = index.dataset.get_vertex(14).unwrap(); - - set_neighbors(&index, 0, vec![12, 72, 5, 9]); - set_neighbors(&index, 1, vec![2, 12, 10, 4]); - set_neighbors(&index, 2, vec![1, 72, 9]); - set_neighbors(&index, 3, vec![13, 6, 5, 11]); - set_neighbors(&index, 4, vec![1, 3, 7, 9]); - set_neighbors(&index, 5, vec![3, 0, 8, 11, 13]); - set_neighbors(&index, 6, vec![3, 72, 7, 10, 13]); - set_neighbors(&index, 7, vec![72, 4, 6]); - set_neighbors(&index, 8, vec![72, 5, 9, 12]); - set_neighbors(&index, 9, vec![8, 4, 0, 2]); - set_neighbors(&index, 10, vec![72, 1, 9, 6]); - set_neighbors(&index, 11, vec![3, 0, 5]); - set_neighbors(&index, 12, vec![1, 0, 8, 9]); - set_neighbors(&index, 13, vec![3, 72, 5, 6]); - set_neighbors(&index, 72, vec![7, 2, 10, 8, 13]); - - let mut scratch = InMemQueryScratch::new( - index.configuration.index_write_parameter.search_list_size, - &index.configuration.index_write_parameter, - false, - ) - .unwrap(); - let visited_nodes = index.search_for_point(&query, &mut scratch).unwrap(); - assert_eq!(visited_nodes.len(), 15); - assert_eq!(scratch.best_candidates.size(), 15); - assert_eq!(scratch.best_candidates[0].id, 2); - assert_eq!(scratch.best_candidates[0].distance, 120899.0_f32); - assert_eq!(scratch.best_candidates[1].id, 8); - assert_eq!(scratch.best_candidates[1].distance, 145538.0_f32); - assert_eq!(scratch.best_candidates[2].id, 72); - assert_eq!(scratch.best_candidates[2].distance, 146046.0_f32); - assert_eq!(scratch.best_candidates[3].id, 4); - assert_eq!(scratch.best_candidates[3].distance, 148462.0_f32); - assert_eq!(scratch.best_candidates[4].id, 7); - assert_eq!(scratch.best_candidates[4].distance, 148912.0_f32); - assert_eq!(scratch.best_candidates[5].id, 10); - assert_eq!(scratch.best_candidates[5].distance, 154570.0_f32); - assert_eq!(scratch.best_candidates[6].id, 1); - assert_eq!(scratch.best_candidates[6].distance, 159448.0_f32); - assert_eq!(scratch.best_candidates[7].id, 12); - assert_eq!(scratch.best_candidates[7].distance, 170698.0_f32); - assert_eq!(scratch.best_candidates[8].id, 9); - assert_eq!(scratch.best_candidates[8].distance, 177205.0_f32); - assert_eq!(scratch.best_candidates[9].id, 0); - assert_eq!(scratch.best_candidates[9].distance, 259996.0_f32); - assert_eq!(scratch.best_candidates[10].id, 6); - assert_eq!(scratch.best_candidates[10].distance, 371819.0_f32); - assert_eq!(scratch.best_candidates[11].id, 5); - assert_eq!(scratch.best_candidates[11].distance, 385240.0_f32); - assert_eq!(scratch.best_candidates[12].id, 3); - assert_eq!(scratch.best_candidates[12].distance, 413899.0_f32); - assert_eq!(scratch.best_candidates[13].id, 13); - assert_eq!(scratch.best_candidates[13].distance, 416386.0_f32); - assert_eq!(scratch.best_candidates[14].id, 11); - assert_eq!(scratch.best_candidates[14].distance, 449266.0_f32); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/aligned_allocator.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/aligned_allocator.rs deleted file mode 100644 index 6164a1f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/aligned_allocator.rs +++ /dev/null @@ -1,281 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -use std::alloc::Layout; -use std::ops::{Deref, DerefMut, Range}; -use std::ptr::copy_nonoverlapping; - -use super::{ANNResult, ANNError}; - -#[derive(Debug)] -/// A box that holds a slice but is aligned to the specified layout. -/// -/// This type is useful for working with types that require a certain alignment, -/// such as SIMD vectors or FFI structs. It allocates memory using the global allocator -/// and frees it when dropped. It also implements Deref and DerefMut to allow access -/// to the underlying slice. -pub struct AlignedBoxWithSlice { - /// The layout of the allocated memory. - layout: Layout, - - /// The slice that points to the allocated memory. - val: Box<[T]>, -} - -impl AlignedBoxWithSlice { - /// Creates a new `AlignedBoxWithSlice` with the given capacity and alignment. - /// The allocated memory are set to 0. - /// - /// # Error - /// - /// Return IndexError if the alignment is not a power of two or if the layout is invalid. - /// - /// This function is unsafe because it allocates uninitialized memory and casts it to - /// a slice of `T`. The caller must ensure that the capacity and alignment are valid - /// for the type `T` and that the memory is initialized before accessing the elements - /// of the slice. - pub fn new(capacity: usize, alignment: usize) -> ANNResult { - let allocsize = capacity.checked_mul(std::mem::size_of::()) - .ok_or_else(|| ANNError::log_index_error("capacity overflow".to_string()))?; - let layout = Layout::from_size_align(allocsize, alignment) - .map_err(ANNError::log_mem_alloc_layout_error)?; - - let val = unsafe { - let mem = std::alloc::alloc_zeroed(layout); - let ptr = mem as *mut T; - let slice = std::slice::from_raw_parts_mut(ptr, capacity); - std::boxed::Box::from_raw(slice) - }; - - Ok(Self { layout, val }) - } - - /// Returns a reference to the slice. - pub fn as_slice(&self) -> &[T] { - &self.val - } - - /// Returns a mutable reference to the slice. - pub fn as_mut_slice(&mut self) -> &mut [T] { - &mut self.val - } - - /// Copies data from the source slice to the destination box. - pub fn memcpy(&mut self, src: &[T]) -> ANNResult<()> { - if src.len() > self.val.len() { - return Err(ANNError::log_index_error(format!("source slice is too large (src:{}, dst:{})", src.len(), self.val.len()))); - } - - // Check that they don't overlap - let src_ptr = src.as_ptr(); - let src_end = unsafe { src_ptr.add(src.len()) }; - let dst_ptr = self.val.as_mut_ptr(); - let dst_end = unsafe { dst_ptr.add(self.val.len()) }; - - if src_ptr < dst_end && src_end > dst_ptr { - return Err(ANNError::log_index_error("Source and destination overlap".to_string())); - } - - unsafe { - copy_nonoverlapping(src.as_ptr(), self.val.as_mut_ptr(), src.len()); - } - - Ok(()) - } - - /// Split the range of memory into nonoverlapping mutable slices. - /// The number of returned slices is (range length / slice_len) and each has a length of slice_len. - pub fn split_into_nonoverlapping_mut_slices(&mut self, range: Range, slice_len: usize) -> ANNResult> { - if range.len() % slice_len != 0 || range.end > self.len() { - return Err(ANNError::log_index_error(format!( - "Cannot split range ({:?}) of AlignedBoxWithSlice (len: {}) into nonoverlapping mutable slices with length {}", - range, - self.len(), - slice_len, - ))); - } - - let mut slices = Vec::with_capacity(range.len() / slice_len); - let mut remaining_slice = &mut self.val[range]; - - while remaining_slice.len() >= slice_len { - let (left, right) = remaining_slice.split_at_mut(slice_len); - slices.push(left); - remaining_slice = right; - } - - Ok(slices) - } -} - - -impl Drop for AlignedBoxWithSlice { - /// Frees the memory allocated for the slice using the global allocator. - fn drop(&mut self) { - let val = std::mem::take(&mut self.val); - let mut val2 = std::mem::ManuallyDrop::new(val); - let ptr = val2.as_mut_ptr(); - - unsafe { - // let nonNull = NonNull::new_unchecked(ptr as *mut u8); - std::alloc::dealloc(ptr as *mut u8, self.layout) - } - } -} - -impl Deref for AlignedBoxWithSlice { - type Target = [T]; - - fn deref(&self) -> &Self::Target { - &self.val - } -} - -impl DerefMut for AlignedBoxWithSlice { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.val - } -} - -#[cfg(test)] -mod tests { - use rand::Rng; - - use crate::utils::is_aligned; - - use super::*; - - #[test] - fn create_alignedvec_works_32() { - (0..100).for_each(|_| { - let size = 1_000_000; - println!("Attempting {}", size); - let data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - assert_eq!(data.len(), size, "Capacity should match"); - - let ptr = data.as_ptr() as usize; - assert_eq!(ptr % 32, 0, "Ptr should be aligned to 32"); - - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], f32::default()); - }); - - drop(data); - }); - } - - #[test] - fn create_alignedvec_works_256() { - let mut rng = rand::thread_rng(); - - (0..100).for_each(|_| { - let n = rng.gen::(); - let size = usize::from(n) + 1; - println!("Attempting {}", size); - let data = AlignedBoxWithSlice::::new(size, 256).unwrap(); - assert_eq!(data.len(), size, "Capacity should match"); - - let ptr = data.as_ptr() as usize; - assert_eq!(ptr % 256, 0, "Ptr should be aligned to 32"); - - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], u8::default()); - }); - - drop(data); - }); - } - - #[test] - fn as_slice_test() { - let size = 1_000_000; - let data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - // assert that the slice is initialized. - (0..size).for_each(|i| { - assert_eq!(data[i], f32::default()); - }); - - let slice = data.as_slice(); - (0..size).for_each(|i| { - assert_eq!(slice[i], f32::default()); - }); - } - - #[test] - fn as_mut_slice_test() { - let size = 1_000_000; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut_slice = data.as_mut_slice(); - (0..size).for_each(|i| { - assert_eq!(mut_slice[i], f32::default()); - }); - } - - #[test] - fn memcpy_test() { - let size = 1_000_000; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut destination = AlignedBoxWithSlice::::new(size-2, 32).unwrap(); - let mut_destination = destination.as_mut_slice(); - data.memcpy(mut_destination).unwrap(); - (0..size-2).for_each(|i| { - assert_eq!(data[i], mut_destination[i]); - }); - } - - #[test] - #[should_panic(expected = "source slice is too large (src:1000000, dst:999998)")] - fn memcpy_panic_test() { - let size = 1_000_000; - let mut data = AlignedBoxWithSlice::::new(size-2, 32).unwrap(); - let mut destination = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let mut_destination = destination.as_mut_slice(); - data.memcpy(mut_destination).unwrap(); - } - - #[test] - fn is_aligned_test() { - assert!(is_aligned(256,256)); - assert!(!is_aligned(255,256)); - } - - #[test] - fn split_into_nonoverlapping_mut_slices_test() { - let size = 10; - let slice_len = 2; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let slices = data.split_into_nonoverlapping_mut_slices(2..8, slice_len).unwrap(); - assert_eq!(slices.len(), 3); - for (i, slice) in slices.into_iter().enumerate() { - assert_eq!(slice.len(), slice_len); - slice[0] = i as f32 + 1.0; - slice[1] = i as f32 + 1.0; - } - let expected_arr = [0.0f32, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 0.0, 0.0]; - assert_eq!(data.as_ref(), &expected_arr); - } - - #[test] - fn split_into_nonoverlapping_mut_slices_error_when_indivisible() { - let size = 10; - let slice_len = 2; - let range = 2..7; - let mut data = AlignedBoxWithSlice::::new(size, 32).unwrap(); - let result = data.split_into_nonoverlapping_mut_slices(range.clone(), slice_len); - let expected_err_str = format!( - "IndexError: Cannot split range ({:?}) of AlignedBoxWithSlice (len: {}) into nonoverlapping mutable slices with length {}", - range, - size, - slice_len, - ); - assert!(result.is_err_and(|e| e.to_string() == expected_err_str)); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/ann_result.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/ann_result.rs deleted file mode 100644 index 69fcf03..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/ann_result.rs +++ /dev/null @@ -1,179 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::alloc::LayoutError; -use std::array::TryFromSliceError; -use std::io; -use std::num::TryFromIntError; - -use logger::error_logger::log_error; -use logger::log_error::LogError; - -/// Result -pub type ANNResult = Result; - -/// DiskANN Error -/// ANNError is `Send` (i.e., safe to send across threads) -#[derive(thiserror::Error, Debug)] -pub enum ANNError { - /// Index construction and search error - #[error("IndexError: {err}")] - IndexError { err: String }, - - /// Index configuration error - #[error("IndexConfigError: {parameter} is invalid, err={err}")] - IndexConfigError { parameter: String, err: String }, - - /// Integer conversion error - #[error("TryFromIntError: {err}")] - TryFromIntError { - #[from] - err: TryFromIntError, - }, - - /// IO error - #[error("IOError: {err}")] - IOError { - #[from] - err: io::Error, - }, - - /// Layout error in memory allocation - #[error("MemoryAllocLayoutError: {err}")] - MemoryAllocLayoutError { - #[from] - err: LayoutError, - }, - - /// PoisonError which can be returned whenever a lock is acquired - /// Both Mutexes and RwLocks are poisoned whenever a thread fails while the lock is held - #[error("LockPoisonError: {err}")] - LockPoisonError { err: String }, - - /// DiskIOAlignmentError which can be returned when calling windows API CreateFileA for the disk index file fails. - #[error("DiskIOAlignmentError: {err}")] - DiskIOAlignmentError { err: String }, - - /// Logging error - #[error("LogError: {err}")] - LogError { - #[from] - err: LogError, - }, - - // PQ construction error - // Error happened when we construct PQ pivot or PQ compressed table - #[error("PQError: {err}")] - PQError { err: String }, - - /// Array conversion error - #[error("Error try creating array from slice: {err}")] - TryFromSliceError { - #[from] - err: TryFromSliceError, - }, -} - -impl ANNError { - /// Create, log and return IndexError - #[inline] - pub fn log_index_error(err: String) -> Self { - let ann_err = ANNError::IndexError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return IndexConfigError - #[inline] - pub fn log_index_config_error(parameter: String, err: String) -> Self { - let ann_err = ANNError::IndexConfigError { parameter, err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return TryFromIntError - #[inline] - pub fn log_try_from_int_error(err: TryFromIntError) -> Self { - let ann_err = ANNError::TryFromIntError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return IOError - #[inline] - pub fn log_io_error(err: io::Error) -> Self { - let ann_err = ANNError::IOError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return DiskIOAlignmentError - /// #[inline] - pub fn log_disk_io_request_alignment_error(err: String) -> Self { - let ann_err: ANNError = ANNError::DiskIOAlignmentError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return IOError - #[inline] - pub fn log_mem_alloc_layout_error(err: LayoutError) -> Self { - let ann_err = ANNError::MemoryAllocLayoutError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return LockPoisonError - #[inline] - pub fn log_lock_poison_error(err: String) -> Self { - let ann_err = ANNError::LockPoisonError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return PQError - #[inline] - pub fn log_pq_error(err: String) -> Self { - let ann_err = ANNError::PQError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } - - /// Create, log and return TryFromSliceError - #[inline] - pub fn log_try_from_slice_error(err: TryFromSliceError) -> Self { - let ann_err = ANNError::TryFromSliceError { err }; - match log_error(ann_err.to_string()) { - Ok(()) => ann_err, - Err(log_err) => ANNError::LogError { err: log_err }, - } - } -} - -#[cfg(test)] -mod ann_result_test { - use super::*; - - #[test] - fn ann_err_is_send() { - fn assert_send() {} - assert_send::(); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/mod.rs deleted file mode 100644 index d9da72b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/common/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod aligned_allocator; -pub use aligned_allocator::AlignedBoxWithSlice; - -mod ann_result; -pub use ann_result::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/ann_disk_index.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/ann_disk_index.rs deleted file mode 100644 index a6e053e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/ann_disk_index.rs +++ /dev/null @@ -1,54 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_docs)] - -//! ANN disk index abstraction - -use vector::FullPrecisionDistance; - -use crate::model::{IndexConfiguration, DiskIndexBuildParameters}; -use crate::storage::DiskIndexStorage; -use crate::model::vertex::{DIM_128, DIM_256, DIM_104}; - -use crate::common::{ANNResult, ANNError}; - -use super::DiskIndex; - -/// ANN disk index abstraction for custom -pub trait ANNDiskIndex : Sync + Send -where T : Default + Copy + Sync + Send + Into - { - /// Build index - fn build(&mut self, codebook_prefix: &str) -> ANNResult<()>; -} - -/// Create Index based on configuration -pub fn create_disk_index<'a, T>( - disk_build_param: Option, - config: IndexConfiguration, - storage: DiskIndexStorage, -) -> ANNResult + 'a>> -where - T: Default + Copy + Sync + Send + Into + 'a, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - match config.aligned_dim { - DIM_104 => { - let index = Box::new(DiskIndex::::new(disk_build_param, config, storage)); - Ok(index as Box>) - }, - DIM_128 => { - let index = Box::new(DiskIndex::::new(disk_build_param, config, storage)); - Ok(index as Box>) - }, - DIM_256 => { - let index = Box::new(DiskIndex::::new(disk_build_param, config, storage)); - Ok(index as Box>) - }, - _ => Err(ANNError::log_index_error(format!("Invalid dimension: {}", config.aligned_dim))), - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/disk_index.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/disk_index.rs deleted file mode 100644 index 16f0d59..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/disk_index.rs +++ /dev/null @@ -1,161 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::mem; - -use logger::logger::indexlog::DiskIndexConstructionCheckpoint; -use vector::FullPrecisionDistance; - -use crate::common::{ANNResult, ANNError}; -use crate::index::{InmemIndex, ANNInmemIndex}; -use crate::instrumentation::DiskIndexBuildLogger; -use crate::model::configuration::DiskIndexBuildParameters; -use crate::model::{IndexConfiguration, MAX_PQ_TRAINING_SET_SIZE, MAX_PQ_CHUNKS, generate_quantized_data, GRAPH_SLACK_FACTOR}; -use crate::storage::DiskIndexStorage; -use crate::utils::set_rayon_num_threads; - -use super::ann_disk_index::ANNDiskIndex; - -pub const OVERHEAD_FACTOR: f64 = 1.1f64; - -pub const MAX_SAMPLE_POINTS_FOR_WARMUP: usize = 100_000; - -pub struct DiskIndex -where - [T; N]: FullPrecisionDistance, -{ - /// Parameters for index construction - /// None for query path - disk_build_param: Option, - - configuration: IndexConfiguration, - - pub storage: DiskIndexStorage, -} - -impl DiskIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - pub fn new( - disk_build_param: Option, - configuration: IndexConfiguration, - storage: DiskIndexStorage, - ) -> Self { - Self { - disk_build_param, - configuration, - storage, - } - } - - pub fn disk_build_param(&self) -> &Option { - &self.disk_build_param - } - - pub fn index_configuration(&self) -> &IndexConfiguration { - &self.configuration - } - - fn build_inmem_index(&self, num_points: usize, data_path: &str, inmem_index_path: &str) -> ANNResult<()> { - let estimated_index_ram = self.estimate_ram_usage(num_points); - if estimated_index_ram >= self.fetch_disk_build_param()?.index_build_ram_limit() * 1024_f64 * 1024_f64 * 1024_f64 { - return Err(ANNError::log_index_error(format!( - "Insufficient memory budget for index build, index_build_ram_limit={}GB estimated_index_ram={}GB", - self.fetch_disk_build_param()?.index_build_ram_limit(), - estimated_index_ram / (1024_f64 * 1024_f64 * 1024_f64), - ))); - } - - let mut index = InmemIndex::::new(self.configuration.clone())?; - index.build(data_path, num_points)?; - index.save(inmem_index_path)?; - - Ok(()) - } - - #[inline] - fn estimate_ram_usage(&self, size: usize) -> f64 { - let degree = self.configuration.index_write_parameter.max_degree as usize; - let datasize = mem::size_of::(); - - let dataset_size = (size * N * datasize) as f64; - let graph_size = (size * degree * mem::size_of::()) as f64 * GRAPH_SLACK_FACTOR; - - OVERHEAD_FACTOR * (dataset_size + graph_size) - } - - #[inline] - fn fetch_disk_build_param(&self) -> ANNResult<&DiskIndexBuildParameters> { - self.disk_build_param - .as_ref() - .ok_or_else(|| ANNError::log_index_config_error( - "disk_build_param".to_string(), - "disk_build_param is None".to_string())) - } -} - -impl ANNDiskIndex for DiskIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - fn build(&mut self, codebook_prefix: &str) -> ANNResult<()> { - if self.configuration.index_write_parameter.num_threads > 0 { - set_rayon_num_threads(self.configuration.index_write_parameter.num_threads); - } - - println!("Starting index build: R={} L={} Query RAM budget={} Indexing RAM budget={} T={}", - self.configuration.index_write_parameter.max_degree, - self.configuration.index_write_parameter.search_list_size, - self.fetch_disk_build_param()?.search_ram_limit(), - self.fetch_disk_build_param()?.index_build_ram_limit(), - self.configuration.index_write_parameter.num_threads - ); - - let mut logger = DiskIndexBuildLogger::new(DiskIndexConstructionCheckpoint::PqConstruction); - - // PQ memory consumption = PQ pivots + PQ compressed table - // PQ pivots: dim * num_centroids * sizeof::() - // PQ compressed table: num_pts * num_pq_chunks * (dim / num_pq_chunks) * sizeof::() - // * Because num_centroids is 256, centroid id can be represented by u8 - let num_points = self.configuration.max_points; - let dim = self.configuration.dim; - let p_val = MAX_PQ_TRAINING_SET_SIZE / (num_points as f64); - let mut num_pq_chunks = ((self.fetch_disk_build_param()?.search_ram_limit() / (num_points as f64)).floor()) as usize; - num_pq_chunks = if num_pq_chunks == 0 { 1 } else { num_pq_chunks }; - num_pq_chunks = if num_pq_chunks > dim { dim } else { num_pq_chunks }; - num_pq_chunks = if num_pq_chunks > MAX_PQ_CHUNKS { MAX_PQ_CHUNKS } else { num_pq_chunks }; - - println!("Compressing {}-dimensional data into {} bytes per vector.", dim, num_pq_chunks); - - // TODO: Decouple PQ from file access - generate_quantized_data::( - p_val, - num_pq_chunks, - codebook_prefix, - self.storage.get_pq_storage(), - )?; - logger.log_checkpoint(DiskIndexConstructionCheckpoint::InmemIndexBuild)?; - - // TODO: Decouple index from file access - let inmem_index_path = self.storage.index_path_prefix().clone() + "_mem.index"; - self.build_inmem_index(num_points, self.storage.dataset_file(), inmem_index_path.as_str())?; - logger.log_checkpoint(DiskIndexConstructionCheckpoint::DiskLayout)?; - - self.storage.create_disk_layout()?; - logger.log_checkpoint(DiskIndexConstructionCheckpoint::None)?; - - let ten_percent_points = ((num_points as f64) * 0.1_f64).ceil(); - let num_sample_points = if ten_percent_points > (MAX_SAMPLE_POINTS_FOR_WARMUP as f64) { MAX_SAMPLE_POINTS_FOR_WARMUP as f64 } else { ten_percent_points }; - let sample_sampling_rate = num_sample_points / (num_points as f64); - self.storage.gen_query_warmup_data(sample_sampling_rate)?; - - self.storage.index_build_cleanup()?; - - Ok(()) - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/mod.rs deleted file mode 100644 index 4f07bd7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/disk_index/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod disk_index; -pub use disk_index::DiskIndex; - -pub mod ann_disk_index; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/ann_inmem_index.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/ann_inmem_index.rs deleted file mode 100644 index dc8dfc8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/ann_inmem_index.rs +++ /dev/null @@ -1,97 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_docs)] - -//! ANN in-memory index abstraction - -use vector::FullPrecisionDistance; - -use crate::model::{vertex::{DIM_128, DIM_256, DIM_104}, IndexConfiguration}; -use crate::common::{ANNResult, ANNError}; - -use super::InmemIndex; - -/// ANN inmem-index abstraction for custom -pub trait ANNInmemIndex : Sync + Send -where T : Default + Copy + Sync + Send + Into - { - /// Build index - fn build(&mut self, filename: &str, num_points_to_load: usize) -> ANNResult<()>; - - /// Save index - fn save(&mut self, filename: &str) -> ANNResult<()>; - - /// Load index - fn load(&mut self, filename: &str, expected_num_points: usize) -> ANNResult<()>; - - /// insert index - fn insert(&mut self, filename: &str, num_points_to_insert: usize) -> ANNResult<()>; - - /// Search the index for K nearest neighbors of query using given L value, for benchmarking purposes - fn search(&self, query : &[T], k_value : usize, l_value : u32, indices : &mut[u32]) -> ANNResult; - - /// Soft deletes the nodes with the ids in the given array. - fn soft_delete(&mut self, vertex_ids_to_delete: Vec, num_points_to_delete: usize) -> ANNResult<()>; -} - -/// Create Index based on configuration -pub fn create_inmem_index<'a, T>(config: IndexConfiguration) -> ANNResult + 'a>> -where - T: Default + Copy + Sync + Send + Into + 'a, - [T; DIM_104]: FullPrecisionDistance, - [T; DIM_128]: FullPrecisionDistance, - [T; DIM_256]: FullPrecisionDistance, -{ - match config.aligned_dim { - DIM_104 => { - let index = Box::new(InmemIndex::::new(config)?); - Ok(index as Box>) - }, - DIM_128 => { - let index = Box::new(InmemIndex::::new(config)?); - Ok(index as Box>) - }, - DIM_256 => { - let index = Box::new(InmemIndex::::new(config)?); - Ok(index as Box>) - }, - _ => Err(ANNError::log_index_error(format!("Invalid dimension: {}", config.aligned_dim))), - } -} - -#[cfg(test)] -mod dataset_test { - use vector::Metric; - - use crate::model::configuration::index_write_parameters::IndexWriteParametersBuilder; - - use super::*; - - #[test] - #[should_panic(expected = "ERROR: Data file fake_file does not exist.")] - fn create_index_test() { - let index_write_parameters = IndexWriteParametersBuilder::new(50, 4) - .with_alpha(1.2) - .with_saturate_graph(false) - .with_num_threads(1) - .build(); - - let config = IndexConfiguration::new( - Metric::L2, - 128, - 256, - 1_000_000, - false, - 0, - false, - 0, - 1f32, - index_write_parameters, - ); - let mut index = create_inmem_index::(config).unwrap(); - index.build("fake_file", 100).unwrap(); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index.rs deleted file mode 100644 index 871d210..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index.rs +++ /dev/null @@ -1,1033 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::cmp; -use std::sync::RwLock; -use std::time::Duration; - -use hashbrown::hash_set::Entry::*; -use hashbrown::HashSet; -use vector::FullPrecisionDistance; - -use crate::common::{ANNError, ANNResult}; -use crate::index::ANNInmemIndex; -use crate::instrumentation::IndexLogger; -use crate::model::graph::AdjacencyList; -use crate::model::{ - ArcConcurrentBoxedQueue, InMemQueryScratch, InMemoryGraph, IndexConfiguration, InmemDataset, - Neighbor, ScratchStoreManager, Vertex, -}; - -use crate::utils::file_util::{file_exists, load_metadata_from_file}; -use crate::utils::rayon_util::execute_with_rayon; -use crate::utils::{set_rayon_num_threads, Timer}; - -/// In-memory Index -pub struct InmemIndex -where - [T; N]: FullPrecisionDistance, -{ - /// Dataset - pub dataset: InmemDataset, - - /// Graph - pub final_graph: InMemoryGraph, - - /// Index configuration - pub configuration: IndexConfiguration, - - /// Start point of the search. When _num_frozen_pts is greater than zero, - /// this is the location of the first frozen point. Otherwise, this is a - /// location of one of the points in index. - pub start: u32, - - /// Max observed out degree - pub max_observed_degree: u32, - - /// Number of active points i.e. existing in the graph - pub num_active_pts: usize, - - /// query scratch queue. - query_scratch_queue: ArcConcurrentBoxedQueue>, - - pub delete_set: RwLock>, -} - -impl InmemIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - /// Create Index obj based on configuration - pub fn new(mut config: IndexConfiguration) -> ANNResult { - // Sanity check. While logically it is correct, max_points = 0 causes - // downstream problems. - if config.max_points == 0 { - config.max_points = 1; - } - - let total_internal_points = config.max_points + config.num_frozen_pts; - - if config.use_pq_dist { - // TODO: pq - todo!("PQ is not supported now"); - } - - let start = config.max_points.try_into()?; - - let query_scratch_queue = ArcConcurrentBoxedQueue::>::new(); - let delete_set = RwLock::new(HashSet::::new()); - - Ok(Self { - dataset: InmemDataset::::new(total_internal_points, config.growth_potential)?, - final_graph: InMemoryGraph::new( - total_internal_points, - config.index_write_parameter.max_degree, - ), - configuration: config, - start, - max_observed_degree: 0, - num_active_pts: 0, - query_scratch_queue, - delete_set, - }) - } - - /// Get distance between two vertices. - pub fn get_distance(&self, id1: u32, id2: u32) -> ANNResult { - self.dataset - .get_distance(id1, id2, self.configuration.dist_metric) - } - - fn build_with_data_populated(&mut self) -> ANNResult<()> { - println!( - "Starting index build with {} points...", - self.num_active_pts - ); - - if self.num_active_pts < 1 { - return Err(ANNError::log_index_error( - "Error: Trying to build an index with 0 points.".to_string(), - )); - } - - if self.query_scratch_queue.size()? == 0 { - self.initialize_query_scratch( - 5 + self.configuration.index_write_parameter.num_threads, - self.configuration.index_write_parameter.search_list_size, - )?; - } - - // TODO: generate_frozen_point() - - self.link()?; - - self.print_stats()?; - - Ok(()) - } - - fn link(&mut self) -> ANNResult<()> { - // visit_order is a vector that is initialized to the entire graph - let mut visit_order = - Vec::with_capacity(self.num_active_pts + self.configuration.num_frozen_pts); - for i in 0..self.num_active_pts { - visit_order.push(i as u32); - } - - // If there are any frozen points, add them all. - for frozen in self.configuration.max_points - ..(self.configuration.max_points + self.configuration.num_frozen_pts) - { - visit_order.push(frozen as u32); - } - - // if there are frozen points, the first such one is set to be the _start - if self.configuration.num_frozen_pts > 0 { - self.start = self.configuration.max_points as u32; - } else { - self.start = self.dataset.calculate_medoid_point_id()?; - } - - let timer = Timer::new(); - - let range = visit_order.len(); - let logger = IndexLogger::new(range); - - execute_with_rayon( - 0..range, - self.configuration.index_write_parameter.num_threads, - |idx| { - self.insert_vertex_id(visit_order[idx])?; - logger.vertex_processed()?; - - Ok(()) - }, - )?; - - self.cleanup_graph(&visit_order)?; - - if self.num_active_pts > 0 { - println!("{}", timer.elapsed_seconds_for_step("Link time: ")); - } - - Ok(()) - } - - fn insert_vertex_id(&self, vertex_id: u32) -> ANNResult<()> { - let mut scratch_manager = - ScratchStoreManager::new(self.query_scratch_queue.clone(), Duration::from_millis(10))?; - let scratch = scratch_manager.scratch_space().ok_or_else(|| { - ANNError::log_index_error( - "ScratchStoreManager doesn't have InMemQueryScratch instance available".to_string(), - ) - })?; - - let new_neighbors = self.search_for_point_and_prune(scratch, vertex_id)?; - self.update_vertex_with_neighbors(vertex_id, new_neighbors)?; - self.update_neighbors_of_vertex(vertex_id, scratch)?; - - Ok(()) - } - - fn update_neighbors_of_vertex( - &self, - vertex_id: u32, - scratch: &mut InMemQueryScratch, - ) -> Result<(), ANNError> { - let vertex = self.final_graph.read_vertex_and_neighbors(vertex_id)?; - assert!(vertex.size() <= self.configuration.index_write_parameter.max_degree as usize); - self.inter_insert( - vertex_id, - vertex.get_neighbors(), - self.configuration.index_write_parameter.max_degree, - scratch, - )?; - Ok(()) - } - - fn update_vertex_with_neighbors( - &self, - vertex_id: u32, - new_neighbors: AdjacencyList, - ) -> Result<(), ANNError> { - let vertex = &mut self.final_graph.write_vertex_and_neighbors(vertex_id)?; - vertex.set_neighbors(new_neighbors); - assert!(vertex.size() <= self.configuration.index_write_parameter.max_degree as usize); - Ok(()) - } - - fn search_for_point_and_prune( - &self, - scratch: &mut InMemQueryScratch, - vertex_id: u32, - ) -> ANNResult { - let mut pruned_list = - AdjacencyList::for_range(self.configuration.index_write_parameter.max_degree as usize); - let vertex = self.dataset.get_vertex(vertex_id)?; - let mut visited_nodes = self.search_for_point(&vertex, scratch)?; - - self.prune_neighbors(vertex_id, &mut visited_nodes, &mut pruned_list, scratch)?; - - if pruned_list.is_empty() { - return Err(ANNError::log_index_error( - "pruned_list is empty.".to_string(), - )); - } - - if self.final_graph.size() - != self.configuration.max_points + self.configuration.num_frozen_pts - { - return Err(ANNError::log_index_error(format!( - "final_graph has {} vertices instead of {}", - self.final_graph.size(), - self.configuration.max_points + self.configuration.num_frozen_pts, - ))); - } - - Ok(pruned_list) - } - - fn search( - &self, - query: &Vertex, - k_value: usize, - l_value: u32, - indices: &mut [u32], - ) -> ANNResult { - if k_value > l_value as usize { - return Err(ANNError::log_index_error(format!( - "Set L: {} to a value of at least K: {}", - l_value, k_value - ))); - } - - let mut scratch_manager = - ScratchStoreManager::new(self.query_scratch_queue.clone(), Duration::from_millis(10))?; - - let scratch = scratch_manager.scratch_space().ok_or_else(|| { - ANNError::log_index_error( - "ScratchStoreManager doesn't have InMemQueryScratch instance available".to_string(), - ) - })?; - - if l_value > scratch.candidate_size { - println!("Attempting to expand query scratch_space. Was created with Lsize: {} but search L is: {}", scratch.candidate_size, l_value); - scratch.resize_for_new_candidate_size(l_value); - println!( - "Resize completed. New scratch size is: {}", - scratch.candidate_size - ); - } - - let cmp = self.search_with_l_override(query, scratch, l_value as usize)?; - let mut pos = 0; - - for i in 0..scratch.best_candidates.size() { - if scratch.best_candidates[i].id < self.configuration.max_points as u32 { - // Filter out the deleted points. - if let Ok(delete_set_guard) = self.delete_set.read() { - if !delete_set_guard.contains(&scratch.best_candidates[i].id) { - indices[pos] = scratch.best_candidates[i].id; - pos += 1; - } - } else { - return Err(ANNError::log_lock_poison_error( - "failed to acquire the lock for delete_set.".to_string(), - )); - } - } - - if pos == k_value { - break; - } - } - - if pos < k_value { - eprintln!( - "Found fewer than K elements for query! Found: {} but K: {}", - pos, k_value - ); - } - - Ok(cmp) - } - - fn cleanup_graph(&mut self, visit_order: &Vec) -> ANNResult<()> { - if self.num_active_pts > 0 { - println!("Starting final cleanup.."); - } - - execute_with_rayon( - 0..visit_order.len(), - self.configuration.index_write_parameter.num_threads, - |idx| { - let vertex_id = visit_order[idx]; - let num_nbrs = self.get_neighbor_count(vertex_id)?; - - if num_nbrs <= self.configuration.index_write_parameter.max_degree as usize { - // Neighbor list is already small enough. - return Ok(()); - } - - let mut scratch_manager = ScratchStoreManager::new( - self.query_scratch_queue.clone(), - Duration::from_millis(10), - )?; - let scratch = scratch_manager.scratch_space().ok_or_else(|| { - ANNError::log_index_error( - "ScratchStoreManager doesn't have InMemQueryScratch instance available" - .to_string(), - ) - })?; - - let mut dummy_pool = self.get_neighbors_for_vertex(vertex_id)?; - - let mut new_out_neighbors = AdjacencyList::for_range( - self.configuration.index_write_parameter.max_degree as usize, - ); - self.prune_neighbors(vertex_id, &mut dummy_pool, &mut new_out_neighbors, scratch)?; - - self.final_graph - .write_vertex_and_neighbors(vertex_id)? - .set_neighbors(new_out_neighbors); - - Ok(()) - }, - ) - } - - /// Get the unique neighbors for a vertex. - /// - /// This code feels out of place here. This should have nothing to do with whether this - /// is in memory index? - /// # Errors - /// - /// This function will return an error if we are not able to get the read lock. - fn get_neighbors_for_vertex(&self, vertex_id: u32) -> ANNResult> { - let binding = self.final_graph.read_vertex_and_neighbors(vertex_id)?; - let neighbors = binding.get_neighbors(); - let dummy_pool = self.get_unique_neighbors(neighbors, vertex_id)?; - - Ok(dummy_pool) - } - - /// Returns a vector of unique neighbors for the given vertex, along with their distances. - /// - /// # Arguments - /// - /// * `neighbors` - A vector of neighbor id index for the given vertex. - /// * `vertex_id` - The given vertex id. - /// - /// # Errors - /// - /// Returns an `ANNError` if there is an error retrieving the vertex or one of its neighbors. - pub fn get_unique_neighbors( - &self, - neighbors: &Vec, - vertex_id: u32, - ) -> Result, ANNError> { - let vertex = self.dataset.get_vertex(vertex_id)?; - - let len = neighbors.len(); - if len == 0 { - return Ok(Vec::new()); - } - - self.dataset.prefetch_vector(neighbors[0]); - - let mut dummy_visited: HashSet = HashSet::with_capacity(len); - let mut dummy_pool: Vec = Vec::with_capacity(len); - - // let slice = ['w', 'i', 'n', 'd', 'o', 'w', 's']; - // for window in slice.windows(2) { - // &println!{"[{}, {}]", window[0], window[1]}; - // } - // prints: [w, i] -> [i, n] -> [n, d] -> [d, o] -> [o, w] -> [w, s] - for current in neighbors.windows(2) { - // Prefetch the next item. - self.dataset.prefetch_vector(current[1]); - let current = current[0]; - - self.insert_neighbor_if_unique( - &mut dummy_visited, - current, - vertex_id, - &vertex, - &mut dummy_pool, - )?; - } - - // Insert the last neighbor - #[allow(clippy::unwrap_used)] - self.insert_neighbor_if_unique( - &mut dummy_visited, - *neighbors.last().unwrap(), // we know len != 0, so this is safe. - vertex_id, - &vertex, - &mut dummy_pool, - )?; - - Ok(dummy_pool) - } - - fn insert_neighbor_if_unique( - &self, - dummy_visited: &mut HashSet, - current: u32, - vertex_id: u32, - vertex: &Vertex<'_, T, N>, - dummy_pool: &mut Vec, - ) -> Result<(), ANNError> { - if current != vertex_id { - if let Vacant(entry) = dummy_visited.entry(current) { - let cur_nbr_vertex = self.dataset.get_vertex(current)?; - let dist = vertex.compare(&cur_nbr_vertex, self.configuration.dist_metric); - dummy_pool.push(Neighbor::new(current, dist)); - entry.insert(); - } - } - - Ok(()) - } - - /// Get count of neighbors for a given vertex. - /// - /// # Errors - /// - /// This function will return an error if we can't get a lock. - fn get_neighbor_count(&self, vertex_id: u32) -> ANNResult { - let num_nbrs = self - .final_graph - .read_vertex_and_neighbors(vertex_id)? - .size(); - Ok(num_nbrs) - } - - fn soft_delete_vertex(&self, vertex_id_to_delete: u32) -> ANNResult<()> { - if vertex_id_to_delete as usize > self.num_active_pts { - return Err(ANNError::log_index_error(format!( - "vertex_id_to_delete: {} is greater than the number of active points in the graph: {}", - vertex_id_to_delete, self.num_active_pts - ))); - } - - let mut delete_set_guard = match self.delete_set.write() { - Ok(guard) => guard, - Err(_) => { - return Err(ANNError::log_index_error(format!( - "Failed to acquire delete_set lock, cannot delete vertex {}", - vertex_id_to_delete - ))); - } - }; - - delete_set_guard.insert(vertex_id_to_delete); - Ok(()) - } - - fn initialize_query_scratch( - &mut self, - num_threads: u32, - search_candidate_size: u32, - ) -> ANNResult<()> { - self.query_scratch_queue.reserve(num_threads as usize)?; - for _ in 0..num_threads { - let scratch = Box::new(InMemQueryScratch::::new( - search_candidate_size, - &self.configuration.index_write_parameter, - false, - )?); - - self.query_scratch_queue.push(scratch)?; - } - - Ok(()) - } - - fn print_stats(&mut self) -> ANNResult<()> { - let mut max = 0; - let mut min = usize::MAX; - let mut total = 0; - let mut cnt = 0; - - for i in 0..self.num_active_pts { - let vertex_id = i.try_into()?; - let pool_size = self - .final_graph - .read_vertex_and_neighbors(vertex_id)? - .size(); - max = cmp::max(max, pool_size); - min = cmp::min(min, pool_size); - total += pool_size; - if pool_size < 2 { - cnt += 1; - } - } - - println!( - "Index built with degree: max: {} avg: {} min: {} count(deg<2): {}", - max, - (total as f32) / ((self.num_active_pts + self.configuration.num_frozen_pts) as f32), - min, - cnt - ); - - match self.delete_set.read() { - Ok(guard) => { - println!( - "Number of soft deleted vertices {}, soft deleted percentage: {}", - guard.len(), - (guard.len() as f32) - / ((self.num_active_pts + self.configuration.num_frozen_pts) as f32), - ); - } - Err(_) => { - return Err(ANNError::log_lock_poison_error( - "Failed to acquire delete_set lock, cannot get the number of deleted vertices" - .to_string(), - )); - } - }; - - self.max_observed_degree = cmp::max(max as u32, self.max_observed_degree); - - Ok(()) - } -} - -impl ANNInmemIndex for InmemIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - fn build(&mut self, filename: &str, num_points_to_load: usize) -> ANNResult<()> { - // TODO: fresh-diskANN - // std::unique_lock ul(_update_lock); - - if !file_exists(filename) { - return Err(ANNError::log_index_error(format!( - "ERROR: Data file {} does not exist.", - filename - ))); - } - - let (file_num_points, file_dim) = load_metadata_from_file(filename)?; - if file_num_points > self.configuration.max_points { - return Err(ANNError::log_index_error(format!( - "ERROR: Driver requests loading {} points and file has {} points, - but index can support only {} points as specified in configuration.", - num_points_to_load, file_num_points, self.configuration.max_points - ))); - } - - if num_points_to_load > file_num_points { - return Err(ANNError::log_index_error(format!( - "ERROR: Driver requests loading {} points and file has only {} points.", - num_points_to_load, file_num_points - ))); - } - - if file_dim != self.configuration.dim { - return Err(ANNError::log_index_error(format!( - "ERROR: Driver requests loading {} dimension, but file has {} dimension.", - self.configuration.dim, file_dim - ))); - } - - if self.configuration.use_pq_dist { - // TODO: PQ - todo!("PQ is not supported now"); - } - - if self.configuration.index_write_parameter.num_threads > 0 { - set_rayon_num_threads(self.configuration.index_write_parameter.num_threads); - } - - self.dataset.build_from_file(filename, num_points_to_load)?; - - println!("Using only first {} from file.", num_points_to_load); - - // TODO: tag_lock - - self.num_active_pts = num_points_to_load; - self.build_with_data_populated()?; - - Ok(()) - } - - fn insert(&mut self, filename: &str, num_points_to_insert: usize) -> ANNResult<()> { - // fresh-diskANN - if !file_exists(filename) { - return Err(ANNError::log_index_error(format!( - "ERROR: Data file {} does not exist.", - filename - ))); - } - - let (file_num_points, file_dim) = load_metadata_from_file(filename)?; - - if num_points_to_insert > file_num_points { - return Err(ANNError::log_index_error(format!( - "ERROR: Driver requests loading {} points and file has only {} points.", - num_points_to_insert, file_num_points - ))); - } - - if file_dim != self.configuration.dim { - return Err(ANNError::log_index_error(format!( - "ERROR: Driver requests loading {} dimension, but file has {} dimension.", - self.configuration.dim, file_dim - ))); - } - - if self.configuration.use_pq_dist { - // TODO: PQ - todo!("PQ is not supported now"); - } - - if self.query_scratch_queue.size()? == 0 { - self.initialize_query_scratch( - 5 + self.configuration.index_write_parameter.num_threads, - self.configuration.index_write_parameter.search_list_size, - )?; - } - - if self.configuration.index_write_parameter.num_threads > 0 { - // set the thread count of Rayon, otherwise it will use threads as many as logical cores. - std::env::set_var( - "RAYON_NUM_THREADS", - self.configuration - .index_write_parameter - .num_threads - .to_string(), - ); - } - - self.dataset - .append_from_file(filename, num_points_to_insert)?; - self.final_graph.extend( - num_points_to_insert, - self.configuration.index_write_parameter.max_degree, - ); - - // TODO: this should not consider frozen points - let previous_last_pt = self.num_active_pts; - self.num_active_pts += num_points_to_insert; - self.configuration.max_points += num_points_to_insert; - - println!("Inserting {} vectors from file.", num_points_to_insert); - - // TODO: tag_lock - let logger = IndexLogger::new(num_points_to_insert); - let timer = Timer::new(); - execute_with_rayon( - previous_last_pt..self.num_active_pts, - self.configuration.index_write_parameter.num_threads, - |idx| { - self.insert_vertex_id(idx as u32)?; - logger.vertex_processed()?; - - Ok(()) - }, - )?; - - let mut visit_order = - Vec::with_capacity(self.num_active_pts + self.configuration.num_frozen_pts); - for i in 0..self.num_active_pts { - visit_order.push(i as u32); - } - - self.cleanup_graph(&visit_order)?; - println!("{}", timer.elapsed_seconds_for_step("Insert time: ")); - - self.print_stats()?; - - Ok(()) - } - - fn save(&mut self, filename: &str) -> ANNResult<()> { - let data_file = filename.to_string() + ".data"; - let delete_file = filename.to_string() + ".delete"; - - self.save_graph(filename)?; - self.save_data(data_file.as_str())?; - self.save_delete_list(delete_file.as_str())?; - - Ok(()) - } - - fn load(&mut self, filename: &str, expected_num_points: usize) -> ANNResult<()> { - self.num_active_pts = expected_num_points; - self.dataset - .build_from_file(&format!("{}.data", filename), expected_num_points)?; - - self.load_graph(filename, expected_num_points)?; - self.load_delete_list(&format!("{}.delete", filename))?; - - if self.query_scratch_queue.size()? == 0 { - self.initialize_query_scratch( - 5 + self.configuration.index_write_parameter.num_threads, - self.configuration.index_write_parameter.search_list_size, - )?; - } - - Ok(()) - } - - fn search( - &self, - query: &[T], - k_value: usize, - l_value: u32, - indices: &mut [u32], - ) -> ANNResult { - let query_vector = Vertex::new(<&[T; N]>::try_from(query)?, 0); - InmemIndex::search(self, &query_vector, k_value, l_value, indices) - } - - fn soft_delete( - &mut self, - vertex_ids_to_delete: Vec, - num_points_to_delete: usize, - ) -> ANNResult<()> { - println!("Deleting {} vectors from file.", num_points_to_delete); - - let logger = IndexLogger::new(num_points_to_delete); - let timer = Timer::new(); - - execute_with_rayon( - 0..num_points_to_delete, - self.configuration.index_write_parameter.num_threads, - |idx: usize| { - self.soft_delete_vertex(vertex_ids_to_delete[idx])?; - logger.vertex_processed()?; - - Ok(()) - }, - )?; - - println!("{}", timer.elapsed_seconds_for_step("Delete time: ")); - self.print_stats()?; - - Ok(()) - } -} - -#[cfg(test)] -mod index_test { - use vector::Metric; - - use super::*; - use crate::{ - model::{ - configuration::index_write_parameters::IndexWriteParametersBuilder, vertex::DIM_128, - }, - test_utils::get_test_file_path, - utils::file_util::load_ids_to_delete_from_file, - utils::round_up, - }; - - const TEST_DATA_FILE: &str = "tests/data/siftsmall_learn_256pts.fbin"; - const TRUTH_GRAPH: &str = "tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2"; - const TEST_DELETE_FILE: &str = "tests/data/delete_set_50pts.bin"; - const TRUTH_GRAPH_WITH_SATURATED: &str = - "tests/data/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_mem.index"; - const R: u32 = 4; - const L: u32 = 50; - const ALPHA: f32 = 1.2; - - /// Build the index with TEST_DATA_FILE and compare the index graph with truth graph TRUTH_GRAPH - /// Change above constants if you want to test with different dataset - macro_rules! index_end_to_end_test_singlethread { - ($saturate_graph:expr, $truth_graph:expr) => {{ - let (data_num, dim) = - load_metadata_from_file(get_test_file_path(TEST_DATA_FILE).as_str()).unwrap(); - - let index_write_parameters = IndexWriteParametersBuilder::new(L, R) - .with_alpha(ALPHA) - .with_num_threads(1) - .with_saturate_graph($saturate_graph) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num, - false, - 0, - false, - 0, - 1.0f32, - index_write_parameters, - ); - let mut index: InmemIndex = InmemIndex::new(config.clone()).unwrap(); - - index - .build(get_test_file_path(TEST_DATA_FILE).as_str(), data_num) - .unwrap(); - - let mut truth_index: InmemIndex = InmemIndex::new(config).unwrap(); - truth_index - .load_graph(get_test_file_path($truth_graph).as_str(), data_num) - .unwrap(); - - compare_graphs(&index, &truth_index); - }}; - } - - #[test] - fn index_end_to_end_test_singlethread() { - index_end_to_end_test_singlethread!(false, TRUTH_GRAPH); - } - - #[test] - fn index_end_to_end_test_singlethread_with_saturate_graph() { - index_end_to_end_test_singlethread!(true, TRUTH_GRAPH_WITH_SATURATED); - } - - #[test] - fn index_end_to_end_test_multithread() { - let (data_num, dim) = - load_metadata_from_file(get_test_file_path(TEST_DATA_FILE).as_str()).unwrap(); - - let index_write_parameters = IndexWriteParametersBuilder::new(L, R) - .with_alpha(ALPHA) - .with_num_threads(8) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num, - false, - 0, - false, - 0, - 1f32, - index_write_parameters, - ); - let mut index: InmemIndex = InmemIndex::new(config).unwrap(); - - index - .build(get_test_file_path(TEST_DATA_FILE).as_str(), data_num) - .unwrap(); - - for i in 0..index.final_graph.size() { - assert_ne!( - index - .final_graph - .read_vertex_and_neighbors(i as u32) - .unwrap() - .size(), - 0 - ); - } - } - - const TEST_DATA_FILE_2: &str = "tests/data/siftsmall_learn_256pts_2.fbin"; - const INSERT_TRUTH_GRAPH: &str = - "tests/data/truth_index_siftsmall_learn_256pts_1+2_R4_L50_A1.2"; - const INSERT_TRUTH_GRAPH_WITH_SATURATED: &str = - "tests/data/truth_index_siftsmall_learn_256pts_1+2_saturated_R4_L50_A1.2"; - - /// Build the index with TEST_DATA_FILE, insert TEST_DATA_FILE_2 and compare the index graph with truth graph TRUTH_GRAPH - /// Change above constants if you want to test with different dataset - macro_rules! index_insert_end_to_end_test_singlethread { - ($saturate_graph:expr, $truth_graph:expr) => {{ - let (data_num, dim) = - load_metadata_from_file(get_test_file_path(TEST_DATA_FILE).as_str()).unwrap(); - - let index_write_parameters = IndexWriteParametersBuilder::new(L, R) - .with_alpha(ALPHA) - .with_num_threads(1) - .with_saturate_graph($saturate_graph) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num, - false, - 0, - false, - 0, - 2.0f32, - index_write_parameters, - ); - let mut index: InmemIndex = InmemIndex::new(config.clone()).unwrap(); - - index - .build(get_test_file_path(TEST_DATA_FILE).as_str(), data_num) - .unwrap(); - index - .insert(get_test_file_path(TEST_DATA_FILE_2).as_str(), data_num) - .unwrap(); - - let config2 = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num * 2, - false, - 0, - false, - 0, - 1.0f32, - index_write_parameters, - ); - let mut truth_index: InmemIndex = InmemIndex::new(config2).unwrap(); - truth_index - .load_graph(get_test_file_path($truth_graph).as_str(), data_num) - .unwrap(); - - compare_graphs(&index, &truth_index); - }}; - } - - /// Build the index with TEST_DATA_FILE, and delete the vertices with id defined in TEST_DELETE_SET - macro_rules! index_delete_end_to_end_test_singlethread { - () => {{ - let (data_num, dim) = - load_metadata_from_file(get_test_file_path(TEST_DATA_FILE).as_str()).unwrap(); - - let index_write_parameters = IndexWriteParametersBuilder::new(L, R) - .with_alpha(ALPHA) - .with_num_threads(1) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num, - false, - 0, - false, - 0, - 2.0f32, - index_write_parameters, - ); - let mut index: InmemIndex = InmemIndex::new(config.clone()).unwrap(); - - index - .build(get_test_file_path(TEST_DATA_FILE).as_str(), data_num) - .unwrap(); - - let (num_points_to_delete, vertex_ids_to_delete) = - load_ids_to_delete_from_file(TEST_DELETE_FILE).unwrap(); - index - .soft_delete(vertex_ids_to_delete, num_points_to_delete) - .unwrap(); - assert!(index.delete_set.read().unwrap().len() == num_points_to_delete); - }}; - } - - #[test] - fn index_insert_end_to_end_test_singlethread() { - index_insert_end_to_end_test_singlethread!(false, INSERT_TRUTH_GRAPH); - } - - #[test] - fn index_delete_end_to_end_test_singlethread() { - index_delete_end_to_end_test_singlethread!(); - } - - #[test] - fn index_insert_end_to_end_test_saturated_singlethread() { - index_insert_end_to_end_test_singlethread!(true, INSERT_TRUTH_GRAPH_WITH_SATURATED); - } - - fn compare_graphs(index: &InmemIndex, truth_index: &InmemIndex) { - assert_eq!(index.start, truth_index.start); - assert_eq!(index.max_observed_degree, truth_index.max_observed_degree); - assert_eq!(index.final_graph.size(), truth_index.final_graph.size()); - - for i in 0..index.final_graph.size() { - assert_eq!( - index - .final_graph - .read_vertex_and_neighbors(i as u32) - .unwrap() - .size(), - truth_index - .final_graph - .read_vertex_and_neighbors(i as u32) - .unwrap() - .size() - ); - assert_eq!( - index - .final_graph - .read_vertex_and_neighbors(i as u32) - .unwrap() - .get_neighbors(), - truth_index - .final_graph - .read_vertex_and_neighbors(i as u32) - .unwrap() - .get_neighbors() - ); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index_storage.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index_storage.rs deleted file mode 100644 index fa14d70..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/inmem_index_storage.rs +++ /dev/null @@ -1,304 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::fs::File; -use std::io::{BufReader, BufWriter, Seek, SeekFrom, Write}; -use std::path::Path; - -use byteorder::{LittleEndian, ReadBytesExt}; -use vector::FullPrecisionDistance; - -use crate::common::{ANNError, ANNResult}; -use crate::model::graph::AdjacencyList; -use crate::model::InMemoryGraph; -use crate::utils::{file_exists, save_data_in_base_dimensions}; - -use super::InmemIndex; - -impl InmemIndex -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - pub fn load_graph(&mut self, filename: &str, expected_num_points: usize) -> ANNResult { - // let file_offset = 0; // will need this for single file format support - - let mut in_file = BufReader::new(File::open(Path::new(filename))?); - // in_file.seek(SeekFrom::Start(file_offset as u64))?; - - let expected_file_size: usize = in_file.read_u64::()? as usize; - self.max_observed_degree = in_file.read_u32::()?; - self.start = in_file.read_u32::()?; - let file_frozen_pts: usize = in_file.read_u64::()? as usize; - - let vamana_metadata_size = 24; - - println!("From graph header, expected_file_size: {}, max_observed_degree: {}, start: {}, file_frozen_pts: {}", - expected_file_size, self.max_observed_degree, self.start, file_frozen_pts); - - if file_frozen_pts != self.configuration.num_frozen_pts { - if file_frozen_pts == 1 { - return Err(ANNError::log_index_config_error( - "num_frozen_pts".to_string(), - "ERROR: When loading index, detected dynamic index, but constructor asks for static index. Exitting.".to_string()) - ); - } else { - return Err(ANNError::log_index_config_error( - "num_frozen_pts".to_string(), - "ERROR: When loading index, detected static index, but constructor asks for dynamic index. Exitting.".to_string()) - ); - } - } - - println!("Loading vamana graph {}...", filename); - - let expected_max_points = expected_num_points - file_frozen_pts; - - // If user provides more points than max_points - // resize the _final_graph to the larger size. - if self.configuration.max_points < expected_max_points { - println!("Number of points in data: {} is greater than max_points: {} Setting max points to: {}", expected_max_points, self.configuration.max_points, expected_max_points); - - self.configuration.max_points = expected_max_points; - self.final_graph = InMemoryGraph::new( - self.configuration.max_points + self.configuration.num_frozen_pts, - self.configuration.index_write_parameter.max_degree, - ); - } - - let mut bytes_read = vamana_metadata_size; - let mut num_edges = 0; - let mut nodes_read = 0; - let mut max_observed_degree = 0; - - while bytes_read != expected_file_size { - let num_nbrs = in_file.read_u32::()?; - max_observed_degree = if num_nbrs > max_observed_degree { - num_nbrs - } else { - max_observed_degree - }; - - if num_nbrs == 0 { - return Err(ANNError::log_index_error(format!( - "ERROR: Point found with no out-neighbors, point# {}", - nodes_read - ))); - } - - num_edges += num_nbrs; - nodes_read += 1; - let mut tmp: Vec = Vec::with_capacity(num_nbrs as usize); - for _ in 0..num_nbrs { - tmp.push(in_file.read_u32::()?); - } - - self.final_graph - .write_vertex_and_neighbors(nodes_read - 1)? - .set_neighbors(AdjacencyList::from(tmp)); - bytes_read += 4 * (num_nbrs as usize + 1); - } - - println!( - "Done. Index has {} nodes and {} out-edges, _start is set to {}", - nodes_read, num_edges, self.start - ); - - self.max_observed_degree = max_observed_degree; - Ok(nodes_read as usize) - } - - /// Save the graph index on a file as an adjacency list. - /// For each point, first store the number of neighbors, - /// and then the neighbor list (each as 4 byte u32) - pub fn save_graph(&mut self, graph_file: &str) -> ANNResult { - let file: File = File::create(graph_file)?; - let mut out = BufWriter::new(file); - - let file_offset: u64 = 0; - out.seek(SeekFrom::Start(file_offset))?; - let mut index_size: u64 = 24; - let mut max_degree: u32 = 0; - out.write_all(&index_size.to_le_bytes())?; - out.write_all(&self.max_observed_degree.to_le_bytes())?; - out.write_all(&self.start.to_le_bytes())?; - out.write_all(&(self.configuration.num_frozen_pts as u64).to_le_bytes())?; - - // At this point, either nd == max_points or any frozen points have - // been temporarily moved to nd, so nd + num_frozen_points is the valid - // location limit - for i in 0..self.num_active_pts + self.configuration.num_frozen_pts { - let idx = i as u32; - let gk: u32 = self.final_graph.read_vertex_and_neighbors(idx)?.size() as u32; - out.write_all(&gk.to_le_bytes())?; - for neighbor in self - .final_graph - .read_vertex_and_neighbors(idx)? - .get_neighbors() - .iter() - { - out.write_all(&neighbor.to_le_bytes())?; - } - max_degree = - if self.final_graph.read_vertex_and_neighbors(idx)?.size() as u32 > max_degree { - self.final_graph.read_vertex_and_neighbors(idx)?.size() as u32 - } else { - max_degree - }; - index_size += (std::mem::size_of::() * (gk as usize + 1)) as u64; - } - out.seek(SeekFrom::Start(file_offset))?; - out.write_all(&index_size.to_le_bytes())?; - out.write_all(&max_degree.to_le_bytes())?; - out.flush()?; - Ok(index_size) - } - - /// Save the data on a file. - pub fn save_data(&mut self, data_file: &str) -> ANNResult { - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid - // location limit. - Ok(save_data_in_base_dimensions( - data_file, - &mut self.dataset.data, - self.num_active_pts + self.configuration.num_frozen_pts, - self.configuration.dim, - self.configuration.aligned_dim, - 0, - )?) - } - - /// Save the delete list to a file only if the delete list length is not zero. - pub fn save_delete_list(&mut self, delete_list_file: &str) -> ANNResult { - let mut delete_file_size = 0; - if let Ok(delete_set) = self.delete_set.read() { - let delete_set_len = delete_set.len() as u32; - - if delete_set_len != 0 { - let file: File = File::create(delete_list_file)?; - let mut writer = BufWriter::new(file); - - // Write the length of the set. - writer.write_all(&delete_set_len.to_le_bytes())?; - delete_file_size += std::mem::size_of::(); - - // Write the elements of the set. - for &item in delete_set.iter() { - writer.write_all(&item.to_be_bytes())?; - delete_file_size += std::mem::size_of::(); - } - - writer.flush()?; - } - } else { - return Err(ANNError::log_lock_poison_error( - "Poisoned lock on delete set. Can't save deleted list.".to_string(), - )); - } - - Ok(delete_file_size) - } - - // load the deleted list from the delete file if it exists. - pub fn load_delete_list(&mut self, delete_list_file: &str) -> ANNResult { - let mut len = 0; - - if file_exists(delete_list_file) { - let file = File::open(delete_list_file)?; - let mut reader = BufReader::new(file); - - len = reader.read_u32::()? as usize; - - if let Ok(mut delete_set) = self.delete_set.write() { - for _ in 0..len { - let item = reader.read_u32::()?; - delete_set.insert(item); - } - } else { - return Err(ANNError::log_lock_poison_error( - "Poisoned lock on delete set. Can't load deleted list.".to_string(), - )); - } - } - - Ok(len) - } -} - -#[cfg(test)] -mod index_test { - use std::fs; - - use vector::Metric; - - use super::*; - use crate::{ - index::ANNInmemIndex, - model::{ - configuration::index_write_parameters::IndexWriteParametersBuilder, vertex::DIM_128, - IndexConfiguration, - }, - utils::{load_metadata_from_file, round_up}, - }; - - const TEST_DATA_FILE: &str = "tests/data/siftsmall_learn_256pts.fbin"; - const R: u32 = 4; - const L: u32 = 50; - const ALPHA: f32 = 1.2; - - #[cfg_attr(not(coverage), test)] - fn save_graph_test() { - let parameters = IndexWriteParametersBuilder::new(50, 4) - .with_alpha(1.2) - .build(); - let config = - IndexConfiguration::new(Metric::L2, 10, 16, 16, false, 0, false, 8, 1f32, parameters); - let mut index = InmemIndex::::new(config).unwrap(); - let final_graph = InMemoryGraph::new(10, 3); - let num_active_pts = 2_usize; - index.final_graph = final_graph; - index.num_active_pts = num_active_pts; - let graph_file = "test_save_graph_data.bin"; - let result = index.save_graph(graph_file); - assert!(result.is_ok()); - - fs::remove_file(graph_file).expect("Failed to delete file"); - } - - #[test] - fn save_data_test() { - let (data_num, dim) = load_metadata_from_file(TEST_DATA_FILE).unwrap(); - - let index_write_parameters = IndexWriteParametersBuilder::new(L, R) - .with_alpha(ALPHA) - .build(); - let config = IndexConfiguration::new( - Metric::L2, - dim, - round_up(dim as u64, 16_u64) as usize, - data_num, - false, - 0, - false, - 0, - 1f32, - index_write_parameters, - ); - let mut index: InmemIndex = InmemIndex::new(config).unwrap(); - - index.build(TEST_DATA_FILE, data_num).unwrap(); - - let data_file = "test.data"; - let result = index.save_data(data_file); - assert_eq!( - result.unwrap(), - 2 * std::mem::size_of::() - + (index.num_active_pts + index.configuration.num_frozen_pts) - * index.configuration.dim - * (std::mem::size_of::()) - ); - fs::remove_file(data_file).expect("Failed to delete file"); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/mod.rs deleted file mode 100644 index f2a091a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/inmem_index/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod inmem_index; -pub use inmem_index::InmemIndex; - -mod inmem_index_storage; - -pub mod ann_inmem_index; - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/mod.rs deleted file mode 100644 index 18c3bd5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/index/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod inmem_index; -pub use inmem_index::ann_inmem_index::*; -pub use inmem_index::InmemIndex; - -mod disk_index; -pub use disk_index::*; - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/disk_index_build_logger.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/disk_index_build_logger.rs deleted file mode 100644 index d349353..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/disk_index_build_logger.rs +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use logger::logger::indexlog::DiskIndexConstructionCheckpoint; -use logger::logger::indexlog::DiskIndexConstructionLog; -use logger::logger::indexlog::Log; -use logger::logger::indexlog::LogLevel; -use logger::message_handler::send_log; - -use crate::{utils::Timer, common::ANNResult}; - -pub struct DiskIndexBuildLogger { - timer: Timer, - checkpoint: DiskIndexConstructionCheckpoint, -} - -impl DiskIndexBuildLogger { - pub fn new(checkpoint: DiskIndexConstructionCheckpoint) -> Self { - Self { - timer: Timer::new(), - checkpoint, - } - } - - pub fn log_checkpoint(&mut self, next_checkpoint: DiskIndexConstructionCheckpoint) -> ANNResult<()> { - if self.checkpoint == DiskIndexConstructionCheckpoint::None { - return Ok(()); - } - - let mut log = Log::default(); - let disk_index_construction_log = DiskIndexConstructionLog { - checkpoint: self.checkpoint as i32, - time_spent_in_seconds: self.timer.elapsed().as_secs_f32(), - g_cycles_spent: self.timer.elapsed_gcycles(), - log_level: LogLevel::Info as i32, - }; - log.disk_index_construction_log = Some(disk_index_construction_log); - - send_log(log)?; - self.checkpoint = next_checkpoint; - self.timer.reset(); - Ok(()) - } -} - -#[cfg(test)] -mod dataset_test { - use super::*; - - #[test] - fn test_log() { - let mut logger = DiskIndexBuildLogger::new(DiskIndexConstructionCheckpoint::PqConstruction); - logger.log_checkpoint(DiskIndexConstructionCheckpoint::InmemIndexBuild).unwrap();logger.log_checkpoint(logger::logger::indexlog::DiskIndexConstructionCheckpoint::DiskLayout).unwrap(); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/index_logger.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/index_logger.rs deleted file mode 100644 index dfc81ad..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/index_logger.rs +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::sync::atomic::{AtomicUsize, Ordering}; - -use logger::logger::indexlog::IndexConstructionLog; -use logger::logger::indexlog::Log; -use logger::logger::indexlog::LogLevel; -use logger::message_handler::send_log; - -use crate::common::ANNResult; -use crate::utils::Timer; - -pub struct IndexLogger { - items_processed: AtomicUsize, - timer: Timer, - range: usize, -} - -impl IndexLogger { - pub fn new(range: usize) -> Self { - Self { - items_processed: AtomicUsize::new(0), - timer: Timer::new(), - range, - } - } - - pub fn vertex_processed(&self) -> ANNResult<()> { - let count = self.items_processed.fetch_add(1, Ordering::Relaxed); - if count % 100_000 == 0 { - let mut log = Log::default(); - let index_construction_log = IndexConstructionLog { - percentage_complete: (100_f32 * count as f32) / (self.range as f32), - time_spent_in_seconds: self.timer.elapsed().as_secs_f32(), - g_cycles_spent: self.timer.elapsed_gcycles(), - log_level: LogLevel::Info as i32, - }; - log.index_construction_log = Some(index_construction_log); - - send_log(log)?; - } - - Ok(()) - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/mod.rs deleted file mode 100644 index 234e53c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/instrumentation/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod index_logger; -pub use index_logger::IndexLogger; - -mod disk_index_build_logger; -pub use disk_index_build_logger::DiskIndexBuildLogger; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/lib.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/lib.rs deleted file mode 100644 index 1f89e33..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/lib.rs +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] -#![cfg_attr(test, allow(clippy::unused_io_amount))] - -pub mod utils; - -pub mod algorithm; - -pub mod model; - -pub mod common; - -pub mod index; - -pub mod storage; - -pub mod instrumentation; - -#[cfg(test)] -pub mod test_utils; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/disk_index_build_parameter.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/disk_index_build_parameter.rs deleted file mode 100644 index 539192a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/disk_index_build_parameter.rs +++ /dev/null @@ -1,85 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Parameters for disk index construction. - -use crate::common::{ANNResult, ANNError}; - -/// Cached nodes size in GB -const SPACE_FOR_CACHED_NODES_IN_GB: f64 = 0.25; - -/// Threshold for caching in GB -const THRESHOLD_FOR_CACHING_IN_GB: f64 = 1.0; - -/// Parameters specific for disk index construction. -#[derive(Clone, Copy, PartialEq, Debug)] -pub struct DiskIndexBuildParameters { - /// Bound on the memory footprint of the index at search time in bytes. - /// Once built, the index will use up only the specified RAM limit, the rest will reside on disk. - /// This will dictate how aggressively we compress the data vectors to store in memory. - /// Larger will yield better performance at search time. - search_ram_limit: f64, - - /// Limit on the memory allowed for building the index in bytes. - index_build_ram_limit: f64, -} - -impl DiskIndexBuildParameters { - /// Create DiskIndexBuildParameters instance - pub fn new(search_ram_limit_gb: f64, index_build_ram_limit_gb: f64) -> ANNResult { - let param = Self { - search_ram_limit: Self::get_memory_budget(search_ram_limit_gb), - index_build_ram_limit: index_build_ram_limit_gb * 1024_f64 * 1024_f64 * 1024_f64, - }; - - if param.search_ram_limit <= 0f64 { - return Err(ANNError::log_index_config_error("search_ram_limit".to_string(), "RAM budget should be > 0".to_string())) - } - - if param.index_build_ram_limit <= 0f64 { - return Err(ANNError::log_index_config_error("index_build_ram_limit".to_string(), "RAM budget should be > 0".to_string())) - } - - Ok(param) - } - - /// Get search_ram_limit - pub fn search_ram_limit(&self) -> f64 { - self.search_ram_limit - } - - /// Get index_build_ram_limit - pub fn index_build_ram_limit(&self) -> f64 { - self.index_build_ram_limit - } - - fn get_memory_budget(mut index_ram_limit_gb: f64) -> f64 { - if index_ram_limit_gb - SPACE_FOR_CACHED_NODES_IN_GB > THRESHOLD_FOR_CACHING_IN_GB { - // slack for space used by cached nodes - index_ram_limit_gb -= SPACE_FOR_CACHED_NODES_IN_GB; - } - - index_ram_limit_gb * 1024_f64 * 1024_f64 * 1024_f64 - } -} - -#[cfg(test)] -mod dataset_test { - use super::*; - - #[test] - fn sufficient_ram_for_caching() { - let param = DiskIndexBuildParameters::new(1.26_f64, 1.0_f64).unwrap(); - assert_eq!(param.search_ram_limit, 1.01_f64 * 1024_f64 * 1024_f64 * 1024_f64); - } - - #[test] - fn insufficient_ram_for_caching() { - let param = DiskIndexBuildParameters::new(0.03_f64, 1.0_f64).unwrap(); - assert_eq!(param.search_ram_limit, 0.03_f64 * 1024_f64 * 1024_f64 * 1024_f64); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_configuration.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_configuration.rs deleted file mode 100644 index 3e8c472..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_configuration.rs +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Index configuration. - -use vector::Metric; - -use super::index_write_parameters::IndexWriteParameters; - -/// The index configuration -#[derive(Debug, Clone)] -pub struct IndexConfiguration { - /// Index write parameter - pub index_write_parameter: IndexWriteParameters, - - /// Distance metric - pub dist_metric: Metric, - - /// Dimension of the raw data - pub dim: usize, - - /// Aligned dimension - round up dim to the nearest multiple of 8 - pub aligned_dim: usize, - - /// Total number of points in given data set - pub max_points: usize, - - /// Number of points which are used as initial candidates when iterating to - /// closest point(s). These are not visible externally and won't be returned - /// by search. DiskANN forces at least 1 frozen point for dynamic index. - /// The frozen points have consecutive locations. - pub num_frozen_pts: usize, - - /// Calculate distance by PQ or not - pub use_pq_dist: bool, - - /// Number of PQ chunks - pub num_pq_chunks: usize, - - /// Use optimized product quantization - /// Currently not supported - pub use_opq: bool, - - /// potential for growth. 1.2 means the index can grow by up to 20%. - pub growth_potential: f32, - - // TODO: below settings are not supported in current iteration - // pub concurrent_consolidate: bool, - // pub has_built: bool, - // pub save_as_one_file: bool, - // pub dynamic_index: bool, - // pub enable_tags: bool, - // pub normalize_vecs: bool, -} - -impl IndexConfiguration { - /// Create IndexConfiguration instance - #[allow(clippy::too_many_arguments)] - pub fn new( - dist_metric: Metric, - dim: usize, - aligned_dim: usize, - max_points: usize, - use_pq_dist: bool, - num_pq_chunks: usize, - use_opq: bool, - num_frozen_pts: usize, - growth_potential: f32, - index_write_parameter: IndexWriteParameters - ) -> Self { - Self { - index_write_parameter, - dist_metric, - dim, - aligned_dim, - max_points, - num_frozen_pts, - use_pq_dist, - num_pq_chunks, - use_opq, - growth_potential, - } - } - - /// Get the size of adjacency list that we build out. - pub fn write_range(&self) -> usize { - self.index_write_parameter.max_degree as usize - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_write_parameters.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_write_parameters.rs deleted file mode 100644 index cb71f42..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/index_write_parameters.rs +++ /dev/null @@ -1,245 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Index write parameters. - -/// Default parameter values. -pub mod default_param_vals { - /// Default value of alpha. - pub const ALPHA: f32 = 1.2; - - /// Default value of number of threads. - pub const NUM_THREADS: u32 = 0; - - /// Default value of number of rounds. - pub const NUM_ROUNDS: u32 = 2; - - /// Default value of max occlusion size. - pub const MAX_OCCLUSION_SIZE: u32 = 750; - - /// Default value of filter list size. - pub const FILTER_LIST_SIZE: u32 = 0; - - /// Default value of number of frozen points. - pub const NUM_FROZEN_POINTS: u32 = 0; - - /// Default value of max degree. - pub const MAX_DEGREE: u32 = 64; - - /// Default value of build list size. - pub const BUILD_LIST_SIZE: u32 = 100; - - /// Default value of saturate graph. - pub const SATURATE_GRAPH: bool = false; - - /// Default value of search list size. - pub const SEARCH_LIST_SIZE: u32 = 100; -} - -/// Index write parameters. -#[derive(Clone, Copy, PartialEq, Debug)] -pub struct IndexWriteParameters { - /// Search list size - L. - pub search_list_size: u32, - - /// Max degree - R. - pub max_degree: u32, - - /// Saturate graph. - pub saturate_graph: bool, - - /// Max occlusion size - C. - pub max_occlusion_size: u32, - - /// Alpha. - pub alpha: f32, - - /// Number of rounds. - pub num_rounds: u32, - - /// Number of threads. - pub num_threads: u32, - - /// Number of frozen points. - pub num_frozen_points: u32, -} - -impl Default for IndexWriteParameters { - /// Create IndexWriteParameters with default values - fn default() -> Self { - Self { - search_list_size: default_param_vals::SEARCH_LIST_SIZE, - max_degree: default_param_vals::MAX_DEGREE, - saturate_graph: default_param_vals::SATURATE_GRAPH, - max_occlusion_size: default_param_vals::MAX_OCCLUSION_SIZE, - alpha: default_param_vals::ALPHA, - num_rounds: default_param_vals::NUM_ROUNDS, - num_threads: default_param_vals::NUM_THREADS, - num_frozen_points: default_param_vals::NUM_FROZEN_POINTS - } - } -} - -/// The builder for IndexWriteParameters. -#[derive(Debug)] -pub struct IndexWriteParametersBuilder { - search_list_size: u32, - max_degree: u32, - max_occlusion_size: Option, - saturate_graph: Option, - alpha: Option, - num_rounds: Option, - num_threads: Option, - // filter_list_size: Option, - num_frozen_points: Option, -} - -impl IndexWriteParametersBuilder { - /// Initialize IndexWriteParametersBuilder - pub fn new(search_list_size: u32, max_degree: u32) -> Self { - Self { - search_list_size, - max_degree, - max_occlusion_size: None, - saturate_graph: None, - alpha: None, - num_rounds: None, - num_threads: None, - // filter_list_size: None, - num_frozen_points: None, - } - } - - /// Set max occlusion size. - pub fn with_max_occlusion_size(mut self, max_occlusion_size: u32) -> Self { - self.max_occlusion_size = Some(max_occlusion_size); - self - } - - /// Set saturate graph. - pub fn with_saturate_graph(mut self, saturate_graph: bool) -> Self { - self.saturate_graph = Some(saturate_graph); - self - } - - /// Set alpha. - pub fn with_alpha(mut self, alpha: f32) -> Self { - self.alpha = Some(alpha); - self - } - - /// Set number of rounds. - pub fn with_num_rounds(mut self, num_rounds: u32) -> Self { - self.num_rounds = Some(num_rounds); - self - } - - /// Set number of threads. - pub fn with_num_threads(mut self, num_threads: u32) -> Self { - self.num_threads = Some(num_threads); - self - } - - /* - pub fn with_filter_list_size(mut self, filter_list_size: u32) -> Self { - self.filter_list_size = Some(filter_list_size); - self - } - */ - - /// Set number of frozen points. - pub fn with_num_frozen_points(mut self, num_frozen_points: u32) -> Self { - self.num_frozen_points = Some(num_frozen_points); - self - } - - /// Build IndexWriteParameters from IndexWriteParametersBuilder. - pub fn build(self) -> IndexWriteParameters { - IndexWriteParameters { - search_list_size: self.search_list_size, - max_degree: self.max_degree, - saturate_graph: self.saturate_graph.unwrap_or(default_param_vals::SATURATE_GRAPH), - max_occlusion_size: self.max_occlusion_size.unwrap_or(default_param_vals::MAX_OCCLUSION_SIZE), - alpha: self.alpha.unwrap_or(default_param_vals::ALPHA), - num_rounds: self.num_rounds.unwrap_or(default_param_vals::NUM_ROUNDS), - num_threads: self.num_threads.unwrap_or(default_param_vals::NUM_THREADS), - // filter_list_size: self.filter_list_size.unwrap_or(default_param_vals::FILTER_LIST_SIZE), - num_frozen_points: self.num_frozen_points.unwrap_or(default_param_vals::NUM_FROZEN_POINTS), - } - } -} - -/// Construct IndexWriteParametersBuilder from IndexWriteParameters. -impl From for IndexWriteParametersBuilder { - fn from(param: IndexWriteParameters) -> Self { - Self { - search_list_size: param.search_list_size, - max_degree: param.max_degree, - max_occlusion_size: Some(param.max_occlusion_size), - saturate_graph: Some(param.saturate_graph), - alpha: Some(param.alpha), - num_rounds: Some(param.num_rounds), - num_threads: Some(param.num_threads), - // filter_list_size: Some(param.filter_list_size), - num_frozen_points: Some(param.num_frozen_points), - } - } -} - -#[cfg(test)] -mod parameters_test { - use crate::model::configuration::index_write_parameters::*; - - #[test] - fn test_default_index_params() { - let wp1 = IndexWriteParameters::default(); - assert_eq!(wp1.search_list_size, default_param_vals::SEARCH_LIST_SIZE); - assert_eq!(wp1.max_degree, default_param_vals::MAX_DEGREE); - assert_eq!(wp1.saturate_graph, default_param_vals::SATURATE_GRAPH); - assert_eq!(wp1.max_occlusion_size, default_param_vals::MAX_OCCLUSION_SIZE); - assert_eq!(wp1.alpha, default_param_vals::ALPHA); - assert_eq!(wp1.num_rounds, default_param_vals::NUM_ROUNDS); - assert_eq!(wp1.num_threads, default_param_vals::NUM_THREADS); - assert_eq!(wp1.num_frozen_points, default_param_vals::NUM_FROZEN_POINTS); - } - - #[test] - fn test_index_write_parameters_builder() { - // default value - let wp1 = IndexWriteParametersBuilder::new(10, 20).build(); - assert_eq!(wp1.search_list_size, 10); - assert_eq!(wp1.max_degree, 20); - assert_eq!(wp1.saturate_graph, default_param_vals::SATURATE_GRAPH); - assert_eq!(wp1.max_occlusion_size, default_param_vals::MAX_OCCLUSION_SIZE); - assert_eq!(wp1.alpha, default_param_vals::ALPHA); - assert_eq!(wp1.num_rounds, default_param_vals::NUM_ROUNDS); - assert_eq!(wp1.num_threads, default_param_vals::NUM_THREADS); - assert_eq!(wp1.num_frozen_points, default_param_vals::NUM_FROZEN_POINTS); - - // build with custom values - let wp2 = IndexWriteParametersBuilder::new(10, 20) - .with_max_occlusion_size(30) - .with_saturate_graph(true) - .with_alpha(0.5) - .with_num_rounds(40) - .with_num_threads(50) - .with_num_frozen_points(60) - .build(); - assert_eq!(wp2.search_list_size, 10); - assert_eq!(wp2.max_degree, 20); - assert!(wp2.saturate_graph); - assert_eq!(wp2.max_occlusion_size, 30); - assert_eq!(wp2.alpha, 0.5); - assert_eq!(wp2.num_rounds, 40); - assert_eq!(wp2.num_threads, 50); - assert_eq!(wp2.num_frozen_points, 60); - - // test from - let wp3 = IndexWriteParametersBuilder::from(wp2).build(); - assert_eq!(wp3, wp2); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/mod.rs deleted file mode 100644 index 201f97e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/configuration/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod index_configuration; -pub use index_configuration::IndexConfiguration; - -pub mod index_write_parameters; -pub use index_write_parameters::*; - -pub mod disk_index_build_parameter; -pub use disk_index_build_parameter::DiskIndexBuildParameters; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/disk_scratch_dataset.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/disk_scratch_dataset.rs deleted file mode 100644 index 0d9a007..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/disk_scratch_dataset.rs +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Disk scratch dataset - -use std::mem::{size_of, size_of_val}; -use std::ptr; - -use crate::common::{AlignedBoxWithSlice, ANNResult}; -use crate::model::MAX_N_CMPS; -use crate::utils::round_up; - -/// DiskScratchDataset alignment -pub const DISK_SCRATCH_DATASET_ALIGN: usize = 256; - -/// Disk scratch dataset storing fp vectors with aligned dim -#[derive(Debug)] -pub struct DiskScratchDataset -{ - /// fp vectors with aligned dim - pub data: AlignedBoxWithSlice, - - /// current index to store the next fp vector - pub cur_index: usize, -} - -impl DiskScratchDataset -{ - /// Create DiskScratchDataset instance - pub fn new() -> ANNResult { - Ok(Self { - // C++ code allocates round_up(MAX_N_CMPS * N, 256) bytes, shouldn't it be round_up(MAX_N_CMPS * N, 256) * size_of:: bytes? - data: AlignedBoxWithSlice::new( - round_up(MAX_N_CMPS * N, DISK_SCRATCH_DATASET_ALIGN), - DISK_SCRATCH_DATASET_ALIGN)?, - cur_index: 0, - }) - } - - /// memcpy from fp vector bytes (its len should be `dim * size_of::()`) to self.data - /// The dest slice is a fp vector with aligned dim - /// * fp_vector_buf's dim might not be aligned dim (N) - /// # Safety - /// Behavior is undefined if any of the following conditions are violated: - /// - /// * `fp_vector_buf`'s len must be `dim * size_of::()` bytes - /// - /// * `fp_vector_buf` must be smaller than or equal to `N * size_of::()` bytes. - /// - /// * `fp_vector_buf` and `self.data` must be nonoverlapping. - pub unsafe fn memcpy_from_fp_vector_buf(&mut self, fp_vector_buf: &[u8]) -> &[T] { - if self.cur_index == MAX_N_CMPS { - self.cur_index = 0; - } - - let aligned_dim_vector = &mut self.data[self.cur_index * N..(self.cur_index + 1) * N]; - - assert!(fp_vector_buf.len() % size_of::() == 0); - assert!(fp_vector_buf.len() <= size_of_val(aligned_dim_vector)); - - // memcpy from fp_vector_buf to aligned_dim_vector - unsafe { - ptr::copy_nonoverlapping( - fp_vector_buf.as_ptr(), - aligned_dim_vector.as_mut_ptr() as *mut u8, - fp_vector_buf.len(), - ); - } - - self.cur_index += 1; - aligned_dim_vector - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/inmem_dataset.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/inmem_dataset.rs deleted file mode 100644 index 6d8b649..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/inmem_dataset.rs +++ /dev/null @@ -1,285 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! In-memory Dataset - -use rayon::prelude::*; -use std::mem; -use vector::{FullPrecisionDistance, Metric}; - -use crate::common::{ANNError, ANNResult, AlignedBoxWithSlice}; -use crate::model::Vertex; -use crate::utils::copy_aligned_data_from_file; - -/// Dataset of all in-memory FP points -#[derive(Debug)] -pub struct InmemDataset -where - [T; N]: FullPrecisionDistance, -{ - /// All in-memory points - pub data: AlignedBoxWithSlice, - - /// Number of points we anticipate to have - pub num_points: usize, - - /// Number of active points i.e. existing in the graph - pub num_active_pts: usize, - - /// Capacity of the dataset - pub capacity: usize, -} - -impl<'a, T, const N: usize> InmemDataset -where - T: Default + Copy + Sync + Send + Into, - [T; N]: FullPrecisionDistance, -{ - /// Create the dataset with size num_points and growth factor. - /// growth factor=1 means no growth (provision 100% space of num_points) - /// growth factor=1.2 means provision 120% space of num_points (20% extra space) - pub fn new(num_points: usize, index_growth_factor: f32) -> ANNResult { - let capacity = (((num_points * N) as f32) * index_growth_factor) as usize; - - Ok(Self { - data: AlignedBoxWithSlice::new(capacity, mem::size_of::() * 16)?, - num_points, - num_active_pts: num_points, - capacity, - }) - } - - /// get immutable data slice - pub fn get_data(&self) -> &[T] { - &self.data - } - - /// Build the dataset from file - pub fn build_from_file(&mut self, filename: &str, num_points_to_load: usize) -> ANNResult<()> { - println!( - "Loading {} vectors from file {} into dataset...", - num_points_to_load, filename - ); - self.num_active_pts = num_points_to_load; - - copy_aligned_data_from_file(filename, self.into_dto(), 0)?; - - println!("Dataset loaded."); - Ok(()) - } - - /// Append the dataset from file - pub fn append_from_file( - &mut self, - filename: &str, - num_points_to_append: usize, - ) -> ANNResult<()> { - println!( - "Appending {} vectors from file {} into dataset...", - num_points_to_append, filename - ); - if self.num_points + num_points_to_append > self.capacity { - return Err(ANNError::log_index_error(format!( - "Cannot append {} points to dataset of capacity {}", - num_points_to_append, self.capacity - ))); - } - - let pts_offset = self.num_active_pts; - copy_aligned_data_from_file(filename, self.into_dto(), pts_offset)?; - - self.num_active_pts += num_points_to_append; - self.num_points += num_points_to_append; - - println!("Dataset appended."); - Ok(()) - } - - /// Get vertex by id - pub fn get_vertex(&'a self, id: u32) -> ANNResult> { - let start = id as usize * N; - let end = start + N; - - if end <= self.data.len() { - let val = <&[T; N]>::try_from(&self.data[start..end]).map_err(|err| { - ANNError::log_index_error(format!("Failed to get vertex {}, err={}", id, err)) - })?; - Ok(Vertex::new(val, id)) - } else { - Err(ANNError::log_index_error(format!( - "Invalid vertex id {}.", - id - ))) - } - } - - /// Get full precision distance between two nodes - pub fn get_distance(&self, id1: u32, id2: u32, metric: Metric) -> ANNResult { - let vertex1 = self.get_vertex(id1)?; - let vertex2 = self.get_vertex(id2)?; - - Ok(vertex1.compare(&vertex2, metric)) - } - - /// find out the medoid, the vertex in the dataset that is closest to the centroid - pub fn calculate_medoid_point_id(&self) -> ANNResult { - Ok(self.find_nearest_point_id(self.calculate_centroid_point()?)) - } - - /// calculate centroid, average of all vertices in the dataset - fn calculate_centroid_point(&self) -> ANNResult<[f32; N]> { - // Allocate and initialize the centroid vector - let mut center: [f32; N] = [0.0; N]; - - // Sum the data points' components - for i in 0..self.num_active_pts { - let vertex = self.get_vertex(i as u32)?; - let vertex_slice = vertex.vector(); - for j in 0..N { - center[j] += vertex_slice[j].into(); - } - } - - // Divide by the number of points to calculate the centroid - let capacity = self.num_active_pts as f32; - for item in center.iter_mut().take(N) { - *item /= capacity; - } - - Ok(center) - } - - /// find out the vertex closest to the given point - fn find_nearest_point_id(&self, point: [f32; N]) -> u32 { - // compute all to one distance - let mut distances = vec![0f32; self.num_active_pts]; - let slice = &self.data[..]; - distances.par_iter_mut().enumerate().for_each(|(i, dist)| { - let start = i * N; - for j in 0..N { - let diff: f32 = (point.as_slice()[j] - slice[start + j].into()) - * (point.as_slice()[j] - slice[start + j].into()); - *dist += diff; - } - }); - - let mut min_idx = 0; - let mut min_dist = f32::MAX; - for (i, distance) in distances.iter().enumerate().take(self.num_active_pts) { - if *distance < min_dist { - min_idx = i; - min_dist = *distance; - } - } - min_idx as u32 - } - - /// Prefetch vertex data in the memory hierarchy - /// NOTE: good efficiency when total_vec_size is integral multiple of 64 - #[inline] - pub fn prefetch_vector(&self, id: u32) { - let start = id as usize * N; - let end = start + N; - - if end <= self.data.len() { - let vec = &self.data[start..end]; - vector::prefetch_vector(vec); - } - } - - /// Convert into dto object - pub fn into_dto(&mut self) -> DatasetDto { - DatasetDto { - data: &mut self.data, - rounded_dim: N, - } - } -} - -/// Dataset dto used for other layer, such as storage -/// N is the aligned dimension -#[derive(Debug)] -pub struct DatasetDto<'a, T> { - /// data slice borrow from dataset - pub data: &'a mut [T], - - /// rounded dimension - pub rounded_dim: usize, -} - -#[cfg(test)] -mod dataset_test { - use std::fs; - - use super::*; - use crate::model::vertex::DIM_128; - - #[test] - fn get_vertex_within_range() { - let num_points = 1_000_000; - let id = 999_999; - let dataset = InmemDataset::::new(num_points, 1f32).unwrap(); - - let vertex = dataset.get_vertex(999_999).unwrap(); - - assert_eq!(vertex.vertex_id(), id); - assert_eq!(vertex.vector().len(), DIM_128); - assert_eq!(vertex.vector().as_ptr(), unsafe { - dataset.data.as_ptr().add((id as usize) * DIM_128) - }); - } - - #[test] - fn get_vertex_out_of_range() { - let num_points = 1_000_000; - let invalid_id = 1_000_000; - let dataset = InmemDataset::::new(num_points, 1f32).unwrap(); - - if dataset.get_vertex(invalid_id).is_ok() { - panic!("id ({}) should be out of range", invalid_id) - }; - } - - #[test] - fn load_data_test() { - let file_name = "dataset_test_load_data_test.bin"; - //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8] - let data: [u8; 72] = [ - 2, 0, 0, 0, 8, 0, 0, 0, 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, - 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, - 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, - 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, 0x00, 0x00, 0x50, 0x41, - 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x00, 0x80, 0x41, - ]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let mut dataset = InmemDataset::::new(2, 1f32).unwrap(); - - match copy_aligned_data_from_file( - file_name, - dataset.into_dto(), - 0, - ) { - Ok((npts, dim)) => { - fs::remove_file(file_name).expect("Failed to delete file"); - assert!(npts == 2); - assert!(dim == 8); - assert!(dataset.data.len() == 16); - - let first_vertex = dataset.get_vertex(0).unwrap(); - let second_vertex = dataset.get_vertex(1).unwrap(); - - assert!(*first_vertex.vector() == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); - assert!(*second_vertex.vector() == [9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]); - } - Err(e) => { - fs::remove_file(file_name).expect("Failed to delete file"); - panic!("{}", e) - } - } - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/mod.rs deleted file mode 100644 index 4e7e683..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/data_store/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod inmem_dataset; -pub use inmem_dataset::InmemDataset; -pub use inmem_dataset::DatasetDto; - -mod disk_scratch_dataset; -pub use disk_scratch_dataset::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/adjacency_list.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/adjacency_list.rs deleted file mode 100644 index 7ad2d7d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/adjacency_list.rs +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Adjacency List - -use std::ops::{Deref, DerefMut}; - -#[derive(Debug, Eq, PartialEq)] -/// Represents the out neighbors of a vertex -pub struct AdjacencyList { - edges: Vec, -} - -/// In-mem index related limits -const GRAPH_SLACK_FACTOR: f32 = 1.3_f32; - -impl AdjacencyList { - /// Create AdjacencyList with capacity slack for a range. - pub fn for_range(range: usize) -> Self { - let capacity = (range as f32 * GRAPH_SLACK_FACTOR).ceil() as usize; - Self { - edges: Vec::with_capacity(capacity), - } - } - - /// Push a node to the list of neighbors for the given node. - pub fn push(&mut self, node_id: u32) { - debug_assert!(self.edges.len() < self.edges.capacity()); - self.edges.push(node_id); - } -} - -impl From> for AdjacencyList { - fn from(edges: Vec) -> Self { - Self { edges } - } -} - -impl Deref for AdjacencyList { - type Target = Vec; - - fn deref(&self) -> &Self::Target { - &self.edges - } -} - -impl DerefMut for AdjacencyList { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.edges - } -} - -impl<'a> IntoIterator for &'a AdjacencyList { - type Item = &'a u32; - type IntoIter = std::slice::Iter<'a, u32>; - - fn into_iter(self) -> Self::IntoIter { - self.edges.iter() - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/disk_graph.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/disk_graph.rs deleted file mode 100644 index 49190b1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/disk_graph.rs +++ /dev/null @@ -1,179 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_docs)] - -//! Disk graph - -use byteorder::{LittleEndian, ByteOrder}; -use vector::FullPrecisionDistance; - -use crate::common::{ANNResult, ANNError}; -use crate::model::data_store::DiskScratchDataset; -use crate::model::Vertex; -use crate::storage::DiskGraphStorage; - -use super::{VertexAndNeighbors, SectorGraph, AdjacencyList}; - -/// Disk graph -pub struct DiskGraph { - /// dim of fp vector in disk sector - dim: usize, - - /// number of nodes per sector - num_nodes_per_sector: u64, - - /// max node length in bytes - max_node_len: u64, - - /// the len of fp vector - fp_vector_len: u64, - - /// list of nodes (vertex_id) to fetch from disk - nodes_to_fetch: Vec, - - /// Sector graph - sector_graph: SectorGraph, -} - -impl<'a> DiskGraph { - /// Create DiskGraph instance - pub fn new( - dim: usize, - num_nodes_per_sector: u64, - max_node_len: u64, - fp_vector_len: u64, - beam_width: usize, - graph_storage: DiskGraphStorage, - ) -> ANNResult { - let graph = Self { - dim, - num_nodes_per_sector, - max_node_len, - fp_vector_len, - nodes_to_fetch: Vec::with_capacity(2 * beam_width), - sector_graph: SectorGraph::new(graph_storage)?, - }; - - Ok(graph) - } - - /// Add vertex_id into the list to fetch from disk - pub fn add_vertex(&mut self, id: u32) { - self.nodes_to_fetch.push(id); - } - - /// Fetch nodes from disk index - pub fn fetch_nodes(&mut self) -> ANNResult<()> { - let sectors_to_fetch: Vec = self.nodes_to_fetch.iter().map(|&id| self.node_sector_index(id)).collect(); - self.sector_graph.read_graph(§ors_to_fetch)?; - - Ok(()) - } - - /// Copy disk fp vector to DiskScratchDataset - /// Return the fp vector with aligned dim from DiskScratchDataset - pub fn copy_fp_vector_to_disk_scratch_dataset( - &self, - node_index: usize, - disk_scratch_dataset: &'a mut DiskScratchDataset - ) -> ANNResult> - where - [T; N]: FullPrecisionDistance, - { - if self.dim > N { - return Err(ANNError::log_index_error(format!( - "copy_sector_fp_to_aligned_dataset: dim {} is greater than aligned dim {}", - self.dim, N))); - } - - let fp_vector_buf = self.node_fp_vector_buf(node_index); - - // Safety condition is met here - let aligned_dim_vector = unsafe { disk_scratch_dataset.memcpy_from_fp_vector_buf(fp_vector_buf) }; - - Vertex::<'a, T, N>::try_from((aligned_dim_vector, self.nodes_to_fetch[node_index])) - .map_err(|err| ANNError::log_index_error(format!("TryFromSliceError: failed to get Vertex for disk index node, err={}", err))) - } - - /// Reset graph - pub fn reset(&mut self) { - self.nodes_to_fetch.clear(); - self.sector_graph.reset(); - } - - fn get_vertex_and_neighbors(&self, node_index: usize) -> VertexAndNeighbors { - let node_disk_buf = self.node_disk_buf(node_index); - let buf = &node_disk_buf[self.fp_vector_len as usize..]; - let num_neighbors = LittleEndian::read_u32(&buf[0..4]) as usize; - let neighbors_buf = &buf[4..4 + num_neighbors * 4]; - - let mut adjacency_list = AdjacencyList::for_range(num_neighbors); - for chunk in neighbors_buf.chunks(4) { - let neighbor_id = LittleEndian::read_u32(chunk); - adjacency_list.push(neighbor_id); - } - - VertexAndNeighbors::new(self.nodes_to_fetch[node_index], adjacency_list) - } - - #[inline] - fn node_sector_index(&self, vertex_id: u32) -> u64 { - vertex_id as u64 / self.num_nodes_per_sector + 1 - } - - #[inline] - fn node_disk_buf(&self, node_index: usize) -> &[u8] { - let vertex_id = self.nodes_to_fetch[node_index]; - - // get sector_buf where this node is located - let sector_buf = self.sector_graph.get_sector_buf(node_index); - let node_offset = (vertex_id as u64 % self.num_nodes_per_sector * self.max_node_len) as usize; - §or_buf[node_offset..node_offset + self.max_node_len as usize] - } - - #[inline] - fn node_fp_vector_buf(&self, node_index: usize) -> &[u8] { - let node_disk_buf = self.node_disk_buf(node_index); - &node_disk_buf[..self.fp_vector_len as usize] - } -} - -/// Iterator for DiskGraph -pub struct DiskGraphIntoIterator<'a> { - graph: &'a DiskGraph, - index: usize, -} - -impl<'a> IntoIterator for &'a DiskGraph -{ - type IntoIter = DiskGraphIntoIterator<'a>; - type Item = ANNResult<(usize, VertexAndNeighbors)>; - - #[inline] - fn into_iter(self) -> Self::IntoIter { - DiskGraphIntoIterator { - graph: self, - index: 0, - } - } -} - -impl<'a> Iterator for DiskGraphIntoIterator<'a> -{ - type Item = ANNResult<(usize, VertexAndNeighbors)>; - - fn next(&mut self) -> Option { - if self.index >= self.graph.nodes_to_fetch.len() { - return None; - } - - let node_index = self.index; - let vertex_and_neighbors = self.graph.get_vertex_and_neighbors(self.index); - - self.index += 1; - Some(Ok((node_index, vertex_and_neighbors))) - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/inmem_graph.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/inmem_graph.rs deleted file mode 100644 index 3d08db8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/inmem_graph.rs +++ /dev/null @@ -1,141 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! In-memory graph - -use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; - -use crate::common::ANNError; - -use super::VertexAndNeighbors; - -/// The entire graph of in-memory index -#[derive(Debug)] -pub struct InMemoryGraph { - /// The entire graph - pub final_graph: Vec>, -} - -impl InMemoryGraph { - /// Create InMemoryGraph instance - pub fn new(size: usize, max_degree: u32) -> Self { - let mut graph = Vec::with_capacity(size); - for id in 0..size { - graph.push(RwLock::new(VertexAndNeighbors::for_range( - id as u32, - max_degree as usize, - ))); - } - Self { final_graph: graph } - } - - /// Size of graph - pub fn size(&self) -> usize { - self.final_graph.len() - } - - /// Extend the graph by size vectors - pub fn extend(&mut self, size: usize, max_degree: u32) { - for id in 0..size { - self.final_graph - .push(RwLock::new(VertexAndNeighbors::for_range( - id as u32, - max_degree as usize, - ))); - } - } - - /// Get read guard of vertex_id - pub fn read_vertex_and_neighbors( - &self, - vertex_id: u32, - ) -> Result, ANNError> { - self.final_graph[vertex_id as usize].read().map_err(|err| { - ANNError::log_lock_poison_error(format!( - "PoisonError: Lock poisoned when reading final_graph for vertex_id {}, err={}", - vertex_id, err - )) - }) - } - - /// Get write guard of vertex_id - pub fn write_vertex_and_neighbors( - &self, - vertex_id: u32, - ) -> Result, ANNError> { - self.final_graph[vertex_id as usize].write().map_err(|err| { - ANNError::log_lock_poison_error(format!( - "PoisonError: Lock poisoned when writing final_graph for vertex_id {}, err={}", - vertex_id, err - )) - }) - } -} - -#[cfg(test)] -mod graph_tests { - use crate::model::{graph::AdjacencyList, GRAPH_SLACK_FACTOR}; - - use super::*; - - #[test] - fn test_new() { - let graph = InMemoryGraph::new(10, 10); - let capacity = (GRAPH_SLACK_FACTOR * 10_f64).ceil() as usize; - - assert_eq!(graph.final_graph.len(), 10); - for i in 0..10 { - let neighbor = graph.final_graph[i].read().unwrap(); - assert_eq!(neighbor.vertex_id, i as u32); - assert_eq!(neighbor.get_neighbors().capacity(), capacity); - } - } - - #[test] - fn test_size() { - let graph = InMemoryGraph::new(10, 10); - assert_eq!(graph.size(), 10); - } - - #[test] - fn test_extend() { - let mut graph = InMemoryGraph::new(10, 10); - graph.extend(10, 10); - - assert_eq!(graph.size(), 20); - - let capacity = (GRAPH_SLACK_FACTOR * 10_f64).ceil() as usize; - let mut id: u32 = 0; - - for i in 10..20 { - let neighbor = graph.final_graph[i].read().unwrap(); - assert_eq!(neighbor.vertex_id, id); - assert_eq!(neighbor.get_neighbors().capacity(), capacity); - id += 1; - } - } - - #[test] - fn test_read_vertex_and_neighbors() { - let graph = InMemoryGraph::new(10, 10); - let neighbor = graph.read_vertex_and_neighbors(0); - assert!(neighbor.is_ok()); - assert_eq!(neighbor.unwrap().vertex_id, 0); - } - - #[test] - fn test_write_vertex_and_neighbors() { - let graph = InMemoryGraph::new(10, 10); - { - let neighbor = graph.write_vertex_and_neighbors(0); - assert!(neighbor.is_ok()); - neighbor.unwrap().add_to_neighbors(10, 10); - } - - let neighbor = graph.read_vertex_and_neighbors(0).unwrap(); - assert_eq!(neighbor.get_neighbors(), &AdjacencyList::from(vec![10_u32])); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/mod.rs deleted file mode 100644 index d1457f1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod inmem_graph; -pub use inmem_graph::InMemoryGraph; - -pub mod vertex_and_neighbors; -pub use vertex_and_neighbors::VertexAndNeighbors; - -mod adjacency_list; -pub use adjacency_list::AdjacencyList; - -mod sector_graph; -pub use sector_graph::*; - -mod disk_graph; -pub use disk_graph::*; - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/sector_graph.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/sector_graph.rs deleted file mode 100644 index e51e0bf..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/sector_graph.rs +++ /dev/null @@ -1,87 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_docs)] - -//! Sector graph - -use std::ops::Deref; - -use crate::common::{AlignedBoxWithSlice, ANNResult, ANNError}; -use crate::model::{MAX_N_SECTOR_READS, SECTOR_LEN, AlignedRead}; -use crate::storage::DiskGraphStorage; - -/// Sector graph read from disk index -pub struct SectorGraph { - /// Sector bytes from disk - /// One sector has num_nodes_per_sector nodes - /// Each node's layout: {full precision vector:[T; DIM]}{num_nbrs: u32}{neighbors: [u32; num_nbrs]} - /// The fp vector is not aligned - sectors_data: AlignedBoxWithSlice, - - /// Graph storage to read sectors - graph_storage: DiskGraphStorage, - - /// Current sector index into which the next read reads data - cur_sector_idx: u64, -} - -impl SectorGraph { - /// Create SectorGraph instance - pub fn new(graph_storage: DiskGraphStorage) -> ANNResult { - Ok(Self { - sectors_data: AlignedBoxWithSlice::new(MAX_N_SECTOR_READS * SECTOR_LEN, SECTOR_LEN)?, - graph_storage, - cur_sector_idx: 0, - }) - } - - /// Reset SectorGraph - pub fn reset(&mut self) { - self.cur_sector_idx = 0; - } - - /// Read sectors into sectors_data - /// They are in the same order as sectors_to_fetch - pub fn read_graph(&mut self, sectors_to_fetch: &[u64]) -> ANNResult<()> { - let cur_sector_idx_usize: usize = self.cur_sector_idx.try_into()?; - if sectors_to_fetch.len() > MAX_N_SECTOR_READS - cur_sector_idx_usize { - return Err(ANNError::log_index_error(format!( - "Trying to read too many sectors. number of sectors to read: {}, max number of sectors can read: {}", - sectors_to_fetch.len(), - MAX_N_SECTOR_READS - cur_sector_idx_usize, - ))); - } - - let mut sector_slices = self.sectors_data.split_into_nonoverlapping_mut_slices( - cur_sector_idx_usize * SECTOR_LEN..(cur_sector_idx_usize + sectors_to_fetch.len()) * SECTOR_LEN, - SECTOR_LEN)?; - - let mut read_requests = Vec::with_capacity(sector_slices.len()); - for (local_sector_idx, slice) in sector_slices.iter_mut().enumerate() { - let sector_id = sectors_to_fetch[local_sector_idx]; - read_requests.push(AlignedRead::new(sector_id * SECTOR_LEN as u64, slice)?); - } - - self.graph_storage.read(&mut read_requests)?; - self.cur_sector_idx += sectors_to_fetch.len() as u64; - - Ok(()) - } - - /// Get sector data by local index - #[inline] - pub fn get_sector_buf(&self, local_sector_idx: usize) -> &[u8] { - &self.sectors_data[local_sector_idx * SECTOR_LEN..(local_sector_idx + 1) * SECTOR_LEN] - } -} - -impl Deref for SectorGraph { - type Target = [u8]; - - fn deref(&self) -> &Self::Target { - &self.sectors_data - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/vertex_and_neighbors.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/vertex_and_neighbors.rs deleted file mode 100644 index a9fa389..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/graph/vertex_and_neighbors.rs +++ /dev/null @@ -1,159 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Vertex and its Adjacency List - -use crate::model::GRAPH_SLACK_FACTOR; - -use super::AdjacencyList; - -/// The out neighbors of vertex_id -#[derive(Debug)] -pub struct VertexAndNeighbors { - /// The id of the vertex - pub vertex_id: u32, - - /// All out neighbors (id) of vertex_id - neighbors: AdjacencyList, -} - -impl VertexAndNeighbors { - /// Create VertexAndNeighbors with id and capacity - pub fn for_range(id: u32, range: usize) -> Self { - Self { - vertex_id: id, - neighbors: AdjacencyList::for_range(range), - } - } - - /// Create VertexAndNeighbors with id and neighbors - pub fn new(vertex_id: u32, neighbors: AdjacencyList) -> Self { - Self { - vertex_id, - neighbors, - } - } - - /// Get size of neighbors - #[inline(always)] - pub fn size(&self) -> usize { - self.neighbors.len() - } - - /// Update the neighbors vector (post a pruning exercise) - #[inline(always)] - pub fn set_neighbors(&mut self, new_neighbors: AdjacencyList) { - // Replace the graph entry with the pruned neighbors - self.neighbors = new_neighbors; - } - - /// Get the neighbors - #[inline(always)] - pub fn get_neighbors(&self) -> &AdjacencyList { - &self.neighbors - } - - /// Adds a node to the list of neighbors for the given node. - /// - /// # Arguments - /// - /// * `node_id` - The ID of the node to add. - /// * `range` - The range of the graph. - /// - /// # Return - /// - /// Returns `None` if the node is already in the list of neighbors, or a `Vec` containing the updated list of neighbors if the list of neighbors is full. - pub fn add_to_neighbors(&mut self, node_id: u32, range: u32) -> Option> { - // Check if n is already in the graph entry - if self.neighbors.contains(&node_id) { - return None; - } - - let neighbor_len = self.neighbors.len(); - - // If not, check if the graph entry has enough space - if neighbor_len < (GRAPH_SLACK_FACTOR * range as f64) as usize { - // If yes, add n to the graph entry - self.neighbors.push(node_id); - return None; - } - - let mut copy_of_neighbors = Vec::with_capacity(neighbor_len + 1); - unsafe { - let dst = copy_of_neighbors.as_mut_ptr(); - std::ptr::copy_nonoverlapping(self.neighbors.as_ptr(), dst, neighbor_len); - dst.add(neighbor_len).write(node_id); - copy_of_neighbors.set_len(neighbor_len + 1); - } - - Some(copy_of_neighbors) - } -} - -#[cfg(test)] -mod vertex_and_neighbors_tests { - use crate::model::GRAPH_SLACK_FACTOR; - - use super::*; - - #[test] - fn test_set_with_capacity() { - let neighbors = VertexAndNeighbors::for_range(20, 10); - assert_eq!(neighbors.vertex_id, 20); - assert_eq!( - neighbors.neighbors.capacity(), - (10_f32 * GRAPH_SLACK_FACTOR as f32).ceil() as usize - ); - } - - #[test] - fn test_size() { - let mut neighbors = VertexAndNeighbors::for_range(20, 10); - - for i in 0..5 { - neighbors.neighbors.push(i); - } - - assert_eq!(neighbors.size(), 5); - } - - #[test] - fn test_set_neighbors() { - let mut neighbors = VertexAndNeighbors::for_range(20, 10); - let new_vec = AdjacencyList::from(vec![1, 2, 3, 4, 5]); - neighbors.set_neighbors(AdjacencyList::from(new_vec.clone())); - - assert_eq!(neighbors.neighbors, new_vec); - } - - #[test] - fn test_get_neighbors() { - let mut neighbors = VertexAndNeighbors::for_range(20, 10); - neighbors.set_neighbors(AdjacencyList::from(vec![1, 2, 3, 4, 5])); - let neighbor_ref = neighbors.get_neighbors(); - - assert!(std::ptr::eq(&neighbors.neighbors, neighbor_ref)) - } - - #[test] - fn test_add_to_neighbors() { - let mut neighbors = VertexAndNeighbors::for_range(20, 10); - - assert_eq!(neighbors.add_to_neighbors(1, 1), None); - assert_eq!(neighbors.neighbors, AdjacencyList::from(vec![1])); - - assert_eq!(neighbors.add_to_neighbors(1, 1), None); - assert_eq!(neighbors.neighbors, AdjacencyList::from(vec![1])); - - let ret = neighbors.add_to_neighbors(2, 1); - assert!(ret.is_some()); - assert_eq!(ret.unwrap(), vec![1, 2]); - assert_eq!(neighbors.neighbors, AdjacencyList::from(vec![1])); - - assert_eq!(neighbors.add_to_neighbors(2, 2), None); - assert_eq!(neighbors.neighbors, AdjacencyList::from(vec![1, 2])); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/mod.rs deleted file mode 100644 index a4f15ee..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod neighbor; -pub use neighbor::Neighbor; -pub use neighbor::NeighborPriorityQueue; - -pub mod data_store; -pub use data_store::InmemDataset; - -pub mod graph; -pub use graph::InMemoryGraph; -pub use graph::VertexAndNeighbors; - -pub mod configuration; -pub use configuration::*; - -pub mod scratch; -pub use scratch::*; - -pub mod vertex; -pub use vertex::Vertex; - -pub mod pq; -pub use pq::*; - -pub mod windows_aligned_file_reader; -pub use windows_aligned_file_reader::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/mod.rs deleted file mode 100644 index cd0dbad..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/mod.rs +++ /dev/null @@ -1,13 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod neighbor; -pub use neighbor::*; - -mod neighbor_priority_queue; -pub use neighbor_priority_queue::*; - -mod sorted_neighbor_vector; -pub use sorted_neighbor_vector::SortedNeighborVector; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor.rs deleted file mode 100644 index 8c712bc..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor.rs +++ /dev/null @@ -1,104 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::cmp::Ordering; - -/// Neighbor node -#[derive(Debug, Clone, Copy)] -pub struct Neighbor { - /// The id of the node - pub id: u32, - - /// The distance from the query node to current node - pub distance: f32, - - /// Whether the current is visited or not - pub visited: bool, -} - -impl Neighbor { - /// Create the neighbor node and it has not been visited - pub fn new (id: u32, distance: f32) -> Self { - Self { - id, - distance, - visited: false - } - } -} - -impl Default for Neighbor { - fn default() -> Self { - Self { id: 0, distance: 0.0_f32, visited: false } - } -} - -impl PartialEq for Neighbor { - #[inline] - fn eq(&self, other: &Self) -> bool { - self.id == other.id - } -} - -impl Eq for Neighbor {} - -impl Ord for Neighbor { - fn cmp(&self, other: &Self) -> Ordering { - let ord = self.distance.partial_cmp(&other.distance).unwrap_or(std::cmp::Ordering::Equal); - - if ord == Ordering::Equal { - return self.id.cmp(&other.id); - } - - ord - } -} - -impl PartialOrd for Neighbor { - #[inline] - fn lt(&self, other: &Self) -> bool { - self.distance < other.distance || (self.distance == other.distance && self.id < other.id) - } - - // Reason for allowing panic = "Does not support comparing Neighbor with partial_cmp" - #[allow(clippy::panic)] - fn partial_cmp(&self, _: &Self) -> Option { - panic!("Neighbor only allows eq and lt") - } -} - -#[cfg(test)] -mod neighbor_test { - use super::*; - - #[test] - fn eq_lt_works() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - let n3 = Neighbor::new(1, 1.1); - - assert!(n1 != n2); - assert!(n1 < n2); - assert!(n1 == n3); - } - - #[test] - #[should_panic] - fn gt_should_panic() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - - assert!(n2 > n1); - } - - #[test] - #[should_panic] - fn le_should_panic() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - - assert!(n1 <= n2); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor_priority_queue.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor_priority_queue.rs deleted file mode 100644 index 81b1610..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/neighbor_priority_queue.rs +++ /dev/null @@ -1,241 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::model::Neighbor; - -/// Neighbor priority Queue based on the distance to the query node -#[derive(Debug)] -pub struct NeighborPriorityQueue { - /// The size of the priority queue - size: usize, - - /// The capacity of the priority queue - capacity: usize, - - /// The current notvisited neighbor whose distance is smallest among all notvisited neighbor - cur: usize, - - /// The neighbor collection - data: Vec, -} - -impl Default for NeighborPriorityQueue { - fn default() -> Self { - Self::new() - } -} - -impl NeighborPriorityQueue { - /// Create NeighborPriorityQueue without capacity - pub fn new() -> Self { - Self { - size: 0, - capacity: 0, - cur: 0, - data: Vec::new(), - } - } - - /// Create NeighborPriorityQueue with capacity - pub fn with_capacity(capacity: usize) -> Self { - Self { - size: 0, - capacity, - cur: 0, - data: vec![Neighbor::default(); capacity + 1], - } - } - - /// Inserts item with order. - /// The item will be dropped if queue is full / already exist in queue / it has a greater distance than the last item. - /// The set cursor that is used to pop() the next item will be set to the lowest index of an uncheck item. - pub fn insert(&mut self, nbr: Neighbor) { - if self.size == self.capacity && self.get_at(self.size - 1) < &nbr { - return; - } - - let mut lo = 0; - let mut hi = self.size; - while lo < hi { - let mid = (lo + hi) >> 1; - if &nbr < self.get_at(mid) { - hi = mid; - } else if self.get_at(mid).id == nbr.id { - // Make sure the same neighbor isn't inserted into the set - return; - } else { - lo = mid + 1; - } - } - - if lo < self.capacity { - self.data.copy_within(lo..self.size, lo + 1); - } - self.data[lo] = Neighbor::new(nbr.id, nbr.distance); - if self.size < self.capacity { - self.size += 1; - } - if lo < self.cur { - self.cur = lo; - } - } - - /// Get the neighbor at index - SAFETY: index must be less than size - fn get_at(&self, index: usize) -> &Neighbor { - unsafe { self.data.get_unchecked(index) } - } - - /// Get the closest and notvisited neighbor - pub fn closest_notvisited(&mut self) -> Neighbor { - self.data[self.cur].visited = true; - let pre = self.cur; - while self.cur < self.size && self.get_at(self.cur).visited { - self.cur += 1; - } - self.data[pre] - } - - /// Whether there is notvisited node or not - pub fn has_notvisited_node(&self) -> bool { - self.cur < self.size - } - - /// Get the size of the NeighborPriorityQueue - pub fn size(&self) -> usize { - self.size - } - - /// Get the capacity of the NeighborPriorityQueue - pub fn capacity(&self) -> usize { - self.capacity - } - - /// Sets an artificial capacity of the NeighborPriorityQueue. For benchmarking purposes only. - pub fn set_capacity(&mut self, capacity: usize) { - if capacity < self.data.len() { - self.capacity = capacity; - } - } - - /// Reserve capacity - pub fn reserve(&mut self, capacity: usize) { - if capacity > self.capacity { - self.data.resize(capacity + 1, Neighbor::default()); - self.capacity = capacity; - } - } - - /// Set size and cur to 0 - pub fn clear(&mut self) { - self.size = 0; - self.cur = 0; - } -} - -impl std::ops::Index for NeighborPriorityQueue { - type Output = Neighbor; - - fn index(&self, i: usize) -> &Self::Output { - &self.data[i] - } -} - -#[cfg(test)] -mod neighbor_priority_queue_test { - use super::*; - - #[test] - fn test_reserve_capacity() { - let mut queue = NeighborPriorityQueue::with_capacity(10); - assert_eq!(queue.capacity(), 10); - queue.reserve(20); - assert_eq!(queue.capacity(), 20); - } - - #[test] - fn test_insert() { - let mut queue = NeighborPriorityQueue::with_capacity(3); - assert_eq!(queue.size(), 0); - queue.insert(Neighbor::new(1, 1.0)); - queue.insert(Neighbor::new(2, 0.5)); - assert_eq!(queue.size(), 2); - queue.insert(Neighbor::new(2, 0.5)); // should be ignored as the same neighbor - assert_eq!(queue.size(), 2); - queue.insert(Neighbor::new(3, 0.9)); - assert_eq!(queue.size(), 3); - assert_eq!(queue[2].id, 1); - queue.insert(Neighbor::new(4, 2.0)); // should be dropped as queue is full and distance is greater than last item - assert_eq!(queue.size(), 3); - assert_eq!(queue[0].id, 2); // node id in queue should be [2,3,1] - assert_eq!(queue[1].id, 3); - assert_eq!(queue[2].id, 1); - println!("{:?}", queue); - } - - #[test] - fn test_index() { - let mut queue = NeighborPriorityQueue::with_capacity(3); - queue.insert(Neighbor::new(1, 1.0)); - queue.insert(Neighbor::new(2, 0.5)); - queue.insert(Neighbor::new(3, 1.5)); - assert_eq!(queue[0].id, 2); - assert_eq!(queue[0].distance, 0.5); - } - - #[test] - fn test_visit() { - let mut queue = NeighborPriorityQueue::with_capacity(3); - queue.insert(Neighbor::new(1, 1.0)); - queue.insert(Neighbor::new(2, 0.5)); - queue.insert(Neighbor::new(3, 1.5)); // node id in queue should be [2,1,3] - assert!(queue.has_notvisited_node()); - let nbr = queue.closest_notvisited(); - assert_eq!(nbr.id, 2); - assert_eq!(nbr.distance, 0.5); - assert!(nbr.visited); - assert!(queue.has_notvisited_node()); - let nbr = queue.closest_notvisited(); - assert_eq!(nbr.id, 1); - assert_eq!(nbr.distance, 1.0); - assert!(nbr.visited); - assert!(queue.has_notvisited_node()); - let nbr = queue.closest_notvisited(); - assert_eq!(nbr.id, 3); - assert_eq!(nbr.distance, 1.5); - assert!(nbr.visited); - assert!(!queue.has_notvisited_node()); - } - - #[test] - fn test_clear_queue() { - let mut queue = NeighborPriorityQueue::with_capacity(3); - queue.insert(Neighbor::new(1, 1.0)); - queue.insert(Neighbor::new(2, 0.5)); - assert_eq!(queue.size(), 2); - assert!(queue.has_notvisited_node()); - queue.clear(); - assert_eq!(queue.size(), 0); - assert!(!queue.has_notvisited_node()); - } - - #[test] - fn test_reserve() { - let mut queue = NeighborPriorityQueue::new(); - queue.reserve(10); - assert_eq!(queue.data.len(), 11); - assert_eq!(queue.capacity, 10); - } - - #[test] - fn test_set_capacity() { - let mut queue = NeighborPriorityQueue::with_capacity(10); - queue.set_capacity(5); - assert_eq!(queue.capacity, 5); - assert_eq!(queue.data.len(), 11); - - queue.set_capacity(11); - assert_eq!(queue.capacity, 5); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/sorted_neighbor_vector.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/sorted_neighbor_vector.rs deleted file mode 100644 index 4c3eff0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/neighbor/sorted_neighbor_vector.rs +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Sorted Neighbor Vector - -use std::ops::{Deref, DerefMut}; - -use super::Neighbor; - -/// A newtype on top of vector of neighbors, is sorted by distance -#[derive(Debug)] -pub struct SortedNeighborVector<'a>(&'a mut Vec); - -impl<'a> SortedNeighborVector<'a> { - /// Create a new SortedNeighborVector - pub fn new(vec: &'a mut Vec) -> Self { - vec.sort_unstable(); - Self(vec) - } -} - -impl<'a> Deref for SortedNeighborVector<'a> { - type Target = Vec; - - fn deref(&self) -> &Self::Target { - self.0 - } -} - -impl<'a> DerefMut for SortedNeighborVector<'a> { - fn deref_mut(&mut self) -> &mut Self::Target { - self.0 - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/fixed_chunk_pq_table.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/fixed_chunk_pq_table.rs deleted file mode 100644 index bfedcae..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/fixed_chunk_pq_table.rs +++ /dev/null @@ -1,483 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations)] - -use hashbrown::HashMap; -use rayon::prelude::{ - IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator, ParallelSliceMut, -}; -use std::arch::x86_64::{_mm_prefetch, _MM_HINT_T0}; - -use crate::{ - common::{ANNError, ANNResult}, - model::NUM_PQ_CENTROIDS, -}; - -/// PQ Pivot table loading and calculate distance -#[derive(Debug)] -pub struct FixedChunkPQTable { - /// pq_tables = float array of size [256 * ndims] - pq_table: Vec, - - /// ndims = true dimension of vectors - dim: usize, - - /// num_pq_chunks = the pq chunk number - num_pq_chunks: usize, - - /// chunk_offsets = the offset of each chunk, start from 0 - chunk_offsets: Vec, - - /// centroid of each dimension - centroids: Vec, - - /// Becasue we're using L2 distance, this is no needed now. - /// Transport of pq_table. transport_pq_table = float array of size [ndims * 256]. - /// e.g. if pa_table is 2 centroids * 3 dims - /// [ 1, 2, 3, - /// 4, 5, 6] - /// then transport_pq_table would be 3 dims * 2 centroids - /// [ 1, 4, - /// 2, 5, - /// 3, 6] - /// transport_pq_table: Vec, - - /// Map dim offset to chunk index e.g., 8 dims in to 2 chunks - /// then would be [(0,0), (1,0), (2,0), (3,0), (4,1), (5,1), (6,1), (7,1)] - dimoffset_chunk_mapping: HashMap, -} - -impl FixedChunkPQTable { - /// Create the FixedChunkPQTable with dim and chunk numbers and pivot file data (pivot table + cenroids + chunk offsets) - pub fn new( - dim: usize, - num_pq_chunks: usize, - pq_table: Vec, - centroids: Vec, - chunk_offsets: Vec, - ) -> Self { - let mut dimoffset_chunk_mapping = HashMap::new(); - for chunk_index in 0..num_pq_chunks { - for dim_offset in chunk_offsets[chunk_index]..chunk_offsets[chunk_index + 1] { - dimoffset_chunk_mapping.insert(dim_offset, chunk_index); - } - } - - Self { - pq_table, - dim, - num_pq_chunks, - chunk_offsets, - centroids, - dimoffset_chunk_mapping, - } - } - - /// Get chunk number - pub fn get_num_chunks(&self) -> usize { - self.num_pq_chunks - } - - /// Shifting the query according to mean or the whole corpus - pub fn preprocess_query(&self, query_vec: &mut [f32]) { - for (query, ¢roid) in query_vec.iter_mut().zip(self.centroids.iter()) { - *query -= centroid; - } - } - - /// Pre-calculated the distance between query and each centroid by l2 distance - /// * `query_vec` - query vector: 1 * dim - /// * `dist_vec` - pre-calculated the distance between query and each centroid: chunk_size * num_centroids - #[allow(clippy::needless_range_loop)] - pub fn populate_chunk_distances(&self, query_vec: &[f32]) -> Vec { - let mut dist_vec = vec![0.0; self.num_pq_chunks * NUM_PQ_CENTROIDS]; - for centroid_index in 0..NUM_PQ_CENTROIDS { - for chunk_index in 0..self.num_pq_chunks { - for dim_offset in - self.chunk_offsets[chunk_index]..self.chunk_offsets[chunk_index + 1] - { - let diff: f32 = self.pq_table[self.dim * centroid_index + dim_offset] - - query_vec[dim_offset]; - dist_vec[chunk_index * NUM_PQ_CENTROIDS + centroid_index] += diff * diff; - } - } - } - dist_vec - } - - /// Pre-calculated the distance between query and each centroid by inner product - /// * `query_vec` - query vector: 1 * dim - /// * `dist_vec` - pre-calculated the distance between query and each centroid: chunk_size * num_centroids - /// - /// Reason to allow clippy::needless_range_loop: - /// The inner loop is operating over a range that is different for each iteration of the outer loop. - /// This isn't a scenario where using iter().enumerate() would be easily applicable, - /// because the inner loop isn't iterating directly over the contents of a slice or array. - /// Thus, using indexing might be the most straightforward way to express this logic. - #[allow(clippy::needless_range_loop)] - pub fn populate_chunk_inner_products(&self, query_vec: &[f32]) -> Vec { - let mut dist_vec = vec![0.0; self.num_pq_chunks * NUM_PQ_CENTROIDS]; - for centroid_index in 0..NUM_PQ_CENTROIDS { - for chunk_index in 0..self.num_pq_chunks { - for dim_offset in - self.chunk_offsets[chunk_index]..self.chunk_offsets[chunk_index + 1] - { - // assumes that we are not shifting the vectors to mean zero, i.e., centroid - // array should be all zeros returning negative to keep the search code - // clean (max inner product vs min distance) - let diff: f32 = self.pq_table[self.dim * centroid_index + dim_offset] - * query_vec[dim_offset]; - dist_vec[chunk_index * NUM_PQ_CENTROIDS + centroid_index] -= diff; - } - } - } - dist_vec - } - - /// Calculate the distance between query and given centroid by l2 distance - /// * `query_vec` - query vector: 1 * dim - /// * `base_vec` - given centroid array: 1 * num_pq_chunks - #[allow(clippy::needless_range_loop)] - pub fn l2_distance(&self, query_vec: &[f32], base_vec: &[u8]) -> f32 { - let mut res_vec: Vec = vec![0.0; self.num_pq_chunks]; - res_vec - .par_iter_mut() - .enumerate() - .for_each(|(chunk_index, chunk_diff)| { - for dim_offset in - self.chunk_offsets[chunk_index]..self.chunk_offsets[chunk_index + 1] - { - let diff = self.pq_table - [self.dim * base_vec[chunk_index] as usize + dim_offset] - - query_vec[dim_offset]; - *chunk_diff += diff * diff; - } - }); - - let res: f32 = res_vec.iter().sum::(); - - res - } - - /// Calculate the distance between query and given centroid by inner product - /// * `query_vec` - query vector: 1 * dim - /// * `base_vec` - given centroid array: 1 * num_pq_chunks - #[allow(clippy::needless_range_loop)] - pub fn inner_product(&self, query_vec: &[f32], base_vec: &[u8]) -> f32 { - let mut res_vec: Vec = vec![0.0; self.num_pq_chunks]; - res_vec - .par_iter_mut() - .enumerate() - .for_each(|(chunk_index, chunk_diff)| { - for dim_offset in - self.chunk_offsets[chunk_index]..self.chunk_offsets[chunk_index + 1] - { - *chunk_diff += self.pq_table - [self.dim * base_vec[chunk_index] as usize + dim_offset] - * query_vec[dim_offset]; - } - }); - - let res: f32 = res_vec.iter().sum::(); - - // returns negative value to simulate distances (max -> min conversion) - -res - } - - /// Revert vector by adding centroid - /// * `base_vec` - given centroid array: 1 * num_pq_chunks - /// * `out_vec` - reverted vector - pub fn inflate_vector(&self, base_vec: &[u8]) -> ANNResult> { - let mut out_vec: Vec = vec![0.0; self.dim]; - for (dim_offset, value) in out_vec.iter_mut().enumerate() { - let chunk_index = - self.dimoffset_chunk_mapping - .get(&dim_offset) - .ok_or(ANNError::log_pq_error( - "ERROR: dim_offset not found in dimoffset_chunk_mapping".to_string(), - ))?; - *value = self.pq_table[self.dim * base_vec[*chunk_index] as usize + dim_offset] - + self.centroids[dim_offset]; - } - - Ok(out_vec) - } -} - -/// Given a batch input nodes, return a batch of PQ distance -/// * `pq_ids` - batch nodes: n_pts * pq_nchunks -/// * `n_pts` - batch number -/// * `pq_nchunks` - pq chunk number number -/// * `pq_dists` - pre-calculated the distance between query and each centroid: chunk_size * num_centroids -/// * `dists_out` - n_pts * 1 -pub fn pq_dist_lookup( - pq_ids: &[u8], - n_pts: usize, - pq_nchunks: usize, - pq_dists: &[f32], -) -> Vec { - let mut dists_out: Vec = vec![0.0; n_pts]; - unsafe { - _mm_prefetch(dists_out.as_ptr() as *const i8, _MM_HINT_T0); - _mm_prefetch(pq_ids.as_ptr() as *const i8, _MM_HINT_T0); - _mm_prefetch(pq_ids.as_ptr().add(64) as *const i8, _MM_HINT_T0); - _mm_prefetch(pq_ids.as_ptr().add(128) as *const i8, _MM_HINT_T0); - } - for chunk in 0..pq_nchunks { - let chunk_dists = &pq_dists[256 * chunk..]; - if chunk < pq_nchunks - 1 { - unsafe { - _mm_prefetch( - chunk_dists.as_ptr().offset(256 * chunk as isize).add(256) as *const i8, - _MM_HINT_T0, - ); - } - } - dists_out - .par_iter_mut() - .enumerate() - .for_each(|(n_iter, dist)| { - let pq_centerid = pq_ids[pq_nchunks * n_iter + chunk]; - *dist += chunk_dists[pq_centerid as usize]; - }); - } - dists_out -} - -pub fn aggregate_coords(ids: &[u32], all_coords: &[u8], ndims: usize) -> Vec { - let mut out: Vec = vec![0u8; ids.len() * ndims]; - let ndim_u32 = ndims as u32; - out.par_chunks_mut(ndims) - .enumerate() - .for_each(|(index, chunk)| { - let id_compressed_pivot = &all_coords - [(ids[index] * ndim_u32) as usize..(ids[index] * ndim_u32 + ndim_u32) as usize]; - let temp_slice = - unsafe { std::slice::from_raw_parts(id_compressed_pivot.as_ptr(), ndims) }; - chunk.copy_from_slice(temp_slice); - }); - - out -} - -#[cfg(test)] -mod fixed_chunk_pq_table_test { - - use super::*; - use crate::common::{ANNError, ANNResult}; - use crate::utils::{convert_types_u32_usize, convert_types_u64_usize, file_exists, load_bin}; - - const DIM: usize = 128; - - #[test] - fn load_pivot_test() { - let pq_pivots_path: &str = "tests/data/siftsmall_learn.bin_pq_pivots.bin"; - let (dim, pq_table, centroids, chunk_offsets) = - load_pq_pivots_bin(pq_pivots_path, &1).unwrap(); - let fixed_chunk_pq_table = - FixedChunkPQTable::new(dim, 1, pq_table, centroids, chunk_offsets); - - assert_eq!(dim, DIM); - assert_eq!(fixed_chunk_pq_table.pq_table.len(), DIM * NUM_PQ_CENTROIDS); - assert_eq!(fixed_chunk_pq_table.centroids.len(), DIM); - - assert_eq!(fixed_chunk_pq_table.chunk_offsets[0], 0); - assert_eq!(fixed_chunk_pq_table.chunk_offsets[1], DIM); - assert_eq!(fixed_chunk_pq_table.chunk_offsets.len(), 2); - } - - #[test] - fn get_num_chunks_test() { - let num_chunks = 7; - let pa_table = vec![0.0; DIM * NUM_PQ_CENTROIDS]; - let centroids = vec![0.0; DIM]; - let chunk_offsets = vec![0, 7, 9, 11, 22, 34, 78, 127]; - let fixed_chunk_pq_table = - FixedChunkPQTable::new(DIM, num_chunks, pa_table, centroids, chunk_offsets); - let chunk: usize = fixed_chunk_pq_table.get_num_chunks(); - assert_eq!(chunk, num_chunks); - } - - #[test] - fn preprocess_query_test() { - let pq_pivots_path: &str = "tests/data/siftsmall_learn.bin_pq_pivots.bin"; - let (dim, pq_table, centroids, chunk_offsets) = - load_pq_pivots_bin(pq_pivots_path, &1).unwrap(); - let fixed_chunk_pq_table = - FixedChunkPQTable::new(dim, 1, pq_table, centroids, chunk_offsets); - - let mut query_vec: Vec = vec![ - 32.39f32, 78.57f32, 50.32f32, 80.46f32, 6.47f32, 69.76f32, 94.2f32, 83.36f32, 5.8f32, - 68.78f32, 42.32f32, 61.77f32, 90.26f32, 60.41f32, 3.86f32, 61.21f32, 16.6f32, 54.46f32, - 7.29f32, 54.24f32, 92.49f32, 30.18f32, 65.36f32, 99.09f32, 3.8f32, 36.4f32, 86.72f32, - 65.18f32, 29.87f32, 62.21f32, 58.32f32, 43.23f32, 94.3f32, 79.61f32, 39.67f32, - 11.18f32, 48.88f32, 38.19f32, 93.95f32, 10.46f32, 36.7f32, 14.75f32, 81.64f32, - 59.18f32, 99.03f32, 74.23f32, 1.26f32, 82.69f32, 35.7f32, 38.39f32, 46.17f32, 64.75f32, - 7.15f32, 36.55f32, 77.32f32, 18.65f32, 32.8f32, 74.84f32, 18.12f32, 20.19f32, 70.06f32, - 48.37f32, 40.18f32, 45.69f32, 88.3f32, 39.15f32, 60.97f32, 71.29f32, 61.79f32, - 47.23f32, 94.71f32, 58.04f32, 52.4f32, 34.66f32, 59.1f32, 47.11f32, 30.2f32, 58.72f32, - 74.35f32, 83.68f32, 66.8f32, 28.57f32, 29.45f32, 52.02f32, 91.95f32, 92.44f32, - 65.25f32, 38.3f32, 35.6f32, 41.67f32, 91.33f32, 76.81f32, 74.88f32, 33.17f32, 48.36f32, - 41.42f32, 23f32, 8.31f32, 81.69f32, 80.08f32, 50.55f32, 54.46f32, 23.79f32, 43.46f32, - 84.5f32, 10.42f32, 29.51f32, 19.73f32, 46.48f32, 35.01f32, 52.3f32, 66.97f32, 4.8f32, - 74.81f32, 2.82f32, 61.82f32, 25.06f32, 17.3f32, 17.29f32, 63.2f32, 64.1f32, 61.68f32, - 37.42f32, 3.39f32, 97.45f32, 5.32f32, 59.02f32, 35.6f32, - ]; - fixed_chunk_pq_table.preprocess_query(&mut query_vec); - assert_eq!(query_vec[0], 32.39f32 - fixed_chunk_pq_table.centroids[0]); - assert_eq!( - query_vec[127], - 35.6f32 - fixed_chunk_pq_table.centroids[127] - ); - } - - #[test] - fn calculate_distances_tests() { - let pq_pivots_path: &str = "tests/data/siftsmall_learn.bin_pq_pivots.bin"; - - let (dim, pq_table, centroids, chunk_offsets) = - load_pq_pivots_bin(pq_pivots_path, &1).unwrap(); - let fixed_chunk_pq_table = - FixedChunkPQTable::new(dim, 1, pq_table, centroids, chunk_offsets); - - let query_vec: Vec = vec![ - 32.39f32, 78.57f32, 50.32f32, 80.46f32, 6.47f32, 69.76f32, 94.2f32, 83.36f32, 5.8f32, - 68.78f32, 42.32f32, 61.77f32, 90.26f32, 60.41f32, 3.86f32, 61.21f32, 16.6f32, 54.46f32, - 7.29f32, 54.24f32, 92.49f32, 30.18f32, 65.36f32, 99.09f32, 3.8f32, 36.4f32, 86.72f32, - 65.18f32, 29.87f32, 62.21f32, 58.32f32, 43.23f32, 94.3f32, 79.61f32, 39.67f32, - 11.18f32, 48.88f32, 38.19f32, 93.95f32, 10.46f32, 36.7f32, 14.75f32, 81.64f32, - 59.18f32, 99.03f32, 74.23f32, 1.26f32, 82.69f32, 35.7f32, 38.39f32, 46.17f32, 64.75f32, - 7.15f32, 36.55f32, 77.32f32, 18.65f32, 32.8f32, 74.84f32, 18.12f32, 20.19f32, 70.06f32, - 48.37f32, 40.18f32, 45.69f32, 88.3f32, 39.15f32, 60.97f32, 71.29f32, 61.79f32, - 47.23f32, 94.71f32, 58.04f32, 52.4f32, 34.66f32, 59.1f32, 47.11f32, 30.2f32, 58.72f32, - 74.35f32, 83.68f32, 66.8f32, 28.57f32, 29.45f32, 52.02f32, 91.95f32, 92.44f32, - 65.25f32, 38.3f32, 35.6f32, 41.67f32, 91.33f32, 76.81f32, 74.88f32, 33.17f32, 48.36f32, - 41.42f32, 23f32, 8.31f32, 81.69f32, 80.08f32, 50.55f32, 54.46f32, 23.79f32, 43.46f32, - 84.5f32, 10.42f32, 29.51f32, 19.73f32, 46.48f32, 35.01f32, 52.3f32, 66.97f32, 4.8f32, - 74.81f32, 2.82f32, 61.82f32, 25.06f32, 17.3f32, 17.29f32, 63.2f32, 64.1f32, 61.68f32, - 37.42f32, 3.39f32, 97.45f32, 5.32f32, 59.02f32, 35.6f32, - ]; - - let dist_vec = fixed_chunk_pq_table.populate_chunk_distances(&query_vec); - assert_eq!(dist_vec.len(), 256); - - // populate_chunk_distances_test - let mut sampled_output = 0.0; - (0..DIM).for_each(|dim_offset| { - let diff = fixed_chunk_pq_table.pq_table[dim_offset] - query_vec[dim_offset]; - sampled_output += diff * diff; - }); - assert_eq!(sampled_output, dist_vec[0]); - - // populate_chunk_inner_products_test - let dist_vec = fixed_chunk_pq_table.populate_chunk_inner_products(&query_vec); - assert_eq!(dist_vec.len(), 256); - - let mut sampled_output = 0.0; - (0..DIM).for_each(|dim_offset| { - sampled_output -= fixed_chunk_pq_table.pq_table[dim_offset] * query_vec[dim_offset]; - }); - assert_eq!(sampled_output, dist_vec[0]); - - // l2_distance_test - let base_vec: Vec = vec![3u8]; - let dist = fixed_chunk_pq_table.l2_distance(&query_vec, &base_vec); - let mut l2_output = 0.0; - (0..DIM).for_each(|dim_offset| { - let diff = fixed_chunk_pq_table.pq_table[3 * DIM + dim_offset] - query_vec[dim_offset]; - l2_output += diff * diff; - }); - assert_eq!(l2_output, dist); - - // inner_product_test - let dist = fixed_chunk_pq_table.inner_product(&query_vec, &base_vec); - let mut l2_output = 0.0; - (0..DIM).for_each(|dim_offset| { - l2_output -= - fixed_chunk_pq_table.pq_table[3 * DIM + dim_offset] * query_vec[dim_offset]; - }); - assert_eq!(l2_output, dist); - - // inflate_vector_test - let inflate_vector = fixed_chunk_pq_table.inflate_vector(&base_vec).unwrap(); - assert_eq!(inflate_vector.len(), DIM); - assert_eq!( - inflate_vector[0], - fixed_chunk_pq_table.pq_table[3 * DIM] + fixed_chunk_pq_table.centroids[0] - ); - assert_eq!( - inflate_vector[1], - fixed_chunk_pq_table.pq_table[3 * DIM + 1] + fixed_chunk_pq_table.centroids[1] - ); - assert_eq!( - inflate_vector[127], - fixed_chunk_pq_table.pq_table[3 * DIM + 127] + fixed_chunk_pq_table.centroids[127] - ); - } - - fn load_pq_pivots_bin( - pq_pivots_path: &str, - num_pq_chunks: &usize, - ) -> ANNResult<(usize, Vec, Vec, Vec)> { - if !file_exists(pq_pivots_path) { - return Err(ANNError::log_pq_error( - "ERROR: PQ k-means pivot file not found.".to_string(), - )); - } - - let (data, offset_num, offset_dim) = load_bin::(pq_pivots_path, 0)?; - let file_offset_data = convert_types_u64_usize(&data, offset_num, offset_dim); - if offset_num != 4 { - let error_message = format!("Error reading pq_pivots file {}. Offsets don't contain correct metadata, # offsets = {}, but expecting 4.", pq_pivots_path, offset_num); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, pq_center_num, dim) = load_bin::(pq_pivots_path, file_offset_data[0])?; - let pq_table = data.to_vec(); - if pq_center_num != NUM_PQ_CENTROIDS { - let error_message = format!( - "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", - pq_pivots_path, pq_center_num, NUM_PQ_CENTROIDS - ); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, centroid_dim, nc) = load_bin::(pq_pivots_path, file_offset_data[1])?; - let centroids = data.to_vec(); - if centroid_dim != dim || nc != 1 { - let error_message = format!("Error reading pq_pivots file {}. file_dim = {}, file_cols = {} but expecting {} entries in 1 dimension.", pq_pivots_path, centroid_dim, nc, dim); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, chunk_offset_num, nc) = load_bin::(pq_pivots_path, file_offset_data[2])?; - let chunk_offsets = convert_types_u32_usize(&data, chunk_offset_num, nc); - if chunk_offset_num != num_pq_chunks + 1 || nc != 1 { - let error_message = format!("Error reading pq_pivots file at chunk offsets; file has nr={}, nc={} but expecting nr={} and nc=1.", chunk_offset_num, nc, num_pq_chunks + 1); - return Err(ANNError::log_pq_error(error_message)); - } - - Ok((dim, pq_table, centroids, chunk_offsets)) - } -} - -#[cfg(test)] -mod pq_index_prune_query_test { - - use super::*; - - #[test] - fn pq_dist_lookup_test() { - let pq_ids: Vec = vec![1u8, 3u8, 2u8, 2u8]; - let mut pq_dists: Vec = Vec::with_capacity(256 * 2); - for _ in 0..pq_dists.capacity() { - pq_dists.push(rand::random()); - } - - let dists_out = pq_dist_lookup(&pq_ids, 2, 2, &pq_dists); - assert_eq!(dists_out.len(), 2); - assert_eq!(dists_out[0], pq_dists[0 + 1] + pq_dists[256 + 3]); - assert_eq!(dists_out[1], pq_dists[0 + 2] + pq_dists[256 + 2]); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/mod.rs deleted file mode 100644 index 85daaa7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod fixed_chunk_pq_table; -pub use fixed_chunk_pq_table::*; - -mod pq_construction; -pub use pq_construction::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/pq_construction.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/pq_construction.rs deleted file mode 100644 index 0a7b078..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/pq/pq_construction.rs +++ /dev/null @@ -1,398 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations)] - -use rayon::prelude::{IndexedParallelIterator, ParallelIterator}; -use rayon::slice::ParallelSliceMut; - -use crate::common::{ANNError, ANNResult}; -use crate::storage::PQStorage; -use crate::utils::{compute_closest_centers, file_exists, k_means_clustering}; - -/// Max size of PQ training set -pub const MAX_PQ_TRAINING_SET_SIZE: f64 = 256_000f64; - -/// Max number of PQ chunks -pub const MAX_PQ_CHUNKS: usize = 512; - -pub const NUM_PQ_CENTROIDS: usize = 256; -/// block size for reading/processing large files and matrices in blocks -const BLOCK_SIZE: usize = 5000000; -const NUM_KMEANS_REPS_PQ: usize = 12; - -/// given training data in train_data of dimensions num_train * dim, generate -/// PQ pivots using k-means algorithm to partition the co-ordinates into -/// num_pq_chunks (if it divides dimension, else rounded) chunks, and runs -/// k-means in each chunk to compute the PQ pivots and stores in bin format in -/// file pq_pivots_path as a s num_centers*dim floating point binary file -/// PQ pivot table layout: {pivot offsets data: METADATA_SIZE}{pivot vector:[dim; num_centroid]}{centroid vector:[dim; 1]}{chunk offsets:[chunk_num+1; 1]} -fn generate_pq_pivots( - train_data: &mut [f32], - num_train: usize, - dim: usize, - num_centers: usize, - num_pq_chunks: usize, - max_k_means_reps: usize, - pq_storage: &mut PQStorage, -) -> ANNResult<()> { - if num_pq_chunks > dim { - return Err(ANNError::log_pq_error( - "Error: number of chunks more than dimension.".to_string(), - )); - } - - if pq_storage.pivot_data_exist() { - let (file_num_centers, file_dim) = pq_storage.read_pivot_metadata()?; - if file_dim == dim && file_num_centers == num_centers { - // PQ pivot file exists. Not generating again. - return Ok(()); - } - } - - // Calculate centroid and center the training data - // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - // Now, we're using L2 as default. - let mut centroid: Vec = vec![0.0; dim]; - for dim_index in 0..dim { - for train_data_index in 0..num_train { - centroid[dim_index] += train_data[train_data_index * dim + dim_index]; - } - centroid[dim_index] /= num_train as f32; - } - for dim_index in 0..dim { - for train_data_index in 0..num_train { - train_data[train_data_index * dim + dim_index] -= centroid[dim_index]; - } - } - - // Calculate each chunk's offset - // If we have 8 dimension and 3 chunk then offsets would be [0,3,6,8] - let mut chunk_offsets: Vec = vec![0; num_pq_chunks + 1]; - let mut chunk_offset: usize = 0; - for chunk_index in 0..num_pq_chunks { - chunk_offset += dim / num_pq_chunks; - if chunk_index < (dim % num_pq_chunks) { - chunk_offset += 1; - } - chunk_offsets[chunk_index + 1] = chunk_offset; - } - - let mut full_pivot_data: Vec = vec![0.0; num_centers * dim]; - for chunk_index in 0..num_pq_chunks { - let chunk_size = chunk_offsets[chunk_index + 1] - chunk_offsets[chunk_index]; - - let mut cur_train_data: Vec = vec![0.0; num_train * chunk_size]; - let mut cur_pivot_data: Vec = vec![0.0; num_centers * chunk_size]; - - cur_train_data - .par_chunks_mut(chunk_size) - .enumerate() - .for_each(|(train_data_index, chunk)| { - for (dim_offset, item) in chunk.iter_mut().enumerate() { - *item = train_data - [train_data_index * dim + chunk_offsets[chunk_index] + dim_offset]; - } - }); - - // Run kmeans to get the centroids of this chunk. - let (_closest_docs, _closest_center, _residual) = k_means_clustering( - &cur_train_data, - num_train, - chunk_size, - &mut cur_pivot_data, - num_centers, - max_k_means_reps, - )?; - - // Copy centroids from this chunk table to full table - for center_index in 0..num_centers { - full_pivot_data[center_index * dim + chunk_offsets[chunk_index] - ..center_index * dim + chunk_offsets[chunk_index + 1]] - .copy_from_slice( - &cur_pivot_data[center_index * chunk_size..(center_index + 1) * chunk_size], - ); - } - } - - pq_storage.write_pivot_data( - &full_pivot_data, - ¢roid, - &chunk_offsets, - num_centers, - dim, - )?; - - Ok(()) -} - -/// streams the base file (data_file), and computes the closest centers in each -/// chunk to generate the compressed data_file and stores it in -/// pq_compressed_vectors_path. -/// If the numbber of centers is < 256, it stores as byte vector, else as -/// 4-byte vector in binary format. -/// Compressed PQ table layout: {num_points: usize}{num_chunks: usize}{compressed pq table: [num_points; num_chunks]} -fn generate_pq_data_from_pivots>( - num_centers: usize, - num_pq_chunks: usize, - pq_storage: &mut PQStorage, -) -> ANNResult<()> { - let (num_points, dim) = pq_storage.read_pq_data_metadata()?; - - let full_pivot_data: Vec; - let centroid: Vec; - let chunk_offsets: Vec; - - if !pq_storage.pivot_data_exist() { - return Err(ANNError::log_pq_error( - "ERROR: PQ k-means pivot file not found.".to_string(), - )); - } else { - (full_pivot_data, centroid, chunk_offsets) = - pq_storage.load_pivot_data(&num_pq_chunks, &num_centers, &dim)?; - } - - pq_storage.write_compressed_pivot_metadata(num_points as i32, num_pq_chunks as i32)?; - - let block_size = if num_points <= BLOCK_SIZE { - num_points - } else { - BLOCK_SIZE - }; - let num_blocks = (num_points / block_size) + (num_points % block_size != 0) as usize; - - for block_index in 0..num_blocks { - let start_index: usize = block_index * block_size; - let end_index: usize = std::cmp::min((block_index + 1) * block_size, num_points); - let cur_block_size: usize = end_index - start_index; - - let mut block_compressed_base: Vec = vec![0; cur_block_size * num_pq_chunks]; - - let block_data: Vec = pq_storage.read_pq_block_data(cur_block_size, dim)?; - - let mut adjusted_block_data: Vec = vec![0.0; cur_block_size * dim]; - - for block_data_index in 0..cur_block_size { - for dim_index in 0..dim { - adjusted_block_data[block_data_index * dim + dim_index] = - block_data[block_data_index * dim + dim_index].into() - centroid[dim_index]; - } - } - - for chunk_index in 0..num_pq_chunks { - let cur_chunk_size = chunk_offsets[chunk_index + 1] - chunk_offsets[chunk_index]; - if cur_chunk_size == 0 { - continue; - } - - let mut cur_pivot_data: Vec = vec![0.0; num_centers * cur_chunk_size]; - let mut cur_data: Vec = vec![0.0; cur_block_size * cur_chunk_size]; - let mut closest_center: Vec = vec![0; cur_block_size]; - - // Divide the data into chunks and process each chunk in parallel. - cur_data - .par_chunks_mut(cur_chunk_size) - .enumerate() - .for_each(|(block_data_index, chunk)| { - for (dim_offset, item) in chunk.iter_mut().enumerate() { - *item = adjusted_block_data - [block_data_index * dim + chunk_offsets[chunk_index] + dim_offset]; - } - }); - - cur_pivot_data - .par_chunks_mut(cur_chunk_size) - .enumerate() - .for_each(|(center_index, chunk)| { - for (din_offset, item) in chunk.iter_mut().enumerate() { - *item = full_pivot_data - [center_index * dim + chunk_offsets[chunk_index] + din_offset]; - } - }); - - // Compute the closet centers - compute_closest_centers( - &cur_data, - cur_block_size, - cur_chunk_size, - &cur_pivot_data, - num_centers, - 1, - &mut closest_center, - None, - None, - )?; - - block_compressed_base - .par_chunks_mut(num_pq_chunks) - .enumerate() - .for_each(|(block_data_index, slice)| { - slice[chunk_index] = closest_center[block_data_index] as usize; - }); - } - - _ = pq_storage.write_compressed_pivot_data( - &block_compressed_base, - num_centers, - cur_block_size, - num_pq_chunks, - ); - } - Ok(()) -} - -/// Save the data on a file. -/// # Arguments -/// * `p_val` - choose how many ratio sample data as trained data to get pivot -/// * `num_pq_chunks` - pq chunk number -/// * `codebook_prefix` - predefined pivots file named -/// * `pq_storage` - pq file access -pub fn generate_quantized_data>( - p_val: f64, - num_pq_chunks: usize, - codebook_prefix: &str, - pq_storage: &mut PQStorage, -) -> ANNResult<()> { - // If predefined pivots already exists, skip training. - if !file_exists(codebook_prefix) { - // Instantiates train data with random sample updates train_data_vector - // Training data with train_size samples loaded. - // Each sampled file has train_dim. - let (mut train_data_vector, train_size, train_dim) = - pq_storage.gen_random_slice::(p_val)?; - - generate_pq_pivots( - &mut train_data_vector, - train_size, - train_dim, - NUM_PQ_CENTROIDS, - num_pq_chunks, - NUM_KMEANS_REPS_PQ, - pq_storage, - )?; - } - generate_pq_data_from_pivots::(NUM_PQ_CENTROIDS, num_pq_chunks, pq_storage)?; - Ok(()) -} - -#[cfg(test)] -mod pq_test { - - use std::fs::File; - use std::io::Write; - - use super::*; - use crate::utils::{convert_types_u32_usize, convert_types_u64_usize, load_bin, METADATA_SIZE}; - - #[test] - fn generate_pq_pivots_test() { - let pivot_file_name = "generate_pq_pivots_test.bin"; - let compressed_file_name = "compressed.bin"; - let pq_training_file_name = "tests/data/siftsmall_learn.bin"; - let mut pq_storage = - PQStorage::new(pivot_file_name, compressed_file_name, pq_training_file_name).unwrap(); - let mut train_data: Vec = vec![ - 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 2.0f32, 2.0f32, 2.0f32, - 2.0f32, 2.0f32, 2.0f32, 2.0f32, 2.0f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, - 2.1f32, 2.1f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, - 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, - ]; - generate_pq_pivots(&mut train_data, 5, 8, 2, 2, 5, &mut pq_storage).unwrap(); - - let (data, nr, nc) = load_bin::(pivot_file_name, 0).unwrap(); - let file_offset_data = convert_types_u64_usize(&data, nr, nc); - assert_eq!(file_offset_data[0], METADATA_SIZE); - assert_eq!(nr, 4); - assert_eq!(nc, 1); - - let (data, nr, nc) = load_bin::(pivot_file_name, file_offset_data[0]).unwrap(); - let full_pivot_data = data.to_vec(); - assert_eq!(full_pivot_data.len(), 16); - assert_eq!(nr, 2); - assert_eq!(nc, 8); - - let (data, nr, nc) = load_bin::(pivot_file_name, file_offset_data[1]).unwrap(); - let centroid = data.to_vec(); - assert_eq!( - centroid[0], - (1.0f32 + 2.0f32 + 2.1f32 + 2.2f32 + 100.0f32) / 5.0f32 - ); - assert_eq!(nr, 8); - assert_eq!(nc, 1); - - let (data, nr, nc) = load_bin::(pivot_file_name, file_offset_data[2]).unwrap(); - let chunk_offsets = convert_types_u32_usize(&data, nr, nc); - assert_eq!(chunk_offsets[0], 0); - assert_eq!(chunk_offsets[1], 4); - assert_eq!(chunk_offsets[2], 8); - assert_eq!(nr, 3); - assert_eq!(nc, 1); - std::fs::remove_file(pivot_file_name).unwrap(); - } - - #[test] - fn generate_pq_data_from_pivots_test() { - let data_file = "generate_pq_data_from_pivots_test_data.bin"; - //npoints=5, dim=8, 5 vectors [1.0;8] [2.0;8] [2.1;8] [2.2;8] [100.0;8] - let mut train_data: Vec = vec![ - 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 1.0f32, 2.0f32, 2.0f32, 2.0f32, - 2.0f32, 2.0f32, 2.0f32, 2.0f32, 2.0f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, 2.1f32, - 2.1f32, 2.1f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, 2.2f32, - 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, 100.0f32, - ]; - let my_nums_unstructured: &[u8] = unsafe { - std::slice::from_raw_parts(train_data.as_ptr() as *const u8, train_data.len() * 4) - }; - let meta: Vec = vec![5, 8]; - let meta_unstructured: &[u8] = - unsafe { std::slice::from_raw_parts(meta.as_ptr() as *const u8, meta.len() * 4) }; - let mut data_file_writer = File::create(data_file).unwrap(); - data_file_writer - .write_all(meta_unstructured) - .expect("Failed to write sample file"); - data_file_writer - .write_all(my_nums_unstructured) - .expect("Failed to write sample file"); - - let pq_pivots_path = "generate_pq_data_from_pivots_test_pivot.bin"; - let pq_compressed_vectors_path = "generate_pq_data_from_pivots_test.bin"; - let mut pq_storage = - PQStorage::new(pq_pivots_path, pq_compressed_vectors_path, data_file).unwrap(); - generate_pq_pivots(&mut train_data, 5, 8, 2, 2, 5, &mut pq_storage).unwrap(); - generate_pq_data_from_pivots::(2, 2, &mut pq_storage).unwrap(); - let (data, nr, nc) = load_bin::(pq_compressed_vectors_path, 0).unwrap(); - assert_eq!(nr, 5); - assert_eq!(nc, 2); - assert_eq!(data[0], data[2]); - assert_ne!(data[0], data[8]); - - std::fs::remove_file(data_file).unwrap(); - std::fs::remove_file(pq_pivots_path).unwrap(); - std::fs::remove_file(pq_compressed_vectors_path).unwrap(); - } - - #[test] - fn pq_end_to_end_validation_with_codebook_test() { - let data_file = "tests/data/siftsmall_learn.bin"; - let pq_pivots_path = "tests/data/siftsmall_learn.bin_pq_pivots.bin"; - let gound_truth_path = "tests/data/siftsmall_learn.bin_pq_compressed.bin"; - let pq_compressed_vectors_path = "validation.bin"; - let mut pq_storage = - PQStorage::new(pq_pivots_path, pq_compressed_vectors_path, data_file).unwrap(); - generate_quantized_data::(0.5, 1, pq_pivots_path, &mut pq_storage).unwrap(); - - let (data, nr, nc) = load_bin::(pq_compressed_vectors_path, 0).unwrap(); - let (gt_data, gt_nr, gt_nc) = load_bin::(gound_truth_path, 0).unwrap(); - assert_eq!(nr, gt_nr); - assert_eq!(nc, gt_nc); - for i in 0..data.len() { - assert_eq!(data[i], gt_data[i]); - } - std::fs::remove_file(pq_compressed_vectors_path).unwrap(); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/concurrent_queue.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/concurrent_queue.rs deleted file mode 100644 index 8c72bab..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/concurrent_queue.rs +++ /dev/null @@ -1,312 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -use std::collections::VecDeque; -use std::ops::Deref; -use std::sync::{Arc, Condvar, Mutex, MutexGuard}; -use std::time::Duration; - -use crate::common::{ANNError, ANNResult}; - -#[derive(Debug)] -/// Query scratch data structures -pub struct ConcurrentQueue { - q: Mutex>, - c: Mutex, - push_cv: Condvar, -} - -impl Default for ConcurrentQueue { - fn default() -> Self { - Self::new() - } -} - -impl ConcurrentQueue { - /// Create a concurrent queue - pub fn new() -> Self { - Self { - q: Mutex::new(VecDeque::new()), - c: Mutex::new(false), - push_cv: Condvar::new(), - } - } - - /// Block the current thread until it is able to acquire the mutex - pub fn reserve(&self, size: usize) -> ANNResult<()> { - let mut guard = lock(&self.q)?; - guard.reserve(size); - Ok(()) - } - - /// queue stats - pub fn size(&self) -> ANNResult { - let guard = lock(&self.q)?; - - Ok(guard.len()) - } - - /// empty the queue - pub fn is_empty(&self) -> ANNResult { - Ok(self.size()? == 0) - } - - /// push back - pub fn push(&self, new_val: T) -> ANNResult<()> { - let mut guard = lock(&self.q)?; - self.push_internal(&mut guard, new_val); - self.push_cv.notify_all(); - Ok(()) - } - - /// push back - fn push_internal(&self, guard: &mut MutexGuard>, new_val: T) { - guard.push_back(new_val); - } - - /// insert into queue - pub fn insert(&self, iter: I) -> ANNResult<()> - where - I: IntoIterator, - { - let mut guard = lock(&self.q)?; - for item in iter { - self.push_internal(&mut guard, item); - } - - self.push_cv.notify_all(); - Ok(()) - } - - /// pop front - pub fn pop(&self) -> ANNResult> { - let mut guard = lock(&self.q)?; - Ok(guard.pop_front()) - } - - /// Empty - is this necessary? - pub fn empty_queue(&self) -> ANNResult<()> { - let mut guard = lock(&self.q)?; - while !guard.is_empty() { - let _ = guard.pop_front(); - } - Ok(()) - } - - /// register for push notifications - pub fn wait_for_push_notify(&self, wait_time: Duration) -> ANNResult<()> { - let guard_lock = lock(&self.c)?; - let _ = self - .push_cv - .wait_timeout(guard_lock, wait_time) - .map_err(|err| { - ANNError::log_lock_poison_error(format!( - "ConcurrentQueue Lock is poisoned, err={}", - err - )) - })?; - Ok(()) - } -} - -fn lock(mutex: &Mutex) -> ANNResult> { - let guard = mutex.lock().map_err(|err| { - ANNError::log_lock_poison_error(format!("ConcurrentQueue lock is poisoned, err={}", err)) - })?; - Ok(guard) -} - -/// A thread-safe queue that holds instances of `T`. -/// Each instance is stored in a `Box` to keep the size of the queue node constant. -#[derive(Debug)] -pub struct ArcConcurrentBoxedQueue { - internal_queue: Arc>>, -} - -impl ArcConcurrentBoxedQueue { - /// Create a new `ArcConcurrentBoxedQueue`. - pub fn new() -> Self { - Self { - internal_queue: Arc::new(ConcurrentQueue::new()), - } - } -} - -impl Default for ArcConcurrentBoxedQueue { - fn default() -> Self { - Self::new() - } -} - -impl Clone for ArcConcurrentBoxedQueue { - /// Create a new `ArcConcurrentBoxedQueue` that shares the same internal queue - /// with the existing one. This allows multiple `ArcConcurrentBoxedQueue` to - /// operate on the same underlying queue. - fn clone(&self) -> Self { - Self { - internal_queue: Arc::clone(&self.internal_queue), - } - } -} - -/// Deref to the ConcurrentQueue. -impl Deref for ArcConcurrentBoxedQueue { - type Target = ConcurrentQueue>; - - fn deref(&self) -> &Self::Target { - &self.internal_queue - } -} - -#[cfg(test)] -mod tests { - use crate::model::ConcurrentQueue; - use std::sync::Arc; - use std::thread; - use std::time::Duration; - - #[test] - fn test_push_pop() { - let queue = ConcurrentQueue::::new(); - - queue.push(1).unwrap(); - queue.push(2).unwrap(); - queue.push(3).unwrap(); - - assert_eq!(queue.pop().unwrap(), Some(1)); - assert_eq!(queue.pop().unwrap(), Some(2)); - assert_eq!(queue.pop().unwrap(), Some(3)); - assert_eq!(queue.pop().unwrap(), None); - } - - #[test] - fn test_size_empty() { - let queue = ConcurrentQueue::new(); - - assert_eq!(queue.size().unwrap(), 0); - assert!(queue.is_empty().unwrap()); - - queue.push(1).unwrap(); - queue.push(2).unwrap(); - - assert_eq!(queue.size().unwrap(), 2); - assert!(!queue.is_empty().unwrap()); - - queue.pop().unwrap(); - queue.pop().unwrap(); - - assert_eq!(queue.size().unwrap(), 0); - assert!(queue.is_empty().unwrap()); - } - - #[test] - fn test_insert() { - let queue = ConcurrentQueue::new(); - - let data = vec![1, 2, 3]; - queue.insert(data.into_iter()).unwrap(); - - assert_eq!(queue.pop().unwrap(), Some(1)); - assert_eq!(queue.pop().unwrap(), Some(2)); - assert_eq!(queue.pop().unwrap(), Some(3)); - assert_eq!(queue.pop().unwrap(), None); - } - - #[test] - fn test_notifications() { - let queue = Arc::new(ConcurrentQueue::new()); - let queue_clone = Arc::clone(&queue); - - let producer = thread::spawn(move || { - for i in 0..3 { - thread::sleep(Duration::from_millis(50)); - queue_clone.push(i).unwrap(); - } - }); - - let consumer = thread::spawn(move || { - let mut values = vec![]; - - for _ in 0..3 { - let mut val = -1; - while val == -1 { - queue - .wait_for_push_notify(Duration::from_millis(10)) - .unwrap(); - val = queue.pop().unwrap().unwrap_or(-1); - } - - values.push(val); - } - - values - }); - - producer.join().unwrap(); - let consumer_results = consumer.join().unwrap(); - - assert_eq!(consumer_results, vec![0, 1, 2]); - } - - #[test] - fn test_multithreaded_push_pop() { - let queue = Arc::new(ConcurrentQueue::new()); - let queue_clone = Arc::clone(&queue); - - let producer = thread::spawn(move || { - for i in 0..10 { - queue_clone.push(i).unwrap(); - thread::sleep(Duration::from_millis(50)); - } - }); - - let consumer = thread::spawn(move || { - let mut values = vec![]; - - for _ in 0..10 { - let mut val = -1; - while val == -1 { - val = queue.pop().unwrap().unwrap_or(-1); - thread::sleep(Duration::from_millis(10)); - } - - values.push(val); - } - - values - }); - - producer.join().unwrap(); - let consumer_results = consumer.join().unwrap(); - - assert_eq!(consumer_results, (0..10).collect::>()); - } - - /// This is a single value test. It avoids the unlimited wait until the collectin got empty on the previous test. - /// It will make sure the signal mutex is matching the waiting mutex. - #[test] - fn test_wait_for_push_notify() { - let queue = Arc::new(ConcurrentQueue::::new()); - let queue_clone = Arc::clone(&queue); - - let producer = thread::spawn(move || { - thread::sleep(Duration::from_millis(100)); - queue_clone.push(1).unwrap(); - }); - - let consumer = thread::spawn(move || { - queue - .wait_for_push_notify(Duration::from_millis(200)) - .unwrap(); - assert_eq!(queue.pop().unwrap(), Some(1)); - }); - - producer.join().unwrap(); - consumer.join().unwrap(); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/inmem_query_scratch.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/inmem_query_scratch.rs deleted file mode 100644 index f0fa432..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/inmem_query_scratch.rs +++ /dev/null @@ -1,186 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Scratch space for in-memory index based search - -use std::cmp::max; -use std::mem; - -use hashbrown::HashSet; - -use crate::common::{ANNError, ANNResult, AlignedBoxWithSlice}; -use crate::model::configuration::index_write_parameters::IndexWriteParameters; -use crate::model::{Neighbor, NeighborPriorityQueue, PQScratch}; - -use super::Scratch; - -/// In-mem index related limits -pub const GRAPH_SLACK_FACTOR: f64 = 1.3_f64; - -/// Max number of points for using bitset -pub const MAX_POINTS_FOR_USING_BITSET: usize = 100000; - -/// TODO: SSD Index related limits -pub const MAX_GRAPH_DEGREE: usize = 512; - -/// TODO: SSD Index related limits -pub const MAX_N_CMPS: usize = 16384; - -/// TODO: SSD Index related limits -pub const SECTOR_LEN: usize = 4096; - -/// TODO: SSD Index related limits -pub const MAX_N_SECTOR_READS: usize = 128; - -/// The alignment required for memory access. This will be multiplied with size of T to get the actual alignment -pub const QUERY_ALIGNMENT_OF_T_SIZE: usize = 16; - -/// Scratch space for in-memory index based search -#[derive(Debug)] -pub struct InMemQueryScratch { - /// Size of the candidate queue - pub candidate_size: u32, - - /// Max degree for each vertex - pub max_degree: u32, - - /// Max occlusion size - pub max_occlusion_size: u32, - - /// Query node - pub query: AlignedBoxWithSlice, - - /// Best candidates, whose size is candidate_queue_size - pub best_candidates: NeighborPriorityQueue, - - /// Occlude factor - pub occlude_factor: Vec, - - /// Visited neighbor id - pub id_scratch: Vec, - - /// The distance between visited neighbor and query node - pub dist_scratch: Vec, - - /// The PQ Scratch, keey it private since this class use the Box to own the memory. Use the function pq_scratch to get its reference - pub pq_scratch: Option>, - - /// Buffers used in process delete, capacity increases as needed - pub expanded_nodes_set: HashSet, - - /// Expanded neighbors - pub expanded_neighbors_vector: Vec, - - /// Occlude list - pub occlude_list_output: Vec, - - /// RobinSet for larger dataset - pub node_visited_robinset: HashSet, -} - -impl InMemQueryScratch { - /// Create InMemQueryScratch instance - pub fn new( - search_candidate_size: u32, - index_write_parameter: &IndexWriteParameters, - init_pq_scratch: bool, - ) -> ANNResult { - let indexing_candidate_size = index_write_parameter.search_list_size; - let max_degree = index_write_parameter.max_degree; - let max_occlusion_size = index_write_parameter.max_occlusion_size; - - if search_candidate_size == 0 || indexing_candidate_size == 0 || max_degree == 0 || N == 0 { - return Err(ANNError::log_index_error(format!( - "In InMemQueryScratch, one of search_candidate_size = {}, indexing_candidate_size = {}, dim = {} or max_degree = {} is zero.", - search_candidate_size, indexing_candidate_size, N, max_degree))); - } - - let query = AlignedBoxWithSlice::new(N, mem::size_of::() * QUERY_ALIGNMENT_OF_T_SIZE)?; - let pq_scratch = if init_pq_scratch { - Some(Box::new(PQScratch::new(MAX_GRAPH_DEGREE, N)?)) - } else { - None - }; - - let occlude_factor = Vec::with_capacity(max_occlusion_size as usize); - - let capacity = (1.5 * GRAPH_SLACK_FACTOR * (max_degree as f64)).ceil() as usize; - let id_scratch = Vec::with_capacity(capacity); - let dist_scratch = Vec::with_capacity(capacity); - - let expanded_nodes_set = HashSet::::new(); - let expanded_neighbors_vector = Vec::::new(); - let occlude_list_output = Vec::::new(); - - let candidate_size = max(search_candidate_size, indexing_candidate_size); - let node_visited_robinset = HashSet::::with_capacity(20 * candidate_size as usize); - let scratch = Self { - candidate_size, - max_degree, - max_occlusion_size, - query, - best_candidates: NeighborPriorityQueue::with_capacity(candidate_size as usize), - occlude_factor, - id_scratch, - dist_scratch, - pq_scratch, - expanded_nodes_set, - expanded_neighbors_vector, - occlude_list_output, - node_visited_robinset, - }; - - Ok(scratch) - } - - /// Resize the scratch with new candidate size - pub fn resize_for_new_candidate_size(&mut self, new_candidate_size: u32) { - if new_candidate_size > self.candidate_size { - let delta = new_candidate_size - self.candidate_size; - self.candidate_size = new_candidate_size; - self.best_candidates.reserve(delta as usize); - self.node_visited_robinset.reserve((20 * delta) as usize); - } - } -} - -impl Scratch for InMemQueryScratch { - fn clear(&mut self) { - self.best_candidates.clear(); - self.occlude_factor.clear(); - - self.node_visited_robinset.clear(); - - self.id_scratch.clear(); - self.dist_scratch.clear(); - - self.expanded_nodes_set.clear(); - self.expanded_neighbors_vector.clear(); - self.occlude_list_output.clear(); - } -} - -#[cfg(test)] -mod inmemory_query_scratch_test { - use crate::model::configuration::index_write_parameters::IndexWriteParametersBuilder; - - use super::*; - - #[test] - fn node_visited_robinset_test() { - let index_write_parameter = IndexWriteParametersBuilder::new(10, 10) - .with_max_occlusion_size(5) - .build(); - - let mut scratch = - InMemQueryScratch::::new(100, &index_write_parameter, false).unwrap(); - - assert_eq!(scratch.node_visited_robinset.len(), 0); - - scratch.clear(); - assert_eq!(scratch.node_visited_robinset.len(), 0); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/mod.rs deleted file mode 100644 index cf9ee29..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/mod.rs +++ /dev/null @@ -1,28 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod scratch_traits; -pub use scratch_traits::*; - -pub mod concurrent_queue; -pub use concurrent_queue::*; - -pub mod pq_scratch; -pub use pq_scratch::*; - - -pub mod inmem_query_scratch; -pub use inmem_query_scratch::*; - -pub mod scratch_store_manager; -pub use scratch_store_manager::*; - -pub mod ssd_query_scratch; -pub use ssd_query_scratch::*; - -pub mod ssd_thread_data; -pub use ssd_thread_data::*; - -pub mod ssd_io_context; -pub use ssd_io_context::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/pq_scratch.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/pq_scratch.rs deleted file mode 100644 index bf9d6c5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/pq_scratch.rs +++ /dev/null @@ -1,105 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -use std::mem::size_of; - -use crate::common::{ANNResult, AlignedBoxWithSlice}; - -const MAX_PQ_CHUNKS: usize = 512; - -#[derive(Debug)] -/// PQ scratch -pub struct PQScratch { - /// Aligned pq table dist scratch, must be at least [256 * NCHUNKS] - pub aligned_pqtable_dist_scratch: AlignedBoxWithSlice, - /// Aligned dist scratch, must be at least diskann MAX_DEGREE - pub aligned_dist_scratch: AlignedBoxWithSlice, - /// Aligned pq coord scratch, must be at least [N_CHUNKS * MAX_DEGREE] - pub aligned_pq_coord_scratch: AlignedBoxWithSlice, - /// Rotated query - pub rotated_query: AlignedBoxWithSlice, - /// Aligned query float - pub aligned_query_float: AlignedBoxWithSlice, -} - -impl PQScratch { - const ALIGNED_ALLOC_256: usize = 256; - - /// Create a new pq scratch - pub fn new(graph_degree: usize, aligned_dim: usize) -> ANNResult { - let aligned_pq_coord_scratch = - AlignedBoxWithSlice::new(graph_degree * MAX_PQ_CHUNKS, PQScratch::ALIGNED_ALLOC_256)?; - let aligned_pqtable_dist_scratch = - AlignedBoxWithSlice::new(256 * MAX_PQ_CHUNKS, PQScratch::ALIGNED_ALLOC_256)?; - let aligned_dist_scratch = - AlignedBoxWithSlice::new(graph_degree, PQScratch::ALIGNED_ALLOC_256)?; - let aligned_query_float = AlignedBoxWithSlice::new(aligned_dim, 8 * size_of::())?; - let rotated_query = AlignedBoxWithSlice::new(aligned_dim, 8 * size_of::())?; - - Ok(Self { - aligned_pqtable_dist_scratch, - aligned_dist_scratch, - aligned_pq_coord_scratch, - rotated_query, - aligned_query_float, - }) - } - - /// Set rotated_query and aligned_query_float values - pub fn set(&mut self, dim: usize, query: &[T], norm: f32) - where - T: Into + Copy, - { - for (d, item) in query.iter().enumerate().take(dim) { - let query_val: f32 = (*item).into(); - if (norm - 1.0).abs() > f32::EPSILON { - self.rotated_query[d] = query_val / norm; - self.aligned_query_float[d] = query_val / norm; - } else { - self.rotated_query[d] = query_val; - self.aligned_query_float[d] = query_val; - } - } - } -} - -#[cfg(test)] -mod tests { - use crate::model::PQScratch; - - #[test] - fn test_pq_scratch() { - let graph_degree = 512; - let aligned_dim = 8; - - let mut pq_scratch: PQScratch = PQScratch::new(graph_degree, aligned_dim).unwrap(); - - // Check alignment - assert_eq!( - (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) % 256, - 0 - ); - assert_eq!((pq_scratch.aligned_dist_scratch.as_ptr() as usize) % 256, 0); - assert_eq!( - (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % 256, - 0 - ); - assert_eq!((pq_scratch.rotated_query.as_ptr() as usize) % 32, 0); - assert_eq!((pq_scratch.aligned_query_float.as_ptr() as usize) % 32, 0); - - // Test set() method - let query = vec![1u8, 2, 3, 4, 5, 6, 7, 8]; - let norm = 2.0f32; - pq_scratch.set::(query.len(), &query, norm); - - (0..query.len()).for_each(|i| { - assert_eq!(pq_scratch.rotated_query[i], query[i] as f32 / norm); - assert_eq!(pq_scratch.aligned_query_float[i], query[i] as f32 / norm); - }); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_store_manager.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_store_manager.rs deleted file mode 100644 index 4e2397f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_store_manager.rs +++ /dev/null @@ -1,84 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::common::ANNResult; - -use super::ArcConcurrentBoxedQueue; -use super::{scratch_traits::Scratch}; -use std::time::Duration; - -pub struct ScratchStoreManager { - scratch: Option>, - scratch_pool: ArcConcurrentBoxedQueue, -} - -impl ScratchStoreManager { - pub fn new(scratch_pool: ArcConcurrentBoxedQueue, wait_time: Duration) -> ANNResult { - let mut scratch = scratch_pool.pop()?; - while scratch.is_none() { - scratch_pool.wait_for_push_notify(wait_time)?; - scratch = scratch_pool.pop()?; - } - - Ok(ScratchStoreManager { - scratch, - scratch_pool, - }) - } - - pub fn scratch_space(&mut self) -> Option<&mut T> { - self.scratch.as_deref_mut() - } -} - -impl Drop for ScratchStoreManager { - fn drop(&mut self) { - if let Some(mut scratch) = self.scratch.take() { - scratch.clear(); - let _ = self.scratch_pool.push(scratch); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[derive(Debug)] - struct MyScratch { - data: Vec, - } - - impl Scratch for MyScratch { - fn clear(&mut self) { - self.data.clear(); - } - } - - #[test] - fn test_scratch_store_manager() { - let wait_time = Duration::from_millis(100); - - let scratch_pool = ArcConcurrentBoxedQueue::new(); - for i in 1..3 { - scratch_pool.push(Box::new(MyScratch { - data: vec![i, 2 * i, 3 * i], - })).unwrap(); - } - - let mut manager = ScratchStoreManager::new(scratch_pool.clone(), wait_time).unwrap(); - let scratch_space = manager.scratch_space().unwrap(); - - assert_eq!(scratch_space.data, vec![1, 2, 3]); - - // At this point, the ScratchStoreManager will go out of scope, - // causing the Drop implementation to be called, which should - // call the clear method on MyScratch. - drop(manager); - - let current_scratch = scratch_pool.pop().unwrap().unwrap(); - assert_eq!(current_scratch.data, vec![2, 4, 6]); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_traits.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_traits.rs deleted file mode 100644 index 71e4b93..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/scratch_traits.rs +++ /dev/null @@ -1,8 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub trait Scratch { - fn clear(&mut self); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_io_context.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_io_context.rs deleted file mode 100644 index d4dff0c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_io_context.rs +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![allow(dead_code)] // Todo: Remove this when the disk index query code is complete. -use crate::common::ANNError; - -use platform::{FileHandle, IOCompletionPort}; - -// The IOContext struct for disk I/O. One for each thread. -pub struct IOContext { - pub status: Status, - pub file_handle: FileHandle, - pub io_completion_port: IOCompletionPort, -} - -impl Default for IOContext { - fn default() -> Self { - IOContext { - status: Status::ReadWait, - file_handle: FileHandle::default(), - io_completion_port: IOCompletionPort::default(), - } - } -} - -impl IOContext { - pub fn new() -> Self { - Self::default() - } -} - -pub enum Status { - ReadWait, - ReadSuccess, - ReadFailed(ANNError), - ProcessComplete, -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_query_scratch.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_query_scratch.rs deleted file mode 100644 index b366693..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_query_scratch.rs +++ /dev/null @@ -1,132 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![allow(dead_code)] // Todo: Remove this when the disk index query code is complete. -use std::mem; -use std::vec::Vec; - -use hashbrown::HashSet; - -use crate::{ - common::{ANNResult, AlignedBoxWithSlice}, - model::{Neighbor, NeighborPriorityQueue}, - model::data_store::DiskScratchDataset, -}; - -use super::{PQScratch, Scratch, MAX_GRAPH_DEGREE, QUERY_ALIGNMENT_OF_T_SIZE}; - -// Scratch space for disk index based search. -pub struct SSDQueryScratch -{ - // Disk scratch dataset storing fp vectors with aligned dim (N) - pub scratch_dataset: DiskScratchDataset, - - // The query scratch. - pub query: AlignedBoxWithSlice, - - /// The PQ Scratch. - pub pq_scratch: Option>, - - // The visited set. - pub id_scratch: HashSet, - - /// Best candidates, whose size is candidate_queue_size - pub best_candidates: NeighborPriorityQueue, - - // Full return set. - pub full_return_set: Vec, -} - -// -impl SSDQueryScratch -{ - pub fn new( - visited_reserve: usize, - candidate_queue_size: usize, - init_pq_scratch: bool, - ) -> ANNResult { - let scratch_dataset = DiskScratchDataset::::new()?; - - let query = AlignedBoxWithSlice::::new(N, mem::size_of::() * QUERY_ALIGNMENT_OF_T_SIZE)?; - - let id_scratch = HashSet::::with_capacity(visited_reserve); - let full_return_set = Vec::::with_capacity(visited_reserve); - let best_candidates = NeighborPriorityQueue::with_capacity(candidate_queue_size); - - let pq_scratch = if init_pq_scratch { - Some(Box::new(PQScratch::new(MAX_GRAPH_DEGREE, N)?)) - } else { - None - }; - - Ok(Self { - scratch_dataset, - query, - pq_scratch, - id_scratch, - best_candidates, - full_return_set, - }) - } - - pub fn pq_scratch(&mut self) -> &Option> { - &self.pq_scratch - } -} - -impl Scratch for SSDQueryScratch -{ - fn clear(&mut self) { - self.id_scratch.clear(); - self.best_candidates.clear(); - self.full_return_set.clear(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_new() { - // Arrange - let visited_reserve = 100; - let candidate_queue_size = 10; - let init_pq_scratch = true; - - // Act - let result = - SSDQueryScratch::::new(visited_reserve, candidate_queue_size, init_pq_scratch); - - // Assert - assert!(result.is_ok()); - - let scratch = result.unwrap(); - - // Assert the properties of the scratch instance - assert!(scratch.pq_scratch.is_some()); - assert!(scratch.id_scratch.is_empty()); - assert!(scratch.best_candidates.size() == 0); - assert!(scratch.full_return_set.is_empty()); - } - - #[test] - fn test_clear() { - // Arrange - let mut scratch = SSDQueryScratch::::new(100, 10, true).unwrap(); - - // Add some data to scratch fields - scratch.id_scratch.insert(1); - scratch.best_candidates.insert(Neighbor::new(2, 0.5)); - scratch.full_return_set.push(Neighbor::new(3, 0.8)); - - // Act - scratch.clear(); - - // Assert - assert!(scratch.id_scratch.is_empty()); - assert!(scratch.best_candidates.size() == 0); - assert!(scratch.full_return_set.is_empty()); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_thread_data.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_thread_data.rs deleted file mode 100644 index e374959..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/scratch/ssd_thread_data.rs +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![allow(dead_code)] // Todo: Remove this when the disk index query code is complete. -use std::sync::Arc; - -use super::{scratch_traits::Scratch, IOContext, SSDQueryScratch}; -use crate::common::ANNResult; - -// The thread data struct for SSD I/O. One for each thread, contains the ScratchSpace and the IOContext. -pub struct SSDThreadData { - pub scratch: SSDQueryScratch, - pub io_context: Option>, -} - -impl SSDThreadData { - pub fn new( - aligned_dim: usize, - visited_reserve: usize, - init_pq_scratch: bool, - ) -> ANNResult { - let scratch = SSDQueryScratch::new(aligned_dim, visited_reserve, init_pq_scratch)?; - Ok(SSDThreadData { - scratch, - io_context: None, - }) - } - - pub fn clear(&mut self) { - self.scratch.clear(); - } -} - -#[cfg(test)] -mod tests { - use crate::model::Neighbor; - - use super::*; - - #[test] - fn test_new() { - // Arrange - let aligned_dim = 10; - let visited_reserve = 100; - let init_pq_scratch = true; - - // Act - let result = SSDThreadData::::new(aligned_dim, visited_reserve, init_pq_scratch); - - // Assert - assert!(result.is_ok()); - - let thread_data = result.unwrap(); - - // Assert the properties of the thread data instance - assert!(thread_data.io_context.is_none()); - - let scratch = &thread_data.scratch; - // Assert the properties of the scratch instance - assert!(scratch.pq_scratch.is_some()); - assert!(scratch.id_scratch.is_empty()); - assert!(scratch.best_candidates.size() == 0); - assert!(scratch.full_return_set.is_empty()); - } - - #[test] - fn test_clear() { - // Arrange - let mut thread_data = SSDThreadData::::new(10, 100, true).unwrap(); - - // Add some data to scratch fields - thread_data.scratch.id_scratch.insert(1); - thread_data - .scratch - .best_candidates - .insert(Neighbor::new(2, 0.5)); - thread_data - .scratch - .full_return_set - .push(Neighbor::new(3, 0.8)); - - // Act - thread_data.clear(); - - // Assert - assert!(thread_data.scratch.id_scratch.is_empty()); - assert!(thread_data.scratch.best_candidates.size() == 0); - assert!(thread_data.scratch.full_return_set.is_empty()); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/dimension.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/dimension.rs deleted file mode 100644 index 32670a8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/dimension.rs +++ /dev/null @@ -1,22 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Vertex dimension - -/// 32 vertex dimension -pub const DIM_32: usize = 32; - -/// 64 vertex dimension -pub const DIM_64: usize = 64; - -/// 104 vertex dimension -pub const DIM_104: usize = 104; - -/// 128 vertex dimension -pub const DIM_128: usize = 128; - -/// 256 vertex dimension -pub const DIM_256: usize = 256; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/mod.rs deleted file mode 100644 index 224d476..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod vertex; -pub use vertex::Vertex; - -mod dimension; -pub use dimension::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/vertex.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/vertex.rs deleted file mode 100644 index 5536974..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/vertex/vertex.rs +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Vertex - -use std::array::TryFromSliceError; - -use vector::{FullPrecisionDistance, Metric}; - -/// Vertex with data type T and dimension N -#[derive(Debug)] -pub struct Vertex<'a, T, const N: usize> -where - [T; N]: FullPrecisionDistance, -{ - /// Vertex value - val: &'a [T; N], - - /// Vertex Id - id: u32, -} - -impl<'a, T, const N: usize> Vertex<'a, T, N> -where - [T; N]: FullPrecisionDistance, -{ - /// Create the vertex with data - pub fn new(val: &'a [T; N], id: u32) -> Self { - Self { - val, - id, - } - } - - /// Compare the vertex with another. - #[inline(always)] - pub fn compare(&self, other: &Vertex<'a, T, N>, metric: Metric) -> f32 { - <[T; N]>::distance_compare(self.val, other.val, metric) - } - - /// Get the vector associated with the vertex. - #[inline] - pub fn vector(&self) -> &[T; N] { - self.val - } - - /// Get the vertex id. - #[inline] - pub fn vertex_id(&self) -> u32 { - self.id - } -} - -impl<'a, T, const N: usize> TryFrom<(&'a [T], u32)> for Vertex<'a, T, N> -where - [T; N]: FullPrecisionDistance, -{ - type Error = TryFromSliceError; - - fn try_from((mem_slice, id): (&'a [T], u32)) -> Result { - let array: &[T; N] = mem_slice.try_into()?; - Ok(Vertex::new(array, id)) - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/mod.rs deleted file mode 100644 index 0e63df0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[allow(clippy::module_inception)] -mod windows_aligned_file_reader; -pub use windows_aligned_file_reader::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/windows_aligned_file_reader.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/windows_aligned_file_reader.rs deleted file mode 100644 index 1cc3dc0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/model/windows_aligned_file_reader/windows_aligned_file_reader.rs +++ /dev/null @@ -1,414 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::sync::Arc; -use std::time::Duration; -use std::{ptr, thread}; - -use crossbeam::sync::ShardedLock; -use hashbrown::HashMap; -use once_cell::sync::Lazy; - -use platform::file_handle::{AccessMode, ShareMode}; -use platform::{ - file_handle::FileHandle, - file_io::{get_queued_completion_status, read_file_to_slice}, - io_completion_port::IOCompletionPort, -}; - -use winapi::{ - shared::{basetsd::ULONG_PTR, minwindef::DWORD}, - um::minwinbase::OVERLAPPED, -}; - -use crate::common::{ANNError, ANNResult}; -use crate::model::IOContext; - -pub const MAX_IO_CONCURRENCY: usize = 128; // To do: explore the optimal value for this. The current value is taken from C++ code. -pub const FILE_ATTRIBUTE_READONLY: DWORD = 0x00000001; -pub const IO_COMPLETION_TIMEOUT: DWORD = u32::MAX; // Infinite timeout. -pub const DISK_IO_ALIGNMENT: usize = 512; -pub const ASYNC_IO_COMPLETION_CHECK_INTERVAL: Duration = Duration::from_micros(5); - -/// Aligned read struct for disk IO, it takes the ownership of the AlignedBoxedSlice and returns the AlignedBoxWithSlice data immutably. -pub struct AlignedRead<'a, T> { - /// where to read from - /// offset needs to be aligned with DISK_IO_ALIGNMENT - offset: u64, - - /// where to read into - /// aligned_buf and its len need to be aligned with DISK_IO_ALIGNMENT - aligned_buf: &'a mut [T], -} - -impl<'a, T> AlignedRead<'a, T> { - pub fn new(offset: u64, aligned_buf: &'a mut [T]) -> ANNResult { - Self::assert_is_aligned(offset as usize)?; - Self::assert_is_aligned(std::mem::size_of_val(aligned_buf))?; - - Ok(Self { - offset, - aligned_buf, - }) - } - - fn assert_is_aligned(val: usize) -> ANNResult<()> { - match val % DISK_IO_ALIGNMENT { - 0 => Ok(()), - _ => Err(ANNError::log_disk_io_request_alignment_error(format!( - "The offset or length of AlignedRead request is not {} bytes aligned", - DISK_IO_ALIGNMENT - ))), - } - } - - pub fn aligned_buf(&self) -> &[T] { - self.aligned_buf - } -} - -pub struct WindowsAlignedFileReader { - file_name: String, - - // ctx_map is the mapping from thread id to io context. It is hashmap behind a sharded lock to allow concurrent access from multiple threads. - // ShardedLock: shardedlock provides an implementation of a reader-writer lock that offers concurrent read access to the shared data while allowing exclusive write access. - // It achieves better scalability by dividing the shared data into multiple shards, and each with its own internal lock. - // Multiple threads can read from different shards simultaneously, reducing contention. - // https://docs.rs/crossbeam/0.8.2/crossbeam/sync/struct.ShardedLock.html - // Comparing to RwLock, ShardedLock provides higher concurrency for read operations and is suitable for read heavy workloads. - // The value of the hashmap is an Arc to allow immutable access to IOContext with automatic reference counting. - ctx_map: Lazy>>>, -} - -impl WindowsAlignedFileReader { - pub fn new(fname: &str) -> ANNResult { - let reader: WindowsAlignedFileReader = WindowsAlignedFileReader { - file_name: fname.to_string(), - ctx_map: Lazy::new(|| ShardedLock::new(HashMap::new())), - }; - - reader.register_thread()?; - Ok(reader) - } - - // Register the io context for a thread if it hasn't been registered. - pub fn register_thread(&self) -> ANNResult<()> { - let mut ctx_map = self.ctx_map.write().map_err(|_| { - ANNError::log_lock_poison_error("unable to acquire read lock on ctx_map".to_string()) - })?; - - let id = thread::current().id(); - if ctx_map.contains_key(&id) { - println!( - "Warning:: Duplicate registration for thread_id : {:?}. Directly call get_ctx to get the thread context data.", - id); - - return Ok(()); - } - - let mut ctx = IOContext::new(); - - match unsafe { FileHandle::new(&self.file_name, AccessMode::Read, ShareMode::Read) } { - Ok(file_handle) => ctx.file_handle = file_handle, - Err(err) => { - return Err(ANNError::log_io_error(err)); - } - } - - // Create a io completion port for the file handle, later it will be used to get the completion status. - match IOCompletionPort::new(&ctx.file_handle, None, 0, 0) { - Ok(io_completion_port) => ctx.io_completion_port = io_completion_port, - Err(err) => { - return Err(ANNError::log_io_error(err)); - } - } - - ctx_map.insert(id, Arc::new(ctx)); - - Ok(()) - } - - // Get the reference counted io context for the current thread. - pub fn get_ctx(&self) -> ANNResult> { - let ctx_map = self.ctx_map.read().map_err(|_| { - ANNError::log_lock_poison_error("unable to acquire read lock on ctx_map".to_string()) - })?; - - let id = thread::current().id(); - match ctx_map.get(&id) { - Some(ctx) => Ok(Arc::clone(ctx)), - None => Err(ANNError::log_index_error(format!( - "unable to find IOContext for thread_id {:?}", - id - ))), - } - } - - // Read the data from the file by sending concurrent io requests in batches. - pub fn read(&self, read_requests: &mut [AlignedRead], ctx: &IOContext) -> ANNResult<()> { - let n_requests = read_requests.len(); - let n_batches = (n_requests + MAX_IO_CONCURRENCY - 1) / MAX_IO_CONCURRENCY; - - let mut overlapped_in_out = - vec![unsafe { std::mem::zeroed::() }; MAX_IO_CONCURRENCY]; - - for batch_idx in 0..n_batches { - let batch_start = MAX_IO_CONCURRENCY * batch_idx; - let batch_size = std::cmp::min(n_requests - batch_start, MAX_IO_CONCURRENCY); - - for j in 0..batch_size { - let req = &mut read_requests[batch_start + j]; - let os = &mut overlapped_in_out[j]; - - match unsafe { - read_file_to_slice(&ctx.file_handle, req.aligned_buf, os, req.offset) - } { - Ok(_) => {} - Err(error) => { - return Err(ANNError::IOError { err: (error) }); - } - } - } - - let mut n_read: DWORD = 0; - let mut n_complete: u64 = 0; - let mut completion_key: ULONG_PTR = 0; - let mut lp_os: *mut OVERLAPPED = ptr::null_mut(); - while n_complete < batch_size as u64 { - match unsafe { - get_queued_completion_status( - &ctx.io_completion_port, - &mut n_read, - &mut completion_key, - &mut lp_os, - IO_COMPLETION_TIMEOUT, - ) - } { - // An IO request completed. - Ok(true) => n_complete += 1, - // No IO request completed, continue to wait. - Ok(false) => { - thread::sleep(ASYNC_IO_COMPLETION_CHECK_INTERVAL); - } - // An error ocurred. - Err(error) => return Err(ANNError::IOError { err: (error) }), - } - } - } - - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use std::{fs::File, io::BufReader}; - - use bincode::deserialize_from; - use serde::{Deserialize, Serialize}; - - use crate::{common::AlignedBoxWithSlice, model::SECTOR_LEN}; - - use super::*; - pub const TEST_INDEX_PATH: &str = - "./tests/data/disk_index_siftsmall_learn_256pts_R4_L50_A1.2_alligned_reader_test.index"; - pub const TRUTH_NODE_DATA_PATH: &str = - "./tests/data/disk_index_node_data_aligned_reader_truth.bin"; - - #[derive(Debug, Serialize, Deserialize)] - struct NodeData { - num_neighbors: u32, - coordinates: Vec, - neighbors: Vec, - } - - impl PartialEq for NodeData { - fn eq(&self, other: &Self) -> bool { - self.num_neighbors == other.num_neighbors - && self.coordinates == other.coordinates - && self.neighbors == other.neighbors - } - } - - #[test] - fn test_new_aligned_file_reader() { - // Replace "test_file_path" with actual file path - let result = WindowsAlignedFileReader::new(TEST_INDEX_PATH); - assert!(result.is_ok()); - - let reader = result.unwrap(); - assert_eq!(reader.file_name, TEST_INDEX_PATH); - } - - #[test] - fn test_read() { - let reader = WindowsAlignedFileReader::new(TEST_INDEX_PATH).unwrap(); - let ctx = reader.get_ctx().unwrap(); - - let read_length = 512; // adjust according to your logic - let num_read = 10; - let mut aligned_mem = AlignedBoxWithSlice::::new(read_length * num_read, 512).unwrap(); - - // create and add AlignedReads to the vector - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); - - let mut aligned_reads: Vec> = mem_slices - .iter_mut() - .enumerate() - .map(|(i, slice)| { - let offset = (i * read_length) as u64; - AlignedRead::new(offset, slice).unwrap() - }) - .collect(); - - let result = reader.read(&mut aligned_reads, &ctx); - assert!(result.is_ok()); - } - - #[test] - fn test_read_disk_index_by_sector() { - let reader = WindowsAlignedFileReader::new(TEST_INDEX_PATH).unwrap(); - let ctx = reader.get_ctx().unwrap(); - - let read_length = SECTOR_LEN; // adjust according to your logic - let num_sector = 10; - let mut aligned_mem = - AlignedBoxWithSlice::::new(read_length * num_sector, 512).unwrap(); - - // Each slice will be used as the buffer for a read request of a sector. - let mut mem_slices = aligned_mem - .split_into_nonoverlapping_mut_slices(0..aligned_mem.len(), read_length) - .unwrap(); - - let mut aligned_reads: Vec> = mem_slices - .iter_mut() - .enumerate() - .map(|(sector_id, slice)| { - let offset = (sector_id * read_length) as u64; - AlignedRead::new(offset, slice).unwrap() - }) - .collect(); - - let result = reader.read(&mut aligned_reads, &ctx); - assert!(result.is_ok()); - - aligned_reads.iter().for_each(|read| { - assert_eq!(read.aligned_buf.len(), SECTOR_LEN); - }); - - let disk_layout_meta = reconstruct_disk_meta(aligned_reads[0].aligned_buf); - assert!(disk_layout_meta.len() > 9); - - let dims = disk_layout_meta[1]; - let num_pts = disk_layout_meta[0]; - let max_node_len = disk_layout_meta[3]; - let max_num_nodes_per_sector = disk_layout_meta[4]; - - assert!(max_node_len * max_num_nodes_per_sector < SECTOR_LEN as u64); - - let num_nbrs_start = (dims as usize) * std::mem::size_of::(); - let nbrs_buf_start = num_nbrs_start + std::mem::size_of::(); - - let mut node_data_array = Vec::with_capacity(max_num_nodes_per_sector as usize * 9); - - // Only validate the first 9 sectors with graph nodes. - (1..9).for_each(|sector_id| { - let sector_data = &mem_slices[sector_id]; - for node_data in sector_data.chunks_exact(max_node_len as usize) { - // Extract coordinates data from the start of the node_data - let coordinates_end = (dims as usize) * std::mem::size_of::(); - let coordinates = node_data[0..coordinates_end] - .chunks_exact(std::mem::size_of::()) - .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) - .collect(); - - // Extract number of neighbors from the node_data - let neighbors_num = u32::from_le_bytes( - node_data[num_nbrs_start..nbrs_buf_start] - .try_into() - .unwrap(), - ); - - let nbors_buf_end = - nbrs_buf_start + (neighbors_num as usize) * std::mem::size_of::(); - - // Extract neighbors from the node data. - let mut neighbors = Vec::new(); - for nbors_data in node_data[nbrs_buf_start..nbors_buf_end] - .chunks_exact(std::mem::size_of::()) - { - let nbors_id = u32::from_le_bytes(nbors_data.try_into().unwrap()); - assert!(nbors_id < num_pts as u32); - neighbors.push(nbors_id); - } - - // Create NodeData struct and push it to the node_data_array - node_data_array.push(NodeData { - num_neighbors: neighbors_num, - coordinates, - neighbors, - }); - } - }); - - // Compare that each node read from the disk index are expected. - let node_data_truth_file = File::open(TRUTH_NODE_DATA_PATH).unwrap(); - let reader = BufReader::new(node_data_truth_file); - - let node_data_vec: Vec = deserialize_from(reader).unwrap(); - for (node_from_node_data_file, node_from_disk_index) in - node_data_vec.iter().zip(node_data_array.iter()) - { - // Verify that the NodeData from the file is equal to the NodeData in node_data_array - assert_eq!(node_from_node_data_file, node_from_disk_index); - } - } - - #[test] - fn test_read_fail_invalid_file() { - let reader = WindowsAlignedFileReader::new("/invalid_path"); - assert!(reader.is_err()); - } - - #[test] - fn test_read_no_requests() { - let reader = WindowsAlignedFileReader::new(TEST_INDEX_PATH).unwrap(); - let ctx = reader.get_ctx().unwrap(); - - let mut read_requests = Vec::>::new(); - let result = reader.read(&mut read_requests, &ctx); - assert!(result.is_ok()); - } - - #[test] - fn test_get_ctx() { - let reader = WindowsAlignedFileReader::new(TEST_INDEX_PATH).unwrap(); - let result = reader.get_ctx(); - assert!(result.is_ok()); - } - - #[test] - fn test_register_thread() { - let reader = WindowsAlignedFileReader::new(TEST_INDEX_PATH).unwrap(); - let result = reader.register_thread(); - assert!(result.is_ok()); - } - - fn reconstruct_disk_meta(buffer: &[u8]) -> Vec { - let size_of_u64 = std::mem::size_of::(); - - let num_values = buffer.len() / size_of_u64; - let mut disk_layout_meta = Vec::with_capacity(num_values); - let meta_data = &buffer[8..]; - - for chunk in meta_data.chunks_exact(size_of_u64) { - let value = u64::from_le_bytes(chunk.try_into().unwrap()); - disk_layout_meta.push(value); - } - - disk_layout_meta - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_graph_storage.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_graph_storage.rs deleted file mode 100644 index 4481752..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_graph_storage.rs +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_docs)] - -//! Disk graph storage - -use std::sync::Arc; - -use crate::{model::{WindowsAlignedFileReader, IOContext, AlignedRead}, common::ANNResult}; - -/// Graph storage for disk index -/// One thread has one storage instance -pub struct DiskGraphStorage { - /// Disk graph reader - disk_graph_reader: Arc, - - /// IOContext of current thread - ctx: Arc, -} - -impl DiskGraphStorage { - /// Create a new DiskGraphStorage instance - pub fn new(disk_graph_reader: Arc) -> ANNResult { - let ctx = disk_graph_reader.get_ctx()?; - Ok(Self { - disk_graph_reader, - ctx, - }) - } - - /// Read disk graph data - pub fn read(&self, read_requests: &mut [AlignedRead]) -> ANNResult<()> { - self.disk_graph_reader.read(read_requests, &self.ctx) - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_index_storage.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_index_storage.rs deleted file mode 100644 index 0c55808..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/disk_index_storage.rs +++ /dev/null @@ -1,363 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use byteorder::{ByteOrder, LittleEndian, ReadBytesExt}; -use std::fs::File; -use std::io::Read; -use std::marker::PhantomData; -use std::{fs, mem}; - -use crate::common::{ANNError, ANNResult}; -use crate::model::NUM_PQ_CENTROIDS; -use crate::storage::PQStorage; -use crate::utils::{convert_types_u32_usize, convert_types_u64_usize, load_bin, save_bin_u64}; -use crate::utils::{ - file_exists, gen_sample_data, get_file_size, round_up, CachedReader, CachedWriter, -}; - -const SECTOR_LEN: usize = 4096; - -/// Todo: Remove the allow(dead_code) when the disk search code is complete -#[allow(dead_code)] -pub struct PQPivotData { - dim: usize, - pq_table: Vec, - centroids: Vec, - chunk_offsets: Vec, -} - -pub struct DiskIndexStorage { - /// Dataset file - dataset_file: String, - - /// Index file path prefix - index_path_prefix: String, - - // TODO: Only a placeholder for T, will be removed later - _marker: PhantomData, - - pq_storage: PQStorage, -} - -impl DiskIndexStorage { - /// Create DiskIndexStorage instance - pub fn new(dataset_file: String, index_path_prefix: String) -> ANNResult { - let pq_storage: PQStorage = PQStorage::new( - &(index_path_prefix.clone() + ".bin_pq_pivots.bin"), - &(index_path_prefix.clone() + ".bin_pq_compressed.bin"), - &dataset_file, - )?; - - Ok(DiskIndexStorage { - dataset_file, - index_path_prefix, - _marker: PhantomData, - pq_storage, - }) - } - - pub fn get_pq_storage(&mut self) -> &mut PQStorage { - &mut self.pq_storage - } - - pub fn dataset_file(&self) -> &String { - &self.dataset_file - } - - pub fn index_path_prefix(&self) -> &String { - &self.index_path_prefix - } - - /// Create disk layout - /// Sector #1: disk_layout_meta - /// Sector #n: num_nodes_per_sector nodes - /// Each node's layout: {full precision vector:[T; DIM]}{num_nbrs: u32}{neighbors: [u32; num_nbrs]} - /// # Arguments - /// * `dataset_file` - dataset file containing full precision vectors - /// * `mem_index_file` - in-memory index graph file - /// * `disk_layout_file` - output disk layout file - pub fn create_disk_layout(&self) -> ANNResult<()> { - let mem_index_file = self.mem_index_file(); - let disk_layout_file = self.disk_index_file(); - - // amount to read or write in one shot - let read_blk_size = 64 * 1024 * 1024; - let write_blk_size = read_blk_size; - let mut dataset_reader = CachedReader::new(self.dataset_file.as_str(), read_blk_size)?; - - let num_pts = dataset_reader.read_u32()? as u64; - let dims = dataset_reader.read_u32()? as u64; - - // Create cached reader + writer - let actual_file_size = get_file_size(mem_index_file.as_str())?; - println!("Vamana index file size={}", actual_file_size); - - let mut vamana_reader = File::open(mem_index_file)?; - let mut diskann_writer = CachedWriter::new(disk_layout_file.as_str(), write_blk_size)?; - - let index_file_size = vamana_reader.read_u64::()?; - if index_file_size != actual_file_size { - println!( - "Vamana Index file size does not match expected size per meta-data. file size from file: {}, actual file size: {}", - index_file_size, actual_file_size - ); - } - - let max_degree = vamana_reader.read_u32::()?; - let medoid = vamana_reader.read_u32::()?; - let vamana_frozen_num = vamana_reader.read_u64::()?; - - let mut vamana_frozen_loc = 0; - if vamana_frozen_num == 1 { - vamana_frozen_loc = medoid; - } - - let max_node_len = ((max_degree as u64 + 1) * (mem::size_of::() as u64)) - + (dims * (mem::size_of::() as u64)); - let num_nodes_per_sector = (SECTOR_LEN as u64) / max_node_len; - - println!("medoid: {}B", medoid); - println!("max_node_len: {}B", max_node_len); - println!("num_nodes_per_sector: {}B", num_nodes_per_sector); - - // SECTOR_LEN buffer for each sector - let mut sector_buf = vec![0u8; SECTOR_LEN]; - let mut node_buf = vec![0u8; max_node_len as usize]; - - let num_nbrs_start = (dims as usize) * mem::size_of::(); - let nbrs_buf_start = num_nbrs_start + mem::size_of::(); - - // number of sectors (1 for meta data) - let num_sectors = round_up(num_pts, num_nodes_per_sector) / num_nodes_per_sector; - let disk_index_file_size = (num_sectors + 1) * (SECTOR_LEN as u64); - - let disk_layout_meta = vec![ - num_pts, - dims, - medoid as u64, - max_node_len, - num_nodes_per_sector, - vamana_frozen_num, - vamana_frozen_loc as u64, - // append_reorder_data - // We are not supporting this. Temporarily write it into the layout so that - // we can leverage C++ query driver to test the disk index - false as u64, - disk_index_file_size, - ]; - - diskann_writer.write(§or_buf)?; - - let mut cur_node_coords = vec![0u8; (dims as usize) * mem::size_of::()]; - let mut cur_node_id = 0u64; - - for sector in 0..num_sectors { - if sector % 100_000 == 0 { - println!("Sector #{} written", sector); - } - sector_buf.fill(0); - - for sector_node_id in 0..num_nodes_per_sector { - if cur_node_id >= num_pts { - break; - } - - node_buf.fill(0); - - // read cur node's num_nbrs - let num_nbrs = vamana_reader.read_u32::()?; - - // sanity checks on num_nbrs - debug_assert!(num_nbrs > 0); - debug_assert!(num_nbrs <= max_degree); - - // write coords of node first - dataset_reader.read(&mut cur_node_coords)?; - node_buf[..cur_node_coords.len()].copy_from_slice(&cur_node_coords); - - // write num_nbrs - LittleEndian::write_u32( - &mut node_buf[num_nbrs_start..(num_nbrs_start + mem::size_of::())], - num_nbrs, - ); - - // write neighbors - let nbrs_buf = &mut node_buf[nbrs_buf_start - ..(nbrs_buf_start + (num_nbrs as usize) * mem::size_of::())]; - vamana_reader.read_exact(nbrs_buf)?; - - // get offset into sector_buf - let sector_node_buf_start = (sector_node_id * max_node_len) as usize; - let sector_node_buf = &mut sector_buf - [sector_node_buf_start..(sector_node_buf_start + max_node_len as usize)]; - sector_node_buf.copy_from_slice(&node_buf[..(max_node_len as usize)]); - - cur_node_id += 1; - } - - // flush sector to disk - diskann_writer.write(§or_buf)?; - } - - diskann_writer.flush()?; - save_bin_u64( - disk_layout_file.as_str(), - &disk_layout_meta, - disk_layout_meta.len(), - 1, - 0, - )?; - - Ok(()) - } - - pub fn index_build_cleanup(&self) -> ANNResult<()> { - fs::remove_file(self.mem_index_file())?; - Ok(()) - } - - pub fn gen_query_warmup_data(&self, sampling_rate: f64) -> ANNResult<()> { - gen_sample_data::( - &self.dataset_file, - &self.warmup_query_prefix(), - sampling_rate, - )?; - Ok(()) - } - - /// Load pre-trained pivot table - pub fn load_pq_pivots_bin( - &self, - num_pq_chunks: &usize, - ) -> ANNResult { - let pq_pivots_path = &self.pq_pivot_file(); - if !file_exists(pq_pivots_path) { - return Err(ANNError::log_pq_error( - "ERROR: PQ k-means pivot file not found.".to_string(), - )); - } - - let (data, offset_num, offset_dim) = load_bin::(pq_pivots_path, 0)?; - let file_offset_data = convert_types_u64_usize(&data, offset_num, offset_dim); - if offset_num != 4 { - let error_message = format!("Error reading pq_pivots file {}. Offsets don't contain correct metadata, # offsets = {}, but expecting 4.", pq_pivots_path, offset_num); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, pivot_num, dim) = load_bin::(pq_pivots_path, file_offset_data[0])?; - let pq_table = data.to_vec(); - if pivot_num != NUM_PQ_CENTROIDS { - let error_message = format!( - "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", - pq_pivots_path, pivot_num, NUM_PQ_CENTROIDS - ); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, centroid_dim, nc) = load_bin::(pq_pivots_path, file_offset_data[1])?; - let centroids = data.to_vec(); - if centroid_dim != dim || nc != 1 { - let error_message = format!("Error reading pq_pivots file {}. file_dim = {}, file_cols = {} but expecting {} entries in 1 dimension.", pq_pivots_path, centroid_dim, nc, dim); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, chunk_offset_num, nc) = load_bin::(pq_pivots_path, file_offset_data[2])?; - let chunk_offsets = convert_types_u32_usize(&data, chunk_offset_num, nc); - if chunk_offset_num != num_pq_chunks + 1 || nc != 1 { - let error_message = format!("Error reading pq_pivots file at chunk offsets; file has nr={}, nc={} but expecting nr={} and nc=1.", chunk_offset_num, nc, num_pq_chunks + 1); - return Err(ANNError::log_pq_error(error_message)); - } - - Ok(PQPivotData { - dim, - pq_table, - centroids, - chunk_offsets - }) - } - - fn mem_index_file(&self) -> String { - self.index_path_prefix.clone() + "_mem.index" - } - - fn disk_index_file(&self) -> String { - self.index_path_prefix.clone() + "_disk.index" - } - - fn warmup_query_prefix(&self) -> String { - self.index_path_prefix.clone() + "_sample" - } - - pub fn pq_pivot_file(&self) -> String { - self.index_path_prefix.clone() + ".bin_pq_pivots.bin" - } - - pub fn compressed_pq_pivot_file(&self) -> String { - self.index_path_prefix.clone() + ".bin_pq_compressed.bin" - } -} - -#[cfg(test)] -mod disk_index_storage_test { - use std::fs; - - use crate::test_utils::get_test_file_path; - - use super::*; - - const TEST_DATA_FILE: &str = "tests/data/siftsmall_learn_256pts.fbin"; - const DISK_INDEX_PATH_PREFIX: &str = "tests/data/disk_index_siftsmall_learn_256pts_R4_L50_A1.2"; - const TRUTH_DISK_LAYOUT: &str = - "tests/data/truth_disk_index_siftsmall_learn_256pts_R4_L50_A1.2_disk.index"; - - #[test] - fn create_disk_layout_test() { - let storage = DiskIndexStorage::::new( - get_test_file_path(TEST_DATA_FILE), - get_test_file_path(DISK_INDEX_PATH_PREFIX), - ).unwrap(); - storage.create_disk_layout().unwrap(); - - let disk_layout_file = storage.disk_index_file(); - let rust_disk_layout = fs::read(disk_layout_file.as_str()).unwrap(); - let truth_disk_layout = fs::read(get_test_file_path(TRUTH_DISK_LAYOUT).as_str()).unwrap(); - - assert!(rust_disk_layout == truth_disk_layout); - - fs::remove_file(disk_layout_file.as_str()).expect("Failed to delete file"); - } - - #[test] - fn load_pivot_test() { - let dim: usize = 128; - let num_pq_chunk: usize = 1; - let pivot_file_prefix: &str = "tests/data/siftsmall_learn"; - let storage = DiskIndexStorage::::new( - get_test_file_path(TEST_DATA_FILE), - pivot_file_prefix.to_string(), - ).unwrap(); - - let pq_pivot_data = - storage.load_pq_pivots_bin(&num_pq_chunk).unwrap(); - - assert_eq!(pq_pivot_data.pq_table.len(), NUM_PQ_CENTROIDS * dim); - assert_eq!(pq_pivot_data.centroids.len(), dim); - - assert_eq!(pq_pivot_data.chunk_offsets[0], 0); - assert_eq!(pq_pivot_data.chunk_offsets[1], dim); - assert_eq!(pq_pivot_data.chunk_offsets.len(), num_pq_chunk + 1); - } - - #[test] - #[should_panic(expected = "ERROR: PQ k-means pivot file not found.")] - fn load_pivot_file_not_exist_test() { - let num_pq_chunk: usize = 1; - let pivot_file_prefix: &str = "tests/data/siftsmall_learn_file_not_exist"; - let storage = DiskIndexStorage::::new( - get_test_file_path(TEST_DATA_FILE), - pivot_file_prefix.to_string(), - ).unwrap(); - let _ = storage.load_pq_pivots_bin(&num_pq_chunk).unwrap(); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/mod.rs deleted file mode 100644 index 03c5b8e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/mod.rs +++ /dev/null @@ -1,12 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -mod disk_index_storage; -pub use disk_index_storage::*; - -mod disk_graph_storage; -pub use disk_graph_storage::*; - -mod pq_storage; -pub use pq_storage::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/pq_storage.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/pq_storage.rs deleted file mode 100644 index b1d3fa0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/storage/pq_storage.rs +++ /dev/null @@ -1,367 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use byteorder::{LittleEndian, ReadBytesExt}; -use rand::distributions::{Distribution, Uniform}; -use std::fs::File; -use std::io::{Read, Seek, SeekFrom, Write}; -use std::mem; - -use crate::common::{ANNError, ANNResult}; -use crate::utils::CachedReader; -use crate::utils::{ - convert_types_u32_usize, convert_types_u64_usize, convert_types_usize_u32, - convert_types_usize_u64, convert_types_usize_u8, save_bin_f32, save_bin_u32, save_bin_u64, -}; -use crate::utils::{file_exists, load_bin, open_file_to_write, METADATA_SIZE}; - -#[derive(Debug)] -pub struct PQStorage { - /// Pivot table path - pivot_file: String, - - /// Compressed pivot path - compressed_pivot_file: String, - - /// Data used to construct PQ table and PQ compressed table - pq_data_file: String, - - /// PQ data reader - pq_data_file_reader: File, -} - -impl PQStorage { - pub fn new( - pivot_file: &str, - compressed_pivot_file: &str, - pq_data_file: &str, - ) -> std::io::Result { - let pq_data_file_reader = File::open(pq_data_file)?; - Ok(Self { - pivot_file: pivot_file.to_string(), - compressed_pivot_file: compressed_pivot_file.to_string(), - pq_data_file: pq_data_file.to_string(), - pq_data_file_reader, - }) - } - - pub fn write_compressed_pivot_metadata(&self, npts: i32, pq_chunk: i32) -> std::io::Result<()> { - let mut writer = open_file_to_write(&self.compressed_pivot_file)?; - writer.write_all(&npts.to_le_bytes())?; - writer.write_all(&pq_chunk.to_le_bytes())?; - Ok(()) - } - - pub fn write_compressed_pivot_data( - &self, - compressed_base: &[usize], - num_centers: usize, - block_size: usize, - num_pq_chunks: usize, - ) -> std::io::Result<()> { - let mut writer = open_file_to_write(&self.compressed_pivot_file)?; - writer.seek(SeekFrom::Start((std::mem::size_of::() * 2) as u64))?; - if num_centers > 256 { - writer.write_all(unsafe { - std::slice::from_raw_parts( - compressed_base.as_ptr() as *const u8, - block_size * num_pq_chunks * std::mem::size_of::(), - ) - })?; - } else { - let compressed_base_u8 = - convert_types_usize_u8(compressed_base, block_size, num_pq_chunks); - writer.write_all(&compressed_base_u8)?; - } - Ok(()) - } - - pub fn write_pivot_data( - &self, - full_pivot_data: &[f32], - centroid: &[f32], - chunk_offsets: &[usize], - num_centers: usize, - dim: usize, - ) -> std::io::Result<()> { - let mut cumul_bytes: Vec = vec![0; 4]; - cumul_bytes[0] = METADATA_SIZE; - cumul_bytes[1] = cumul_bytes[0] - + save_bin_f32( - &self.pivot_file, - full_pivot_data, - num_centers, - dim, - cumul_bytes[0], - )?; - cumul_bytes[2] = - cumul_bytes[1] + save_bin_f32(&self.pivot_file, centroid, dim, 1, cumul_bytes[1])?; - - // Because the writer only can write u32, u64 but not usize, so we need to convert the type first. - let chunk_offsets_u64 = convert_types_usize_u32(chunk_offsets, chunk_offsets.len(), 1); - cumul_bytes[3] = cumul_bytes[2] - + save_bin_u32( - &self.pivot_file, - &chunk_offsets_u64, - chunk_offsets.len(), - 1, - cumul_bytes[2], - )?; - - let cumul_bytes_u64 = convert_types_usize_u64(&cumul_bytes, 4, 1); - save_bin_u64(&self.pivot_file, &cumul_bytes_u64, cumul_bytes.len(), 1, 0)?; - - Ok(()) - } - - pub fn pivot_data_exist(&self) -> bool { - file_exists(&self.pivot_file) - } - - pub fn read_pivot_metadata(&self) -> std::io::Result<(usize, usize)> { - let (_, file_num_centers, file_dim) = load_bin::(&self.pivot_file, METADATA_SIZE)?; - Ok((file_num_centers, file_dim)) - } - - pub fn load_pivot_data( - &self, - num_pq_chunks: &usize, - num_centers: &usize, - dim: &usize, - ) -> ANNResult<(Vec, Vec, Vec)> { - // Load file offset data. File saved as offset data(4*1) -> pivot data(centroid num*dim) -> centroid of dim data(dim*1) -> chunk offset data(chunksize+1*1) - // Because we only can write u64 rather than usize, so the file stored as u64 type. Need to convert to usize when use. - let (data, offset_num, nc) = load_bin::(&self.pivot_file, 0)?; - let file_offset_data = convert_types_u64_usize(&data, offset_num, nc); - if offset_num != 4 { - let error_message = format!("Error reading pq_pivots file {}. Offsets don't contain correct metadata, # offsets = {}, but expecting 4.", &self.pivot_file, offset_num); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, pivot_num, pivot_dim) = load_bin::(&self.pivot_file, file_offset_data[0])?; - let full_pivot_data = data; - if pivot_num != *num_centers || pivot_dim != *dim { - let error_message = format!("Error reading pq_pivots file {}. file_num_centers = {}, file_dim = {} but expecting {} centers in {} dimensions.", &self.pivot_file, pivot_num, pivot_dim, num_centers, dim); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, centroid_dim, nc) = load_bin::(&self.pivot_file, file_offset_data[1])?; - let centroid = data; - if centroid_dim != *dim || nc != 1 { - let error_message = format!("Error reading pq_pivots file {}. file_dim = {}, file_cols = {} but expecting {} entries in 1 dimension.", &self.pivot_file, centroid_dim, nc, dim); - return Err(ANNError::log_pq_error(error_message)); - } - - let (data, chunk_offset_number, nc) = - load_bin::(&self.pivot_file, file_offset_data[2])?; - let chunk_offsets = convert_types_u32_usize(&data, chunk_offset_number, nc); - if chunk_offset_number != *num_pq_chunks + 1 || nc != 1 { - let error_message = format!("Error reading pq_pivots file at chunk offsets; file has nr={}, nc={} but expecting nr={} and nc=1.", chunk_offset_number, nc, num_pq_chunks + 1); - return Err(ANNError::log_pq_error(error_message)); - } - Ok((full_pivot_data, centroid, chunk_offsets)) - } - - pub fn read_pq_data_metadata(&mut self) -> std::io::Result<(usize, usize)> { - let npts_i32 = self.pq_data_file_reader.read_i32::()?; - let dim_i32 = self.pq_data_file_reader.read_i32::()?; - let num_points = npts_i32 as usize; - let dim = dim_i32 as usize; - Ok((num_points, dim)) - } - - pub fn read_pq_block_data( - &mut self, - cur_block_size: usize, - dim: usize, - ) -> std::io::Result> { - let mut buf = vec![0u8; cur_block_size * dim * std::mem::size_of::()]; - self.pq_data_file_reader.read_exact(&mut buf)?; - - let ptr = buf.as_ptr() as *const T; - let block_data = unsafe { std::slice::from_raw_parts(ptr, cur_block_size * dim) }; - Ok(block_data.to_vec()) - } - - /// streams data from the file, and samples each vector with probability p_val - /// and returns a matrix of size slice_size* ndims as floating point type. - /// the slice_size and ndims are set inside the function. - /// # Arguments - /// * `file_name` - filename where the data is - /// * `p_val` - possibility to sample data - /// * `sampled_vectors` - sampled vector chose by p_val possibility - /// * `slice_size` - how many sampled data return - /// * `dim` - each sample data dimension - pub fn gen_random_slice>( - &self, - mut p_val: f64, - ) -> ANNResult<(Vec, usize, usize)> { - let read_blk_size = 64 * 1024 * 1024; - let mut reader = CachedReader::new(&self.pq_data_file, read_blk_size)?; - - let npts = reader.read_u32()? as usize; - let dim = reader.read_u32()? as usize; - let mut sampled_vectors: Vec = Vec::new(); - let mut slice_size = 0; - p_val = if p_val < 1f64 { p_val } else { 1f64 }; - - let mut generator = rand::thread_rng(); - let distribution = Uniform::from(0.0..1.0); - - for _ in 0..npts { - let mut cur_vector_bytes = vec![0u8; dim * mem::size_of::()]; - reader.read(&mut cur_vector_bytes)?; - let random_value = distribution.sample(&mut generator); - if random_value < p_val { - let ptr = cur_vector_bytes.as_ptr() as *const T; - let cur_vector_t = unsafe { std::slice::from_raw_parts(ptr, dim) }; - sampled_vectors.extend(cur_vector_t.iter().map(|&t| t.into())); - slice_size += 1; - } - } - - Ok((sampled_vectors, slice_size, dim)) - } -} - -#[cfg(test)] -mod pq_storage_tests { - use rand::Rng; - - use super::*; - use crate::utils::gen_random_slice; - - const DATA_FILE: &str = "tests/data/siftsmall_learn.bin"; - const PQ_PIVOT_PATH: &str = "tests/data/siftsmall_learn.bin_pq_pivots.bin"; - const PQ_COMPRESSED_PATH: &str = "tests/data/empty_pq_compressed.bin"; - - #[test] - fn new_test() { - let result = PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, DATA_FILE); - assert!(result.is_ok()); - } - - #[test] - fn write_compressed_pivot_metadata_test() { - let compress_pivot_path = "write_compressed_pivot_metadata_test.bin"; - let result = PQStorage::new(PQ_PIVOT_PATH, compress_pivot_path, DATA_FILE).unwrap(); - - _ = result.write_compressed_pivot_metadata(100, 20); - let mut result_reader = File::open(compress_pivot_path).unwrap(); - let npts_i32 = result_reader.read_i32::().unwrap(); - let dim_i32 = result_reader.read_i32::().unwrap(); - - assert_eq!(npts_i32, 100); - assert_eq!(dim_i32, 20); - - std::fs::remove_file(compress_pivot_path).unwrap(); - } - - #[test] - fn write_compressed_pivot_data_test() { - let compress_pivot_path = "write_compressed_pivot_data_test.bin"; - let result = PQStorage::new(PQ_PIVOT_PATH, compress_pivot_path, DATA_FILE).unwrap(); - - let mut rng = rand::thread_rng(); - - let num_centers = 256; - let block_size = 4; - let num_pq_chunks = 2; - let compressed_base: Vec = (0..block_size * num_pq_chunks) - .map(|_| rng.gen_range(0..num_centers)) - .collect(); - _ = result.write_compressed_pivot_data( - &compressed_base, - num_centers, - block_size, - num_pq_chunks, - ); - - let mut result_reader = File::open(compress_pivot_path).unwrap(); - _ = result_reader.read_i32::().unwrap(); - _ = result_reader.read_i32::().unwrap(); - let mut buf = vec![0u8; block_size * num_pq_chunks * std::mem::size_of::()]; - result_reader.read_exact(&mut buf).unwrap(); - - let ptr = buf.as_ptr() as *const u8; - let block_data = unsafe { std::slice::from_raw_parts(ptr, block_size * num_pq_chunks) }; - - for index in 0..block_data.len() { - assert_eq!(compressed_base[index], block_data[index] as usize); - } - std::fs::remove_file(compress_pivot_path).unwrap(); - } - - #[test] - fn pivot_data_exist_test() { - let result = PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, DATA_FILE).unwrap(); - assert!(result.pivot_data_exist()); - - let pivot_path = "not_exist_pivot_path.bin"; - let result = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, DATA_FILE).unwrap(); - assert!(!result.pivot_data_exist()); - } - - #[test] - fn read_pivot_metadata_test() { - let result = PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, DATA_FILE).unwrap(); - let (npt, dim) = result.read_pivot_metadata().unwrap(); - - assert_eq!(npt, 256); - assert_eq!(dim, 128); - } - - #[test] - fn load_pivot_data_test() { - let result = PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, DATA_FILE).unwrap(); - let (pq_pivot_data, centroids, chunk_offsets) = - result.load_pivot_data(&1, &256, &128).unwrap(); - - assert_eq!(pq_pivot_data.len(), 256 * 128); - assert_eq!(centroids.len(), 128); - assert_eq!(chunk_offsets.len(), 2); - } - - #[test] - fn read_pq_data_metadata_test() { - let mut result = PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, DATA_FILE).unwrap(); - let (npt, dim) = result.read_pq_data_metadata().unwrap(); - - assert_eq!(npt, 25000); - assert_eq!(dim, 128); - } - - #[test] - fn gen_random_slice_test() { - let file_name = "gen_random_slice_test.bin"; - //npoints=2, dim=8 - let data: [u8; 72] = [ - 2, 0, 0, 0, 8, 0, 0, 0, 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, - 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, - 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, - 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, 0x00, 0x00, 0x50, 0x41, - 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x00, 0x80, 0x41, - ]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let (sampled_vectors, slice_size, ndims) = - gen_random_slice::(file_name, 1f64).unwrap(); - let mut start = 8; - (0..sampled_vectors.len()).for_each(|i| { - assert_eq!(sampled_vectors[i].to_le_bytes(), data[start..start + 4]); - start += 4; - }); - assert_eq!(sampled_vectors.len(), 16); - assert_eq!(slice_size, 2); - assert_eq!(ndims, 8); - - let (sampled_vectors, slice_size, ndims) = - gen_random_slice::(file_name, 0f64).unwrap(); - assert_eq!(sampled_vectors.len(), 0); - assert_eq!(slice_size, 0); - assert_eq!(ndims, 8); - - std::fs::remove_file(file_name).expect("Failed to delete file"); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/inmem_index_initialization.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/inmem_index_initialization.rs deleted file mode 100644 index db3b581..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/inmem_index_initialization.rs +++ /dev/null @@ -1,74 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use vector::Metric; - -use crate::index::InmemIndex; -use crate::model::configuration::index_write_parameters::IndexWriteParametersBuilder; -use crate::model::{IndexConfiguration}; -use crate::model::vertex::DIM_128; -use crate::utils::{file_exists, load_metadata_from_file}; - -use super::get_test_file_path; - -// f32, 128 DIM and 256 points source data -const TEST_DATA_FILE: &str = "tests/data/siftsmall_learn_256pts.fbin"; -const NUM_POINTS_TO_LOAD: usize = 256; - -pub fn create_index_with_test_data() -> InmemIndex { - let index_write_parameters = IndexWriteParametersBuilder::new(50, 4).with_alpha(1.2).build(); - let config = IndexConfiguration::new( - Metric::L2, - 128, - 128, - 256, - false, - 0, - false, - 0, - 1.0f32, - index_write_parameters); - let mut index: InmemIndex = InmemIndex::new(config).unwrap(); - - build_test_index(&mut index, get_test_file_path(TEST_DATA_FILE).as_str(), NUM_POINTS_TO_LOAD); - - index.start = index.dataset.calculate_medoid_point_id().unwrap(); - - index -} - -fn build_test_index(index: &mut InmemIndex, filename: &str, num_points_to_load: usize) { - if !file_exists(filename) { - panic!("ERROR: Data file {} does not exist.", filename); - } - - let (file_num_points, file_dim) = load_metadata_from_file(filename).unwrap(); - if file_num_points > index.configuration.max_points { - panic!( - "ERROR: Driver requests loading {} points and file has {} points, - but index can support only {} points as specified in configuration.", - num_points_to_load, file_num_points, index.configuration.max_points - ); - } - - if num_points_to_load > file_num_points { - panic!( - "ERROR: Driver requests loading {} points and file has only {} points.", - num_points_to_load, file_num_points - ); - } - - if file_dim != index.configuration.dim { - panic!( - "ERROR: Driver requests loading {} dimension, but file has {} dimension.", - index.configuration.dim, file_dim - ); - } - - index.dataset.build_from_file(filename, num_points_to_load).unwrap(); - - println!("Using only first {} from file.", num_points_to_load); - - index.num_active_pts = num_points_to_load; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/mod.rs deleted file mode 100644 index fc8de5f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/test_utils/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod inmem_index_initialization; - -/// test files should be placed under tests folder -pub fn get_test_file_path(relative_path: &str) -> String { - format!("{}/{}", env!("CARGO_MANIFEST_DIR"), relative_path) -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/bit_vec_extension.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/bit_vec_extension.rs deleted file mode 100644 index 9571a72..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/bit_vec_extension.rs +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::cmp::Ordering; - -use bit_vec::BitVec; - -pub trait BitVecExtension { - fn resize(&mut self, new_len: usize, value: bool); -} - -impl BitVecExtension for BitVec { - fn resize(&mut self, new_len: usize, value: bool) { - let old_len = self.len(); - match new_len.cmp(&old_len) { - Ordering::Less => self.truncate(new_len), - Ordering::Greater => self.grow(new_len - old_len, value), - Ordering::Equal => {} - } - } -} - -#[cfg(test)] -mod bit_vec_extension_test { - use super::*; - - #[test] - fn resize_test() { - let mut bitset = BitVec::new(); - - bitset.resize(10, false); - assert_eq!(bitset.len(), 10); - assert!(bitset.none()); - - bitset.resize(11, true); - assert_eq!(bitset.len(), 11); - assert!(bitset[10]); - - bitset.resize(5, false); - assert_eq!(bitset.len(), 5); - assert!(bitset.none()); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_reader.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_reader.rs deleted file mode 100644 index 1a21f1a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_reader.rs +++ /dev/null @@ -1,160 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::fs::File; -use std::io::{Seek, Read}; - -use crate::common::{ANNResult, ANNError}; - -/// Sequential cached reads -pub struct CachedReader { - /// File reader - reader: File, - - /// # bytes to cache in one shot read - cache_size: u64, - - /// Underlying buf for cache - cache_buf: Vec, - - /// Offset into cache_buf for cur_pos - cur_off: u64, - - /// File size - fsize: u64, -} - -impl CachedReader { - pub fn new(filename: &str, cache_size: u64) -> std::io::Result { - let mut reader = File::open(filename)?; - let metadata = reader.metadata()?; - let fsize = metadata.len(); - - let cache_size = cache_size.min(fsize); - let mut cache_buf = vec![0; cache_size as usize]; - reader.read_exact(&mut cache_buf)?; - println!("Opened: {}, size: {}, cache_size: {}", filename, fsize, cache_size); - - Ok(Self { - reader, - cache_size, - cache_buf, - cur_off: 0, - fsize, - }) - } - - pub fn get_file_size(&self) -> u64 { - self.fsize - } - - pub fn read(&mut self, read_buf: &mut [u8]) -> ANNResult<()> { - let n_bytes = read_buf.len() as u64; - if n_bytes <= (self.cache_size - self.cur_off) { - // case 1: cache contains all data - read_buf.copy_from_slice(&self.cache_buf[(self.cur_off as usize)..(self.cur_off as usize + n_bytes as usize)]); - self.cur_off += n_bytes; - } else { - // case 2: cache contains some data - let cached_bytes = self.cache_size - self.cur_off; - if n_bytes - cached_bytes > self.fsize - self.reader.stream_position()? { - return Err(ANNError::log_index_error(format!( - "Reading beyond end of file, n_bytes: {} cached_bytes: {} fsize: {} current pos: {}", - n_bytes, cached_bytes, self.fsize, self.reader.stream_position()?)) - ); - } - - read_buf[..cached_bytes as usize].copy_from_slice(&self.cache_buf[self.cur_off as usize..]); - // go to disk and fetch more data - self.reader.read_exact(&mut read_buf[cached_bytes as usize..])?; - // reset cur off - self.cur_off = self.cache_size; - - let size_left = self.fsize - self.reader.stream_position()?; - if size_left >= self.cache_size { - self.reader.read_exact(&mut self.cache_buf)?; - self.cur_off = 0; - } - // note that if size_left < cache_size, then cur_off = cache_size, - // so subsequent reads will all be directly from file - } - Ok(()) - } - - pub fn read_u32(&mut self) -> ANNResult { - let mut bytes = [0u8; 4]; - self.read(&mut bytes)?; - Ok(u32::from_le_bytes(bytes)) - } -} - -#[cfg(test)] -mod cached_reader_test { - use std::fs; - - use super::*; - - #[test] - fn cached_reader_works() { - let file_name = "cached_reader_works_test.bin"; - //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8] - let data: [u8; 72] = [2, 0, 1, 2, 8, 0, 1, 3, - 0x00, 0x01, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11, 0x80, 0x41]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let mut reader = CachedReader::new(file_name, 8).unwrap(); - assert_eq!(reader.get_file_size(), 72); - assert_eq!(reader.cache_size, 8); - - let mut all_from_cache_buf = vec![0; 4]; - reader.read(all_from_cache_buf.as_mut_slice()).unwrap(); - assert_eq!(all_from_cache_buf, [2, 0, 1, 2]); - assert_eq!(reader.cur_off, 4); - - let mut partial_from_cache_buf = vec![0; 6]; - reader.read(partial_from_cache_buf.as_mut_slice()).unwrap(); - assert_eq!(partial_from_cache_buf, [8, 0, 1, 3, 0x00, 0x01]); - assert_eq!(reader.cur_off, 0); - - let mut over_cache_size_buf = vec![0; 60]; - reader.read(over_cache_size_buf.as_mut_slice()).unwrap(); - assert_eq!( - over_cache_size_buf, - [0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11] - ); - - let mut remaining_less_than_cache_size_buf = vec![0; 2]; - reader.read(remaining_less_than_cache_size_buf.as_mut_slice()).unwrap(); - assert_eq!(remaining_less_than_cache_size_buf, [0x80, 0x41]); - assert_eq!(reader.cur_off, reader.cache_size); - - fs::remove_file(file_name).expect("Failed to delete file"); - } - - #[test] - #[should_panic(expected = "n_bytes: 73 cached_bytes: 8 fsize: 72 current pos: 8")] - fn failed_for_reading_beyond_end_of_file() { - let file_name = "failed_for_reading_beyond_end_of_file_test.bin"; - //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8] - let data: [u8; 72] = [2, 0, 1, 2, 8, 0, 1, 3, - 0x00, 0x01, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11, 0x80, 0x41]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let mut reader = CachedReader::new(file_name, 8).unwrap(); - fs::remove_file(file_name).expect("Failed to delete file"); - - let mut over_size_buf = vec![0; 73]; - reader.read(over_size_buf.as_mut_slice()).unwrap(); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_writer.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_writer.rs deleted file mode 100644 index d3929be..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/cached_writer.rs +++ /dev/null @@ -1,142 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::io::{Write, Seek, SeekFrom}; -use std::fs::{OpenOptions, File}; -use std::path::Path; - -pub struct CachedWriter { - /// File writer - writer: File, - - /// # bytes to cache for one shot write - cache_size: u64, - - /// Underlying buf for cache - cache_buf: Vec, - - /// Offset into cache_buf for cur_pos - cur_off: u64, - - /// File size - fsize: u64, -} - -impl CachedWriter { - pub fn new(filename: &str, cache_size: u64) -> std::io::Result { - let writer = OpenOptions::new() - .write(true) - .create(true) - .open(Path::new(filename))?; - - if cache_size == 0 { - return Err(std::io::Error::new(std::io::ErrorKind::Other, "Cache size must be greater than 0")); - } - - println!("Opened: {}, cache_size: {}", filename, cache_size); - Ok(Self { - writer, - cache_size, - cache_buf: vec![0; cache_size as usize], - cur_off: 0, - fsize: 0, - }) - } - - pub fn flush(&mut self) -> std::io::Result<()> { - // dump any remaining data in memory - if self.cur_off > 0 { - self.flush_cache()?; - } - - self.writer.flush()?; - println!("Finished writing {}B", self.fsize); - Ok(()) - } - - pub fn get_file_size(&self) -> u64 { - self.fsize - } - - /// Writes n_bytes from write_buf to the underlying cache - pub fn write(&mut self, write_buf: &[u8]) -> std::io::Result<()> { - let n_bytes = write_buf.len() as u64; - if n_bytes <= (self.cache_size - self.cur_off) { - // case 1: cache can take all data - self.cache_buf[(self.cur_off as usize)..((self.cur_off + n_bytes) as usize)].copy_from_slice(&write_buf[..n_bytes as usize]); - self.cur_off += n_bytes; - } else { - // case 2: cache cant take all data - // go to disk and write existing cache data - self.writer.write_all(&self.cache_buf[..self.cur_off as usize])?; - self.fsize += self.cur_off; - // write the new data to disk - self.writer.write_all(write_buf)?; - self.fsize += n_bytes; - // clear cache data and reset cur_off - self.cache_buf.fill(0); - self.cur_off = 0; - } - Ok(()) - } - - pub fn reset(&mut self) -> std::io::Result<()> { - self.flush_cache()?; - self.writer.seek(SeekFrom::Start(0))?; - Ok(()) - } - - fn flush_cache(&mut self) -> std::io::Result<()> { - self.writer.write_all(&self.cache_buf[..self.cur_off as usize])?; - self.fsize += self.cur_off; - self.cache_buf.fill(0); - self.cur_off = 0; - Ok(()) - } -} - -impl Drop for CachedWriter { - fn drop(&mut self) { - let _ = self.flush(); - } -} - -#[cfg(test)] -mod cached_writer_test { - use std::fs; - - use super::*; - - #[test] - fn cached_writer_works() { - let file_name = "cached_writer_works_test.bin"; - //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8] - let data: [u8; 72] = [2, 0, 1, 2, 8, 0, 1, 3, - 0x00, 0x01, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x11, 0x80, 0x41]; - - let mut writer = CachedWriter::new(file_name, 8).unwrap(); - assert_eq!(writer.get_file_size(), 0); - assert_eq!(writer.cache_size, 8); - assert_eq!(writer.get_file_size(), 0); - - let cache_all_buf = &data[0..4]; - writer.write(cache_all_buf).unwrap(); - assert_eq!(&writer.cache_buf[..4], cache_all_buf); - assert_eq!(&writer.cache_buf[4..], vec![0; 4]); - assert_eq!(writer.cur_off, 4); - assert_eq!(writer.get_file_size(), 0); - - let write_all_buf = &data[4..10]; - writer.write(write_all_buf).unwrap(); - assert_eq!(writer.cache_buf, vec![0; 8]); - assert_eq!(writer.cur_off, 0); - assert_eq!(writer.get_file_size(), 10); - - fs::remove_file(file_name).expect("Failed to delete file"); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/file_util.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/file_util.rs deleted file mode 100644 index f187d01..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/file_util.rs +++ /dev/null @@ -1,377 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! File operations - -use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; -use std::{mem, io}; -use std::fs::{self, File, OpenOptions}; -use std::io::{Read, BufReader, Write, Seek, SeekFrom}; -use std::path::Path; - -use crate::model::data_store::DatasetDto; - -/// Read metadata of data file. -pub fn load_metadata_from_file(file_name: &str) -> std::io::Result<(usize, usize)> { - let file = File::open(file_name)?; - let mut reader = BufReader::new(file); - - let npoints = reader.read_i32::()? as usize; - let ndims = reader.read_i32::()? as usize; - - Ok((npoints, ndims)) -} - -/// Read the deleted vertex ids from file. -pub fn load_ids_to_delete_from_file(file_name: &str) -> std::io::Result<(usize, Vec)> { - // The first 4 bytes are the number of vector ids. - // The rest of the file are the vector ids in the format of usize. - // The vector ids are sorted in ascending order. - let mut file = File::open(file_name)?; - let num_ids = file.read_u32::()? as usize; - - let mut ids = Vec::with_capacity(num_ids); - for _ in 0..num_ids { - let id = file.read_u32::()?; - ids.push(id); - } - - Ok((num_ids, ids)) -} - -/// Copy data from file -/// # Arguments -/// * `bin_file` - filename where the data is -/// * `data` - destination dataset dto to which the data is copied -/// * `pts_offset` - offset of points. data will be loaded after this point in dataset -/// * `npts` - number of points read from bin_file -/// * `dim` - point dimension read from bin_file -/// * `rounded_dim` - rounded dimension (padding zero if it's > dim) -/// # Return -/// * `npts` - number of points read from bin_file -/// * `dim` - point dimension read from bin_file -pub fn copy_aligned_data_from_file( - bin_file: &str, - dataset_dto: DatasetDto, - pts_offset: usize, -) -> std::io::Result<(usize, usize)> { - let mut reader = File::open(bin_file)?; - - let npts = reader.read_i32::()? as usize; - let dim = reader.read_i32::()? as usize; - let rounded_dim = dataset_dto.rounded_dim; - let offset = pts_offset * rounded_dim; - - for i in 0..npts { - let data_slice = &mut dataset_dto.data[offset + i * rounded_dim..offset + i * rounded_dim + dim]; - let mut buf = vec![0u8; dim * mem::size_of::()]; - reader.read_exact(&mut buf)?; - - let ptr = buf.as_ptr() as *const T; - let temp_slice = unsafe { std::slice::from_raw_parts(ptr, dim) }; - data_slice.copy_from_slice(temp_slice); - - (i * rounded_dim + dim..i * rounded_dim + rounded_dim).for_each(|j| { - dataset_dto.data[j] = T::default(); - }); - } - - Ok((npts, dim)) -} - -/// Open a file to write -/// # Arguments -/// * `writer` - mutable File reference -/// * `file_name` - file name -#[inline] -pub fn open_file_to_write(file_name: &str) -> std::io::Result { - OpenOptions::new() - .write(true) - .create(true) - .open(Path::new(file_name)) -} - -/// Delete a file -/// # Arguments -/// * `file_name` - file name -pub fn delete_file(file_name: &str) -> std::io::Result<()> { - if file_exists(file_name) { - fs::remove_file(file_name)?; - } - - Ok(()) -} - -/// Check whether file exists or not -pub fn file_exists(filename: &str) -> bool { - std::path::Path::new(filename).exists() -} - -/// Save data to file -/// # Arguments -/// * `filename` - filename where the data is -/// * `data` - information data -/// * `npts` - number of points -/// * `ndims` - point dimension -/// * `aligned_dim` - aligned dimension -/// * `offset` - data offset in file -pub fn save_data_in_base_dimensions( - filename: &str, - data: &mut [T], - npts: usize, - ndims: usize, - aligned_dim: usize, - offset: usize, -) -> std::io::Result { - let mut writer = open_file_to_write(filename)?; - let npts_i32 = npts as i32; - let ndims_i32 = ndims as i32; - let bytes_written = 2 * std::mem::size_of::() + npts * ndims * (std::mem::size_of::()); - - writer.seek(std::io::SeekFrom::Start(offset as u64))?; - writer.write_all(&npts_i32.to_le_bytes())?; - writer.write_all(&ndims_i32.to_le_bytes())?; - let data_ptr = data.as_ptr() as *const u8; - for i in 0..npts { - let middle_offset = i * aligned_dim * std::mem::size_of::(); - let middle_slice = unsafe { std::slice::from_raw_parts(data_ptr.add(middle_offset), ndims * std::mem::size_of::()) }; - writer.write_all(middle_slice)?; - } - writer.flush()?; - Ok(bytes_written) -} - -/// Read data file -/// # Arguments -/// * `bin_file` - filename where the data is -/// * `file_offset` - data offset in file -/// * `data` - information data -/// * `npts` - number of points -/// * `ndims` - point dimension -pub fn load_bin( - bin_file: &str, - file_offset: usize) -> std::io::Result<(Vec, usize, usize)> -{ - let mut reader = File::open(bin_file)?; - reader.seek(std::io::SeekFrom::Start(file_offset as u64))?; - let npts = reader.read_i32::()? as usize; - let dim = reader.read_i32::()? as usize; - - let size = npts * dim * std::mem::size_of::(); - let mut buf = vec![0u8; size]; - reader.read_exact(&mut buf)?; - - let ptr = buf.as_ptr() as *const T; - let data = unsafe { std::slice::from_raw_parts(ptr, npts * dim)}; - - Ok((data.to_vec(), npts, dim)) -} - -/// Get file size -pub fn get_file_size(filename: &str) -> io::Result { - let reader = File::open(filename)?; - let metadata = reader.metadata()?; - Ok(metadata.len()) -} - -macro_rules! save_bin { - ($name:ident, $t:ty, $write_func:ident) => { - /// Write data into file - pub fn $name(filename: &str, data: &[$t], num_pts: usize, dims: usize, offset: usize) -> std::io::Result { - let mut writer = open_file_to_write(filename)?; - - println!("Writing bin: {}", filename); - writer.seek(SeekFrom::Start(offset as u64))?; - let num_pts_i32 = num_pts as i32; - let dims_i32 = dims as i32; - let bytes_written = num_pts * dims * mem::size_of::<$t>() + 2 * mem::size_of::(); - - writer.write_i32::(num_pts_i32)?; - writer.write_i32::(dims_i32)?; - println!("bin: #pts = {}, #dims = {}, size = {}B", num_pts, dims, bytes_written); - - for item in data.iter() { - writer.$write_func::(*item)?; - } - - writer.flush()?; - - println!("Finished writing bin."); - Ok(bytes_written) - } - }; -} - -save_bin!(save_bin_f32, f32, write_f32); -save_bin!(save_bin_u64, u64, write_u64); -save_bin!(save_bin_u32, u32, write_u32); - -#[cfg(test)] -mod file_util_test { - use crate::model::data_store::InmemDataset; - use std::fs; - use super::*; - - pub const DIM_8: usize = 8; - - #[test] - fn load_metadata_test() { - let file_name = "test_load_metadata_test.bin"; - let data = [200, 0, 0, 0, 128, 0, 0, 0]; // 200 and 128 in little endian bytes - std::fs::write(file_name, data).expect("Failed to write sample file"); - match load_metadata_from_file(file_name) { - Ok((npoints, ndims)) => { - assert!(npoints == 200); - assert!(ndims == 128); - }, - Err(_e) => {}, - } - fs::remove_file(file_name).expect("Failed to delete file"); - } - - #[test] - fn load_data_test() { - let file_name = "test_load_data_test.bin"; - //npoints=2, dim=8, 2 vectors [1.0;8] [2.0;8] - let data: [u8; 72] = [2, 0, 0, 0, 8, 0, 0, 0, - 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x00, 0x80, 0x41]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let mut dataset = InmemDataset::::new(2, 1f32).unwrap(); - - match copy_aligned_data_from_file(file_name, dataset.into_dto(), 0) { - Ok((num_points, dim)) => { - fs::remove_file(file_name).expect("Failed to delete file"); - assert!(num_points == 2); - assert!(dim == 8); - assert!(dataset.data.len() == 16); - - let first_vertex = dataset.get_vertex(0).unwrap(); - let second_vertex = dataset.get_vertex(1).unwrap(); - - assert!(*first_vertex.vector() == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); - assert!(*second_vertex.vector() == [9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]); - }, - Err(e) => { - fs::remove_file(file_name).expect("Failed to delete file"); - panic!("{}", e) - }, - } - } - - #[test] - fn open_file_to_write_test() { - let file_name = "test_open_file_to_write_test.bin"; - let mut writer = File::create(file_name).unwrap(); - let data = [200, 0, 0, 0, 128, 0, 0, 0]; - writer.write(&data).expect("Failed to write sample file"); - - let _ = open_file_to_write(file_name); - - fs::remove_file(file_name).expect("Failed to delete file"); - } - - #[test] - fn delete_file_test() { - let file_name = "test_delete_file_test.bin"; - let mut file = File::create(file_name).unwrap(); - writeln!(file, "test delete file").unwrap(); - - let result = delete_file(file_name); - - assert!(result.is_ok()); - assert!(fs::metadata(file_name).is_err()); - } - - #[test] - fn save_data_in_base_dimensions_test() { - //npoints=2, dim=8 - let mut data: [u8; 72] = [2, 0, 0, 0, 8, 0, 0, 0, - 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x00, 0x80, 0x41]; - let num_points = 2; - let dim = DIM_8; - let data_file = "save_data_in_base_dimensions_test.data"; - match save_data_in_base_dimensions(data_file, &mut data, num_points, dim, DIM_8, 0) { - Ok(num) => { - assert!(file_exists(data_file)); - assert_eq!(num, 2 * std::mem::size_of::() + num_points * dim * std::mem::size_of::()); - fs::remove_file(data_file).expect("Failed to delete file"); - }, - Err(e) => { - fs::remove_file(data_file).expect("Failed to delete file"); - panic!("{}", e) - } - } - } - - #[test] - fn save_bin_test() { - let filename = "save_bin_test"; - let data = vec![0u64, 1u64, 2u64]; - let num_pts = data.len(); - let dims = 1; - let bytes_written = save_bin_u64(filename, &data, num_pts, dims, 0).unwrap(); - assert_eq!(bytes_written, 32); - - let mut file = File::open(filename).unwrap(); - let mut buffer = vec![]; - - let npts_read = file.read_i32::().unwrap() as usize; - let dims_read = file.read_i32::().unwrap() as usize; - - file.read_to_end(&mut buffer).unwrap(); - let data_read: Vec = buffer - .chunks_exact(8) - .map(|b| u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])) - .collect(); - - std::fs::remove_file(filename).unwrap(); - - assert_eq!(num_pts, npts_read); - assert_eq!(dims, dims_read); - assert_eq!(data, data_read); - } - - #[test] - fn load_bin_test() { - let file_name = "load_bin_test"; - let data = vec![0u64, 1u64, 2u64]; - let num_pts = data.len(); - let dims = 1; - let bytes_written = save_bin_u64(file_name, &data, num_pts, dims, 0).unwrap(); - assert_eq!(bytes_written, 32); - - let (load_data, load_num_pts, load_dims) = load_bin::(file_name, 0).unwrap(); - assert_eq!(load_num_pts, num_pts); - assert_eq!(load_dims, dims); - assert_eq!(load_data, data); - std::fs::remove_file(file_name).unwrap(); - } - - #[test] - fn load_bin_offset_test() { - let offset:usize = 32; - let file_name = "load_bin_offset_test"; - let data = vec![0u64, 1u64, 2u64]; - let num_pts = data.len(); - let dims = 1; - let bytes_written = save_bin_u64(file_name, &data, num_pts, dims, offset).unwrap(); - assert_eq!(bytes_written, 32); - - let (load_data, load_num_pts, load_dims) = load_bin::(file_name, offset).unwrap(); - assert_eq!(load_num_pts, num_pts); - assert_eq!(load_dims, dims); - assert_eq!(load_data, data); - std::fs::remove_file(file_name).unwrap(); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/hashset_u32.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/hashset_u32.rs deleted file mode 100644 index 15db687..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/hashset_u32.rs +++ /dev/null @@ -1,46 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use hashbrown::HashSet; -use std::{hash::BuildHasherDefault, ops::{Deref, DerefMut}}; -use fxhash::FxHasher; - -lazy_static::lazy_static! { - /// Singleton hasher. - static ref HASHER: BuildHasherDefault = { - BuildHasherDefault::::default() - }; -} - -pub struct HashSetForU32 { - hashset: HashSet::>, -} - -impl HashSetForU32 { - pub fn with_capacity(capacity: usize) -> HashSetForU32 { - let hashset = HashSet::>::with_capacity_and_hasher(capacity, HASHER.clone()); - HashSetForU32 { - hashset - } - } -} - -impl Deref for HashSetForU32 { - type Target = HashSet::>; - - fn deref(&self) -> &Self::Target { - &self.hashset - } -} - -impl DerefMut for HashSetForU32 { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.hashset - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/kmeans.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/kmeans.rs deleted file mode 100644 index d1edffa..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/kmeans.rs +++ /dev/null @@ -1,430 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -use rand::{distributions::Uniform, prelude::Distribution, thread_rng}; -use rayon::prelude::*; -use std::cmp::min; - -use crate::common::ANNResult; -use crate::utils::math_util::{calc_distance, compute_closest_centers, compute_vecs_l2sq}; - -/// Run Lloyds one iteration -/// Given data in row-major num_points * dim, and centers in row-major -/// num_centers * dim and squared lengths of ata points, output the closest -/// center to each data point, update centers, and also return inverted index. -/// If closest_centers == NULL, will allocate memory and return. -/// Similarly, if closest_docs == NULL, will allocate memory and return. -#[allow(clippy::too_many_arguments)] -fn lloyds_iter( - data: &[f32], - num_points: usize, - dim: usize, - centers: &mut [f32], - num_centers: usize, - docs_l2sq: &[f32], - mut closest_docs: &mut Vec>, - closest_center: &mut [u32], -) -> ANNResult { - let compute_residual = true; - - closest_docs.iter_mut().for_each(|doc| doc.clear()); - - compute_closest_centers( - data, - num_points, - dim, - centers, - num_centers, - 1, - closest_center, - Some(&mut closest_docs), - Some(docs_l2sq), - )?; - - centers.fill(0.0); - - centers - .par_chunks_mut(dim) - .enumerate() - .for_each(|(c, center)| { - let mut cluster_sum = vec![0.0; dim]; - for &doc_index in &closest_docs[c] { - let current = &data[doc_index * dim..(doc_index + 1) * dim]; - for (j, current_val) in current.iter().enumerate() { - cluster_sum[j] += *current_val as f64; - } - } - if !closest_docs[c].is_empty() { - for (i, sum_val) in cluster_sum.iter().enumerate() { - center[i] = (*sum_val / closest_docs[c].len() as f64) as f32; - } - } - }); - - let mut residual = 0.0; - if compute_residual { - let buf_pad: usize = 32; - let chunk_size: usize = 2 * 8192; - let nchunks = - num_points / chunk_size + (if num_points % chunk_size == 0 { 0 } else { 1 } as usize); - - let mut residuals: Vec = vec![0.0; nchunks * buf_pad]; - - residuals - .par_iter_mut() - .enumerate() - .for_each(|(chunk, res)| { - for d in (chunk * chunk_size)..min(num_points, (chunk + 1) * chunk_size) { - *res += calc_distance( - &data[d * dim..(d + 1) * dim], - ¢ers[closest_center[d] as usize * dim..], - dim, - ); - } - }); - - for chunk in 0..nchunks { - residual += residuals[chunk * buf_pad]; - } - } - - Ok(residual) -} - -/// Run Lloyds until max_reps or stopping criterion -/// If you pass NULL for closest_docs and closest_center, it will NOT return -/// the results, else it will assume appropriate allocation as closest_docs = -/// new vec [num_centers], and closest_center = new size_t[num_points] -/// Final centers are output in centers as row-major num_centers * dim. -fn run_lloyds( - data: &[f32], - num_points: usize, - dim: usize, - centers: &mut [f32], - num_centers: usize, - max_reps: usize, -) -> ANNResult<(Vec>, Vec, f32)> { - let mut residual = f32::MAX; - - let mut closest_docs = vec![Vec::new(); num_centers]; - let mut closest_center = vec![0; num_points]; - - let mut docs_l2sq = vec![0.0; num_points]; - compute_vecs_l2sq(&mut docs_l2sq, data, num_points, dim); - - let mut old_residual; - - for i in 0..max_reps { - old_residual = residual; - - residual = lloyds_iter( - data, - num_points, - dim, - centers, - num_centers, - &docs_l2sq, - &mut closest_docs, - &mut closest_center, - )?; - - if (i != 0 && (old_residual - residual) / residual < 0.00001) || (residual < f32::EPSILON) { - println!( - "Residuals unchanged: {} becomes {}. Early termination.", - old_residual, residual - ); - break; - } - } - - Ok((closest_docs, closest_center, residual)) -} - -/// Assume memory allocated for pivot_data as new float[num_centers * dim] -/// and select randomly num_centers points as pivots -fn selecting_pivots( - data: &[f32], - num_points: usize, - dim: usize, - pivot_data: &mut [f32], - num_centers: usize, -) { - let mut picked = Vec::new(); - let mut rng = thread_rng(); - let distribution = Uniform::from(0..num_points); - - for j in 0..num_centers { - let mut tmp_pivot = distribution.sample(&mut rng); - while picked.contains(&tmp_pivot) { - tmp_pivot = distribution.sample(&mut rng); - } - picked.push(tmp_pivot); - let data_offset = tmp_pivot * dim; - let pivot_offset = j * dim; - pivot_data[pivot_offset..pivot_offset + dim] - .copy_from_slice(&data[data_offset..data_offset + dim]); - } -} - -/// Select pivots in k-means++ algorithm -/// Points that are farther away from the already chosen centroids -/// have a higher probability of being selected as the next centroid. -/// The k-means++ algorithm helps avoid poor initial centroid -/// placement that can result in suboptimal clustering. -fn k_meanspp_selecting_pivots( - data: &[f32], - num_points: usize, - dim: usize, - pivot_data: &mut [f32], - num_centers: usize, -) { - if num_points > (1 << 23) { - println!("ERROR: n_pts {} currently not supported for k-means++, maximum is 8388608. Falling back to random pivot selection.", num_points); - selecting_pivots(data, num_points, dim, pivot_data, num_centers); - return; - } - - let mut picked: Vec = Vec::new(); - let mut rng = thread_rng(); - let real_distribution = Uniform::from(0.0..1.0); - let int_distribution = Uniform::from(0..num_points); - - let init_id = int_distribution.sample(&mut rng); - let mut num_picked = 1; - - picked.push(init_id); - let init_data_offset = init_id * dim; - pivot_data[0..dim].copy_from_slice(&data[init_data_offset..init_data_offset + dim]); - - let mut dist = vec![0.0; num_points]; - - dist.par_iter_mut().enumerate().for_each(|(i, dist_i)| { - *dist_i = calc_distance( - &data[i * dim..(i + 1) * dim], - &data[init_id * dim..(init_id + 1) * dim], - dim, - ); - }); - - let mut dart_val: f64; - let mut tmp_pivot = 0; - let mut sum_flag = false; - - while num_picked < num_centers { - dart_val = real_distribution.sample(&mut rng); - - let mut sum: f64 = 0.0; - for item in dist.iter().take(num_points) { - sum += *item as f64; - } - if sum == 0.0 { - sum_flag = true; - } - - dart_val *= sum; - - let mut prefix_sum: f64 = 0.0; - for (i, pivot) in dist.iter().enumerate().take(num_points) { - tmp_pivot = i; - if dart_val >= prefix_sum && dart_val < (prefix_sum + *pivot as f64) { - break; - } - - prefix_sum += *pivot as f64; - } - - if picked.contains(&tmp_pivot) && !sum_flag { - continue; - } - - picked.push(tmp_pivot); - let pivot_offset = num_picked * dim; - let data_offset = tmp_pivot * dim; - pivot_data[pivot_offset..pivot_offset + dim] - .copy_from_slice(&data[data_offset..data_offset + dim]); - - dist.par_iter_mut().enumerate().for_each(|(i, dist_i)| { - *dist_i = (*dist_i).min(calc_distance( - &data[i * dim..(i + 1) * dim], - &data[tmp_pivot * dim..(tmp_pivot + 1) * dim], - dim, - )); - }); - - num_picked += 1; - } -} - -/// k-means algorithm interface -pub fn k_means_clustering( - data: &[f32], - num_points: usize, - dim: usize, - centers: &mut [f32], - num_centers: usize, - max_reps: usize, -) -> ANNResult<(Vec>, Vec, f32)> { - k_meanspp_selecting_pivots(data, num_points, dim, centers, num_centers); - let (closest_docs, closest_center, residual) = - run_lloyds(data, num_points, dim, centers, num_centers, max_reps)?; - Ok((closest_docs, closest_center, residual)) -} - -#[cfg(test)] -mod kmeans_test { - use super::*; - use approx::assert_relative_eq; - use rand::Rng; - - #[test] - fn lloyds_iter_test() { - let dim = 2; - let num_points = 10; - let num_centers = 3; - - let data: Vec = (1..=num_points * dim).map(|x| x as f32).collect(); - let mut centers = [1.0, 2.0, 7.0, 8.0, 19.0, 20.0]; - - let mut closest_docs: Vec> = vec![vec![]; num_centers]; - let mut closest_center: Vec = vec![0; num_points]; - let docs_l2sq: Vec = data - .chunks(dim) - .map(|chunk| chunk.iter().map(|val| val.powi(2)).sum()) - .collect(); - - let residual = lloyds_iter( - &data, - num_points, - dim, - &mut centers, - num_centers, - &docs_l2sq, - &mut closest_docs, - &mut closest_center, - ) - .unwrap(); - - let expected_centers: [f32; 6] = [2.0, 3.0, 9.0, 10.0, 17.0, 18.0]; - let expected_closest_docs: Vec> = - vec![vec![0, 1], vec![2, 3, 4, 5, 6], vec![7, 8, 9]]; - let expected_closest_center: [u32; 10] = [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]; - let expected_residual: f32 = 100.0; - - // sort data for assert - centers.sort_by(|a, b| a.partial_cmp(b).unwrap()); - for inner_vec in &mut closest_docs { - inner_vec.sort(); - } - closest_center.sort_by(|a, b| a.partial_cmp(b).unwrap()); - - assert_eq!(centers, expected_centers); - assert_eq!(closest_docs, expected_closest_docs); - assert_eq!(closest_center, expected_closest_center); - assert_relative_eq!(residual, expected_residual, epsilon = 1.0e-6_f32); - } - - #[test] - fn run_lloyds_test() { - let dim = 2; - let num_points = 10; - let num_centers = 3; - let max_reps = 5; - - let data: Vec = (1..=num_points * dim).map(|x| x as f32).collect(); - let mut centers = [1.0, 2.0, 7.0, 8.0, 19.0, 20.0]; - - let (mut closest_docs, mut closest_center, residual) = - run_lloyds(&data, num_points, dim, &mut centers, num_centers, max_reps).unwrap(); - - let expected_centers: [f32; 6] = [3.0, 4.0, 10.0, 11.0, 17.0, 18.0]; - let expected_closest_docs: Vec> = - vec![vec![0, 1, 2], vec![3, 4, 5, 6], vec![7, 8, 9]]; - let expected_closest_center: [u32; 10] = [0, 0, 0, 1, 1, 1, 1, 2, 2, 2]; - let expected_residual: f32 = 72.0; - - // sort data for assert - centers.sort_by(|a, b| a.partial_cmp(b).unwrap()); - for inner_vec in &mut closest_docs { - inner_vec.sort(); - } - closest_center.sort_by(|a, b| a.partial_cmp(b).unwrap()); - - assert_eq!(centers, expected_centers); - assert_eq!(closest_docs, expected_closest_docs); - assert_eq!(closest_center, expected_closest_center); - assert_relative_eq!(residual, expected_residual, epsilon = 1.0e-6_f32); - } - - #[test] - fn selecting_pivots_test() { - let dim = 2; - let num_points = 10; - let num_centers = 3; - - // Generate some random data points - let mut rng = rand::thread_rng(); - let data: Vec = (0..num_points * dim).map(|_| rng.gen()).collect(); - - let mut pivot_data = vec![0.0; num_centers * dim]; - - selecting_pivots(&data, num_points, dim, &mut pivot_data, num_centers); - - // Verify that each pivot point corresponds to a point in the data - for i in 0..num_centers { - let pivot_offset = i * dim; - let pivot = &pivot_data[pivot_offset..(pivot_offset + dim)]; - - // Make sure the pivot is found in the data - let mut found = false; - for j in 0..num_points { - let data_offset = j * dim; - let point = &data[data_offset..(data_offset + dim)]; - - if pivot == point { - found = true; - break; - } - } - assert!(found, "Pivot not found in data"); - } - } - - #[test] - fn k_meanspp_selecting_pivots_test() { - let dim = 2; - let num_points = 10; - let num_centers = 3; - - // Generate some random data points - let mut rng = rand::thread_rng(); - let data: Vec = (0..num_points * dim).map(|_| rng.gen()).collect(); - - let mut pivot_data = vec![0.0; num_centers * dim]; - - k_meanspp_selecting_pivots(&data, num_points, dim, &mut pivot_data, num_centers); - - // Verify that each pivot point corresponds to a point in the data - for i in 0..num_centers { - let pivot_offset = i * dim; - let pivot = &pivot_data[pivot_offset..pivot_offset + dim]; - - // Make sure the pivot is found in the data - let mut found = false; - for j in 0..num_points { - let data_offset = j * dim; - let point = &data[data_offset..data_offset + dim]; - - if pivot == point { - found = true; - break; - } - } - assert!(found, "Pivot not found in data"); - } - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/math_util.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/math_util.rs deleted file mode 100644 index ef30c76..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/math_util.rs +++ /dev/null @@ -1,481 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Aligned allocator - -extern crate cblas; -extern crate openblas_src; - -use cblas::{sgemm, snrm2, Layout, Transpose}; -use rayon::prelude::*; -use std::{ - cmp::{min, Ordering}, - collections::BinaryHeap, - sync::{Arc, Mutex}, -}; - -use crate::common::{ANNError, ANNResult}; - -struct PivotContainer { - piv_id: usize, - piv_dist: f32, -} - -impl PartialOrd for PivotContainer { - fn partial_cmp(&self, other: &Self) -> Option { - other.piv_dist.partial_cmp(&self.piv_dist) - } -} - -impl Ord for PivotContainer { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - // Treat NaN as less than all other values. - // piv_dist should never be NaN. - self.partial_cmp(other).unwrap_or(Ordering::Less) - } -} - -impl PartialEq for PivotContainer { - fn eq(&self, other: &Self) -> bool { - self.piv_dist == other.piv_dist - } -} - -impl Eq for PivotContainer {} - -/// Calculate the Euclidean distance between two vectors -pub fn calc_distance(vec_1: &[f32], vec_2: &[f32], dim: usize) -> f32 { - let mut dist = 0.0; - for j in 0..dim { - let diff = vec_1[j] - vec_2[j]; - dist += diff * diff; - } - dist -} - -/// Compute L2-squared norms of data stored in row-major num_points * dim, -/// need to be pre-allocated -pub fn compute_vecs_l2sq(vecs_l2sq: &mut [f32], data: &[f32], num_points: usize, dim: usize) { - assert_eq!(vecs_l2sq.len(), num_points); - - vecs_l2sq - .par_iter_mut() - .enumerate() - .for_each(|(n_iter, vec_l2sq)| { - let slice = &data[n_iter * dim..(n_iter + 1) * dim]; - let norm = unsafe { snrm2(dim as i32, slice, 1) }; - *vec_l2sq = norm * norm; - }); -} - -/// Calculate k closest centers to data of num_points * dim (row-major) -/// Centers is num_centers * dim (row-major) -/// data_l2sq has pre-computed squared norms of data -/// centers_l2sq has pre-computed squared norms of centers -/// Pre-allocated center_index will contain id of nearest center -/// Pre-allocated dist_matrix should be num_points * num_centers and contain squared distances -/// Default value of k is 1 -/// Ideally used only by compute_closest_centers -#[allow(clippy::too_many_arguments)] -pub fn compute_closest_centers_in_block( - data: &[f32], - num_points: usize, - dim: usize, - centers: &[f32], - num_centers: usize, - docs_l2sq: &[f32], - centers_l2sq: &[f32], - center_index: &mut [u32], - dist_matrix: &mut [f32], - k: usize, -) -> ANNResult<()> { - if k > num_centers { - return Err(ANNError::log_index_error(format!( - "ERROR: k ({}) > num_centers({})", - k, num_centers - ))); - } - - let ones_a: Vec = vec![1.0; num_centers]; - let ones_b: Vec = vec![1.0; num_points]; - - unsafe { - sgemm( - Layout::RowMajor, - Transpose::None, - Transpose::Ordinary, - num_points as i32, - num_centers as i32, - 1, - 1.0, - docs_l2sq, - 1, - &ones_a, - 1, - 0.0, - dist_matrix, - num_centers as i32, - ); - } - - unsafe { - sgemm( - Layout::RowMajor, - Transpose::None, - Transpose::Ordinary, - num_points as i32, - num_centers as i32, - 1, - 1.0, - &ones_b, - 1, - centers_l2sq, - 1, - 1.0, - dist_matrix, - num_centers as i32, - ); - } - - unsafe { - sgemm( - Layout::RowMajor, - Transpose::None, - Transpose::Ordinary, - num_points as i32, - num_centers as i32, - dim as i32, - -2.0, - data, - dim as i32, - centers, - dim as i32, - 1.0, - dist_matrix, - num_centers as i32, - ); - } - - if k == 1 { - center_index - .par_iter_mut() - .enumerate() - .for_each(|(i, center_idx)| { - let mut min = f32::MAX; - let current = &dist_matrix[i * num_centers..(i + 1) * num_centers]; - let mut min_idx = 0; - for (j, &distance) in current.iter().enumerate() { - if distance < min { - min = distance; - min_idx = j; - } - } - *center_idx = min_idx as u32; - }); - } else { - center_index - .par_chunks_mut(k) - .enumerate() - .for_each(|(i, center_chunk)| { - let current = &dist_matrix[i * num_centers..(i + 1) * num_centers]; - let mut top_k_queue = BinaryHeap::new(); - for (j, &distance) in current.iter().enumerate() { - let this_piv = PivotContainer { - piv_id: j, - piv_dist: distance, - }; - if top_k_queue.len() < k { - top_k_queue.push(this_piv); - } else { - // Safe unwrap, top_k_queue is not empty - #[allow(clippy::unwrap_used)] - let mut top = top_k_queue.peek_mut().unwrap(); - if this_piv.piv_dist < top.piv_dist { - *top = this_piv; - } - } - } - for (_j, center_idx) in center_chunk.iter_mut().enumerate() { - if let Some(this_piv) = top_k_queue.pop() { - *center_idx = this_piv.piv_id as u32; - } else { - break; - } - } - }); - } - - Ok(()) -} - -/// Given data in num_points * new_dim row major -/// Pivots stored in full_pivot_data as num_centers * new_dim row major -/// Calculate the k closest pivot for each point and store it in vector -/// closest_centers_ivf (row major, num_points*k) (which needs to be allocated -/// outside) Additionally, if inverted index is not null (and pre-allocated), -/// it will return inverted index for each center, assuming each of the inverted -/// indices is an empty vector. Additionally, if pts_norms_squared is not null, -/// then it will assume that point norms are pre-computed and use those values -#[allow(clippy::too_many_arguments)] -pub fn compute_closest_centers( - data: &[f32], - num_points: usize, - dim: usize, - pivot_data: &[f32], - num_centers: usize, - k: usize, - closest_centers_ivf: &mut [u32], - mut inverted_index: Option<&mut Vec>>, - pts_norms_squared: Option<&[f32]>, -) -> ANNResult<()> { - if k > num_centers { - return Err(ANNError::log_index_error(format!( - "ERROR: k ({}) > num_centers({})", - k, num_centers - ))); - } - - let _is_norm_given_for_pts = pts_norms_squared.is_some(); - - let mut pivs_norms_squared = vec![0.0; num_centers]; - - let mut pts_norms_squared = if let Some(pts_norms) = pts_norms_squared { - pts_norms.to_vec() - } else { - let mut norms_squared = vec![0.0; num_points]; - compute_vecs_l2sq(&mut norms_squared, data, num_points, dim); - norms_squared - }; - - compute_vecs_l2sq(&mut pivs_norms_squared, pivot_data, num_centers, dim); - - let par_block_size = num_points; - let n_blocks = if num_points % par_block_size == 0 { - num_points / par_block_size - } else { - num_points / par_block_size + 1 - }; - - let mut closest_centers = vec![0u32; par_block_size * k]; - let mut distance_matrix = vec![0.0; num_centers * par_block_size]; - - for cur_blk in 0..n_blocks { - let data_cur_blk = &data[cur_blk * par_block_size * dim..]; - let num_pts_blk = min(par_block_size, num_points - cur_blk * par_block_size); - let pts_norms_blk = &mut pts_norms_squared[cur_blk * par_block_size..]; - - compute_closest_centers_in_block( - data_cur_blk, - num_pts_blk, - dim, - pivot_data, - num_centers, - pts_norms_blk, - &pivs_norms_squared, - &mut closest_centers, - &mut distance_matrix, - k, - )?; - - closest_centers_ivf.clone_from_slice(&closest_centers); - - if let Some(inverted_index_inner) = inverted_index.as_mut() { - let inverted_index_arc = Arc::new(Mutex::new(inverted_index_inner)); - - (0..num_points) - .into_par_iter() - .try_for_each(|j| -> ANNResult<()> { - let this_center_id = closest_centers[j] as usize; - let mut guard = inverted_index_arc.lock().map_err(|err| { - ANNError::log_index_error(format!( - "PoisonError: Lock poisoned when acquiring inverted_index_arc, err={}", - err - )) - })?; - guard[this_center_id].push(j); - - Ok(()) - })?; - } - } - - Ok(()) -} - -/// If to_subtract is true, will subtract nearest center from each row. -/// Else will add. -/// Output will be in data_load itself. -/// Nearest centers need to be provided in closest_centers. -pub fn process_residuals( - data_load: &mut [f32], - num_points: usize, - dim: usize, - cur_pivot_data: &[f32], - num_centers: usize, - closest_centers: &[u32], - to_subtract: bool, -) { - println!( - "Processing residuals of {} points in {} dimensions using {} centers", - num_points, dim, num_centers - ); - - data_load - .par_chunks_mut(dim) - .enumerate() - .for_each(|(n_iter, chunk)| { - let cur_pivot_index = closest_centers[n_iter] as usize * dim; - for d_iter in 0..dim { - if to_subtract { - chunk[d_iter] -= cur_pivot_data[cur_pivot_index + d_iter]; - } else { - chunk[d_iter] += cur_pivot_data[cur_pivot_index + d_iter]; - } - } - }); -} - -#[cfg(test)] -mod math_util_test { - use super::*; - use approx::assert_abs_diff_eq; - - #[test] - fn calc_distance_test() { - let vec1 = vec![1.0, 2.0, 3.0]; - let vec2 = vec![4.0, 5.0, 6.0]; - let dim = vec1.len(); - - let dist = calc_distance(&vec1, &vec2, dim); - - let expected = 27.0; - - assert_eq!(dist, expected); - } - - #[test] - fn compute_vecs_l2sq_test() { - let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; - let num_points = 2; - let dim = 3; - let mut vecs_l2sq = vec![0.0; num_points]; - - compute_vecs_l2sq(&mut vecs_l2sq, &data, num_points, dim); - - let expected = vec![14.0, 77.0]; - - assert_eq!(vecs_l2sq.len(), num_points); - assert_abs_diff_eq!(vecs_l2sq[0], expected[0], epsilon = 1e-6); - assert_abs_diff_eq!(vecs_l2sq[1], expected[1], epsilon = 1e-6); - } - - #[test] - fn compute_closest_centers_in_block_test() { - let num_points = 10; - let dim = 5; - let num_centers = 3; - let data = vec![ - 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, - 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, - 31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0, 40.0, 41.0, 42.0, 43.0, 44.0, - 45.0, 46.0, 47.0, 48.0, 49.0, 50.0, - ]; - let centers = vec![ - 1.0, 2.0, 3.0, 4.0, 5.0, 21.0, 22.0, 23.0, 24.0, 25.0, 31.0, 32.0, 33.0, 34.0, 35.0, - ]; - let mut docs_l2sq = vec![0.0; num_points]; - compute_vecs_l2sq(&mut docs_l2sq, &data, num_points, dim); - let mut centers_l2sq = vec![0.0; num_centers]; - compute_vecs_l2sq(&mut centers_l2sq, ¢ers, num_centers, dim); - let mut center_index = vec![0; num_points]; - let mut dist_matrix = vec![0.0; num_points * num_centers]; - let k = 1; - - compute_closest_centers_in_block( - &data, - num_points, - dim, - ¢ers, - num_centers, - &docs_l2sq, - ¢ers_l2sq, - &mut center_index, - &mut dist_matrix, - k, - ) - .unwrap(); - - assert_eq!(center_index.len(), num_points); - let expected_center_index = vec![0, 0, 0, 1, 1, 1, 2, 2, 2, 2]; - assert_abs_diff_eq!(*center_index, expected_center_index); - - assert_eq!(dist_matrix.len(), num_points * num_centers); - let expected_dist_matrix = vec![ - 0.0, 2000.0, 4500.0, 125.0, 1125.0, 3125.0, 500.0, 500.0, 2000.0, 1125.0, 125.0, - 1125.0, 2000.0, 0.0, 500.0, 3125.0, 125.0, 125.0, 4500.0, 500.0, 0.0, 6125.0, 1125.0, - 125.0, 8000.0, 2000.0, 500.0, 10125.0, 3125.0, 1125.0, - ]; - assert_abs_diff_eq!(*dist_matrix, expected_dist_matrix, epsilon = 1e-2); - } - - #[test] - fn test_compute_closest_centers() { - let num_points = 4; - let dim = 3; - let num_centers = 2; - let mut data = vec![ - 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, - ]; - let pivot_data = vec![1.0, 2.0, 3.0, 10.0, 11.0, 12.0]; - let k = 1; - - let mut closest_centers_ivf = vec![0u32; num_points * k]; - let mut inverted_index: Vec> = vec![vec![], vec![]]; - - compute_closest_centers( - &data, - num_points, - dim, - &pivot_data, - num_centers, - k, - &mut closest_centers_ivf, - Some(&mut inverted_index), - None, - ) - .unwrap(); - - assert_eq!(closest_centers_ivf, vec![0, 0, 1, 1]); - - for vec in inverted_index.iter_mut() { - vec.sort_unstable(); - } - assert_eq!(inverted_index, vec![vec![0, 1], vec![2, 3]]); - } - - #[test] - fn process_residuals_test() { - let mut data_load = vec![1.0, 2.0, 3.0, 4.0]; - let num_points = 2; - let dim = 2; - let cur_pivot_data = vec![0.5, 1.5, 2.5, 3.5]; - let num_centers = 2; - let closest_centers = vec![0, 1]; - let to_subtract = true; - - process_residuals( - &mut data_load, - num_points, - dim, - &cur_pivot_data, - num_centers, - &closest_centers, - to_subtract, - ); - - assert_eq!(data_load, vec![0.5, 0.5, 0.5, 0.5]); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/mod.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/mod.rs deleted file mode 100644 index df174f8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/mod.rs +++ /dev/null @@ -1,34 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -pub mod file_util; -pub use file_util::*; - -#[allow(clippy::module_inception)] -pub mod utils; -pub use utils::*; - -pub mod bit_vec_extension; -pub use bit_vec_extension::*; - -pub mod rayon_util; -pub use rayon_util::*; - -pub mod timer; -pub use timer::*; - -pub mod cached_reader; -pub use cached_reader::*; - -pub mod cached_writer; -pub use cached_writer::*; - -pub mod partition; -pub use partition::*; - -pub mod math_util; -pub use math_util::*; - -pub mod kmeans; -pub use kmeans::*; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/partition.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/partition.rs deleted file mode 100644 index dbe6862..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/partition.rs +++ /dev/null @@ -1,151 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::mem; -use std::{fs::File, path::Path}; -use std::io::{Write, Seek, SeekFrom}; -use rand::distributions::{Distribution, Uniform}; - -use crate::common::ANNResult; - -use super::CachedReader; - -/// streams data from the file, and samples each vector with probability p_val -/// and returns a matrix of size slice_size* ndims as floating point type. -/// the slice_size and ndims are set inside the function. -/// # Arguments -/// * `file_name` - filename where the data is -/// * `p_val` - possibility to sample data -/// * `sampled_vectors` - sampled vector chose by p_val possibility -/// * `slice_size` - how many sampled data return -/// * `dim` - each sample data dimension -pub fn gen_random_slice>(data_file: &str, mut p_val: f64) -> ANNResult<(Vec, usize, usize)> { - let read_blk_size = 64 * 1024 * 1024; - let mut reader = CachedReader::new(data_file, read_blk_size)?; - - let npts = reader.read_u32()? as usize; - let dim = reader.read_u32()? as usize; - let mut sampled_vectors: Vec = Vec::new(); - let mut slice_size = 0; - p_val = if p_val < 1f64 { p_val } else { 1f64 }; - - let mut generator = rand::thread_rng(); - let distribution = Uniform::from(0.0..1.0); - - for _ in 0..npts { - let mut cur_vector_bytes = vec![0u8; dim * mem::size_of::()]; - reader.read(&mut cur_vector_bytes)?; - let random_value = distribution.sample(&mut generator); - if random_value < p_val { - let ptr = cur_vector_bytes.as_ptr() as *const T; - let cur_vector_t = unsafe { std::slice::from_raw_parts(ptr, dim) }; - sampled_vectors.extend(cur_vector_t.iter().map(|&t| t.into())); - slice_size += 1; - } - } - - Ok((sampled_vectors, slice_size, dim)) -} - -/// Generate random sample data and write into output_file -pub fn gen_sample_data(data_file: &str, output_file: &str, sampling_rate: f64) -> ANNResult<()> { - let read_blk_size = 64 * 1024 * 1024; - let mut reader = CachedReader::new(data_file, read_blk_size)?; - - let sample_data_path = format!("{}_data.bin", output_file); - let sample_ids_path = format!("{}_ids.bin", output_file); - let mut sample_data_writer = File::create(Path::new(&sample_data_path))?; - let mut sample_id_writer = File::create(Path::new(&sample_ids_path))?; - - let mut num_sampled_pts = 0u32; - let one_const = 1u32; - let mut generator = rand::thread_rng(); - let distribution = Uniform::from(0.0..1.0); - - let npts_u32 = reader.read_u32()?; - let dim_u32 = reader.read_u32()?; - let dim = dim_u32 as usize; - sample_data_writer.write_all(&num_sampled_pts.to_le_bytes())?; - sample_data_writer.write_all(&dim_u32.to_le_bytes())?; - sample_id_writer.write_all(&num_sampled_pts.to_le_bytes())?; - sample_id_writer.write_all(&one_const.to_le_bytes())?; - - for id in 0..npts_u32 { - let mut cur_row_bytes = vec![0u8; dim * mem::size_of::()]; - reader.read(&mut cur_row_bytes)?; - let random_value = distribution.sample(&mut generator); - if random_value < sampling_rate { - sample_data_writer.write_all(&cur_row_bytes)?; - sample_id_writer.write_all(&id.to_le_bytes())?; - num_sampled_pts += 1; - } - } - - sample_data_writer.seek(SeekFrom::Start(0))?; - sample_data_writer.write_all(&num_sampled_pts.to_le_bytes())?; - sample_id_writer.seek(SeekFrom::Start(0))?; - sample_id_writer.write_all(&num_sampled_pts.to_le_bytes())?; - println!("Wrote {} points to sample file: {}", num_sampled_pts, sample_data_path); - - Ok(()) -} - -#[cfg(test)] -mod partition_test { - use std::{fs, io::Read}; - use byteorder::{ReadBytesExt, LittleEndian}; - - use crate::utils::file_exists; - - use super::*; - - #[test] - fn gen_sample_data_test() { - let file_name = "gen_sample_data_test.bin"; - //npoints=2, dim=8 - let data: [u8; 72] = [2, 0, 0, 0, 8, 0, 0, 0, - 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x40, 0x40, 0x00, 0x00, 0x80, 0x40, - 0x00, 0x00, 0xa0, 0x40, 0x00, 0x00, 0xc0, 0x40, 0x00, 0x00, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x41, - 0x00, 0x00, 0x10, 0x41, 0x00, 0x00, 0x20, 0x41, 0x00, 0x00, 0x30, 0x41, 0x00, 0x00, 0x40, 0x41, - 0x00, 0x00, 0x50, 0x41, 0x00, 0x00, 0x60, 0x41, 0x00, 0x00, 0x70, 0x41, 0x00, 0x00, 0x80, 0x41]; - std::fs::write(file_name, data).expect("Failed to write sample file"); - - let sample_file_prefix = file_name.to_string() + "_sample"; - gen_sample_data::(file_name, sample_file_prefix.as_str(), 1f64).unwrap(); - - let sample_data_path = format!("{}_data.bin", sample_file_prefix); - let sample_ids_path = format!("{}_ids.bin", sample_file_prefix); - assert!(file_exists(sample_data_path.as_str())); - assert!(file_exists(sample_ids_path.as_str())); - - let mut data_file_reader = File::open(sample_data_path.as_str()).unwrap(); - let mut ids_file_reader = File::open(sample_ids_path.as_str()).unwrap(); - - let mut num_sampled_pts = data_file_reader.read_u32::().unwrap(); - assert_eq!(num_sampled_pts, 2); - num_sampled_pts = ids_file_reader.read_u32::().unwrap(); - assert_eq!(num_sampled_pts, 2); - - let dim = data_file_reader.read_u32::().unwrap() as usize; - assert_eq!(dim, 8); - assert_eq!(ids_file_reader.read_u32::().unwrap(), 1); - - let mut start = 8; - for i in 0..num_sampled_pts { - let mut data_bytes = vec![0u8; dim * 4]; - data_file_reader.read_exact(&mut data_bytes).unwrap(); - assert_eq!(data_bytes, data[start..start + dim * 4]); - - let id = ids_file_reader.read_u32::().unwrap(); - assert_eq!(id, i); - - start += dim * 4; - } - - fs::remove_file(file_name).expect("Failed to delete file"); - fs::remove_file(sample_data_path.as_str()).expect("Failed to delete file"); - fs::remove_file(sample_ids_path.as_str()).expect("Failed to delete file"); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/rayon_util.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/rayon_util.rs deleted file mode 100644 index f8174ee..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/rayon_util.rs +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::ops::Range; -use rayon::prelude::{IntoParallelIterator, ParallelIterator}; - -use crate::common::ANNResult; - -/// based on thread_num, execute the task in parallel using Rayon or serial -#[inline] -pub fn execute_with_rayon(range: Range, num_threads: u32, f: F) -> ANNResult<()> -where F: Fn(usize) -> ANNResult<()> + Sync + Send + Copy -{ - if num_threads == 1 { - for i in range { - f(i)?; - } - Ok(()) - } else { - range.into_par_iter().try_for_each(f) - } -} - -/// set the thread count of Rayon, otherwise it will use threads as many as logical cores. -#[inline] -pub fn set_rayon_num_threads(num_threads: u32) { - std::env::set_var( - "RAYON_NUM_THREADS", - num_threads.to_string(), - ); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/timer.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/timer.rs deleted file mode 100644 index 2f4b38b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/timer.rs +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use platform::*; -use std::time::{Duration, Instant}; - -#[derive(Clone)] -pub struct Timer { - check_point: Instant, - pid: Option, - cycles: Option, -} - -impl Default for Timer { - fn default() -> Self { - Self::new() - } -} - -impl Timer { - pub fn new() -> Timer { - let pid = get_process_handle(); - let cycles = get_process_cycle_time(pid); - Timer { - check_point: Instant::now(), - pid, - cycles, - } - } - - pub fn reset(&mut self) { - self.check_point = Instant::now(); - self.cycles = get_process_cycle_time(self.pid); - } - - pub fn elapsed(&self) -> Duration { - Instant::now().duration_since(self.check_point) - } - - pub fn elapsed_seconds(&self) -> f64 { - self.elapsed().as_secs_f64() - } - - pub fn elapsed_gcycles(&self) -> f32 { - let cur_cycles = get_process_cycle_time(self.pid); - if let (Some(cur_cycles), Some(cycles)) = (cur_cycles, self.cycles) { - let spent_cycles = - ((cur_cycles - cycles) as f64 * 1.0f64) / (1024 * 1024 * 1024) as f64; - return spent_cycles as f32; - } - - 0.0 - } - - pub fn elapsed_seconds_for_step(&self, step: &str) -> String { - format!( - "Time for {}: {:.3} seconds, {:.3}B cycles", - step, - self.elapsed_seconds(), - self.elapsed_gcycles() - ) - } -} - -#[cfg(test)] -mod timer_tests { - use super::*; - use std::{thread, time}; - - #[test] - fn test_new() { - let timer = Timer::new(); - assert!(timer.check_point.elapsed().as_secs() < 1); - if cfg!(windows) { - assert!(timer.pid.is_some()); - assert!(timer.cycles.is_some()); - } - else { - assert!(timer.pid.is_none()); - assert!(timer.cycles.is_none()); - } - } - - #[test] - fn test_reset() { - let mut timer = Timer::new(); - thread::sleep(time::Duration::from_millis(100)); - timer.reset(); - assert!(timer.check_point.elapsed().as_millis() < 10); - } - - #[test] - fn test_elapsed() { - let timer = Timer::new(); - thread::sleep(time::Duration::from_millis(100)); - assert!(timer.elapsed().as_millis() > 100); - assert!(timer.elapsed_seconds() > 0.1); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/utils.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/utils.rs deleted file mode 100644 index 2e80676..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/src/utils/utils.rs +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::sync::Mutex; -use num_traits::Num; - -/// Non recursive mutex -pub type NonRecursiveMutex = Mutex<()>; - -/// Round up X to the nearest multiple of Y -#[inline] -pub fn round_up(x: T, y: T) -> T -where T : Num + Copy -{ - div_round_up(x, y) * y -} - -/// Rounded-up division -#[inline] -pub fn div_round_up(x: T, y: T) -> T -where T : Num + Copy -{ - (x / y) + if x % y != T::zero() {T::one()} else {T::zero()} -} - -/// Round down X to the nearest multiple of Y -#[inline] -pub fn round_down(x: T, y: T) -> T -where T : Num + Copy -{ - (x / y) * y -} - -/// Is aligned -#[inline] -pub fn is_aligned(x: T, y: T) -> bool -where T : Num + Copy -{ - x % y == T::zero() -} - -#[inline] -pub fn is_512_aligned(x: u64) -> bool { - is_aligned(x, 512) -} - -#[inline] -pub fn is_4096_aligned(x: u64) -> bool { - is_aligned(x, 4096) -} - -/// all metadata of individual sub-component files is written in first 4KB for unified files -pub const METADATA_SIZE: usize = 4096; - -pub const BUFFER_SIZE_FOR_CACHED_IO: usize = 1024 * 1048576; - -pub const PBSTR: &str = "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"; - -pub const PBWIDTH: usize = 60; - -macro_rules! convert_types { - ($name:ident, $intput_type:ty, $output_type:ty) => { - /// Write data into file - pub fn $name(srcmat: &[$intput_type], npts: usize, dim: usize) -> Vec<$output_type> { - let mut destmat: Vec<$output_type> = Vec::new(); - for i in 0..npts { - for j in 0..dim { - destmat.push(srcmat[i * dim + j] as $output_type); - } - } - destmat - } - }; -} -convert_types!(convert_types_usize_u8, usize, u8); -convert_types!(convert_types_usize_u32, usize, u32); -convert_types!(convert_types_usize_u64, usize, u64); -convert_types!(convert_types_u64_usize, u64, usize); -convert_types!(convert_types_u32_usize, u32, usize); - -#[cfg(test)] -mod file_util_test { - use super::*; - use std::any::type_name; - - #[test] - fn round_up_test() { - assert_eq!(round_up(252, 8), 256); - assert_eq!(round_up(256, 8), 256); - } - - #[test] - fn div_round_up_test() { - assert_eq!(div_round_up(252, 8), 32); - assert_eq!(div_round_up(256, 8), 32); - } - - #[test] - fn round_down_test() { - assert_eq!(round_down(252, 8), 248); - assert_eq!(round_down(256, 8), 256); - } - - #[test] - fn is_aligned_test() { - assert!(!is_aligned(252, 8)); - assert!(is_aligned(256, 8)); - } - - #[test] - fn is_512_aligned_test() { - assert!(!is_512_aligned(520)); - assert!(is_512_aligned(512)); - } - - #[test] - fn is_4096_aligned_test() { - assert!(!is_4096_aligned(4090)); - assert!(is_4096_aligned(4096)); - } - - #[test] - fn convert_types_test() { - let data = vec![0u64, 1u64, 2u64]; - let output = convert_types_u64_usize(&data, 3, 1); - assert_eq!(output.len(), 3); - assert_eq!(type_of(output[0]), "usize"); - assert_eq!(output[0], 0usize); - - let data = vec![0usize, 1usize, 2usize]; - let output = convert_types_usize_u8(&data, 3, 1); - assert_eq!(output.len(), 3); - assert_eq!(type_of(output[0]), "u8"); - assert_eq!(output[0], 0u8); - - let data = vec![0usize, 1usize, 2usize]; - let output = convert_types_usize_u64(&data, 3, 1); - assert_eq!(output.len(), 3); - assert_eq!(type_of(output[0]), "u64"); - assert_eq!(output[0], 0u64); - - let data = vec![0u32, 1u32, 2u32]; - let output = convert_types_u32_usize(&data, 3, 1); - assert_eq!(output.len(), 3); - assert_eq!(type_of(output[0]), "usize"); - assert_eq!(output[0],0usize); - } - - fn type_of(_: T) -> &'static str { - type_name::() - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts.fbin b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts.fbin deleted file mode 100644 index 357a9db..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts.fbin and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts_2.fbin b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts_2.fbin deleted file mode 100644 index 9528e4b..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/siftsmall_learn_256pts_2.fbin and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_R4_L50_A1.2 b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_R4_L50_A1.2 deleted file mode 100644 index 9c803c3..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_R4_L50_A1.2 and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_saturated_R4_L50_A1.2 b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_saturated_R4_L50_A1.2 deleted file mode 100644 index a9dac10..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_1+2_saturated_R4_L50_A1.2 and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2 b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2 deleted file mode 100644 index 8170090..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2 and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2.data b/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2.data deleted file mode 100644 index 357a9db..0000000 Binary files a/packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/tests/data/truth_index_siftsmall_learn_256pts_R4_L50_A1.2.data and /dev/null differ diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/Cargo.toml deleted file mode 100644 index e750d95..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/Cargo.toml +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "logger" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -lazy_static = "1.4.0" -log="0.4.17" -once_cell = "1.17.1" -prost = "0.11.9" -prost-types = "0.11.9" -thiserror = "1.0.40" -win_etw_macros="0.1.8" -win_etw_provider="0.1.8" - -[build-dependencies] -prost-build = "0.11.9" - -[[example]] -name="trace_example" -path= "src/examples/trace_example.rs" - -[target."cfg(target_os=\"windows\")".build-dependencies.vcpkg] -version = "0.2" - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/build.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/build.rs deleted file mode 100644 index 76058f7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/build.rs +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::env; - -extern crate prost_build; - -fn main() { - let protopkg = vcpkg::find_package("protobuf").unwrap(); - let protobuf_path = protopkg.link_paths[0].parent().unwrap(); - - let protobuf_bin_path = protobuf_path - .join("tools") - .join("protobuf") - .join("protoc.exe") - .to_str() - .unwrap() - .to_string(); - env::set_var("PROTOC", protobuf_bin_path); - - let protobuf_inc_path = protobuf_path - .join("include") - .join("google") - .join("protobuf") - .to_str() - .unwrap() - .to_string(); - env::set_var("PROTOC_INCLUDE", protobuf_inc_path); - - prost_build::compile_protos(&["src/indexlog.proto"], &["src/"]).unwrap(); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/error_logger.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/error_logger.rs deleted file mode 100644 index 50069b4..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/error_logger.rs +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::log_error::LogError; -use crate::logger::indexlog::{ErrorLog, Log, LogLevel}; -use crate::message_handler::send_log; - -pub fn log_error(error_message: String) -> Result<(), LogError> { - let mut log = Log::default(); - let error_log = ErrorLog { - log_level: LogLevel::Error as i32, - error_message, - }; - log.error_log = Some(error_log); - - send_log(log) -} - -#[cfg(test)] -mod error_logger_test { - use super::*; - - #[test] - fn log_error_works() { - log_error(String::from("Error")).unwrap(); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/examples/trace_example.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/examples/trace_example.rs deleted file mode 100644 index 7933a56..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/examples/trace_example.rs +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use log::{debug, info, log_enabled, warn, Level}; -use logger::trace_logger::TraceLogger; - -// cargo run --example trace_example - -fn main() { - static LOGGER: TraceLogger = TraceLogger {}; - log::set_logger(&LOGGER) - .map(|()| log::set_max_level(log::LevelFilter::Trace)) - .unwrap(); - - info!("Rust logging n = {}", 42); - warn!("This is too much fun!"); - debug!("Maybe we can make this code work"); - - let error_is_enabled = log_enabled!(Level::Error); - let warn_is_enabled = log_enabled!(Level::Warn); - let info_is_enabled = log_enabled!(Level::Info); - let debug_is_enabled = log_enabled!(Level::Debug); - let trace_is_enabled = log_enabled!(Level::Trace); - println!( - "is_enabled? error: {:5?}, warn: {:5?}, info: {:5?}, debug: {:5?}, trace: {:5?}", - error_is_enabled, warn_is_enabled, info_is_enabled, debug_is_enabled, trace_is_enabled, - ); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/indexlog.proto b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/indexlog.proto deleted file mode 100644 index 68310ae..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/indexlog.proto +++ /dev/null @@ -1,50 +0,0 @@ -syntax = "proto3"; - -package diskann_logger; - -message Log { - IndexConstructionLog IndexConstructionLog = 1; - DiskIndexConstructionLog DiskIndexConstructionLog = 2; - ErrorLog ErrorLog = 3; - TraceLog TraceLog = 100; -} - -enum LogLevel { - UNSPECIFIED = 0; - Error = 1; - Warn = 2; - Info = 3; - Debug = 4; - Trace = 5; -} - -message IndexConstructionLog { - float PercentageComplete = 1; - float TimeSpentInSeconds = 2; - float GCyclesSpent = 3; - LogLevel LogLevel = 4; -} - -message DiskIndexConstructionLog { - DiskIndexConstructionCheckpoint checkpoint = 1; - float TimeSpentInSeconds = 2; - float GCyclesSpent = 3; - LogLevel LogLevel = 4; -} - -enum DiskIndexConstructionCheckpoint { - None = 0; - PqConstruction = 1; - InmemIndexBuild = 2; - DiskLayout = 3; -} - -message TraceLog { - string LogLine = 1; - LogLevel LogLevel = 2; -} - -message ErrorLog { - string ErrorMessage = 1; - LogLevel LogLevel = 2; -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/lib.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/lib.rs deleted file mode 100644 index 6cfe2d5..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/lib.rs +++ /dev/null @@ -1,19 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] - -pub mod logger { - pub mod indexlog { - include!(concat!(env!("OUT_DIR"), "/diskann_logger.rs")); - } -} - -pub mod error_logger; -pub mod log_error; -pub mod message_handler; -pub mod trace_logger; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/log_error.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/log_error.rs deleted file mode 100644 index 149d094..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/log_error.rs +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::sync::mpsc::SendError; - -use crate::logger::indexlog::Log; - -#[derive(thiserror::Error, Debug, Clone)] -pub enum LogError { - /// Sender failed to send message to the channel - #[error("IOError: {err}")] - SendError { - #[from] - err: SendError, - }, - - /// PoisonError which can be returned whenever a lock is acquired - /// Both Mutexes and RwLocks are poisoned whenever a thread fails while the lock is held - #[error("LockPoisonError: {err}")] - LockPoisonError { err: String }, - - /// Failed to create EtwPublisher - #[error("EtwProviderError: {err:?}")] - ETWProviderError { err: win_etw_provider::Error }, -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/message_handler.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/message_handler.rs deleted file mode 100644 index 37f352a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/message_handler.rs +++ /dev/null @@ -1,167 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::log_error::LogError; -use crate::logger::indexlog::DiskIndexConstructionCheckpoint; -use crate::logger::indexlog::Log; -use crate::logger::indexlog::LogLevel; - -use std::sync::mpsc::{self, Sender}; -use std::sync::Mutex; -use std::thread; - -use win_etw_macros::trace_logging_provider; - -trait MessagePublisher { - fn publish(&self, log_level: LogLevel, message: &str); -} - -// ETW provider - the GUID specified here is that of the default provider for Geneva Metric Extensions -// We are just using it as a placeholder until we have a version of OpenTelemetry exporter for Rust -#[trace_logging_provider(guid = "edc24920-e004-40f6-a8e1-0e6e48f39d84")] -trait EtwTraceProvider { - fn write(msg: &str); -} - -struct EtwPublisher { - provider: EtwTraceProvider, - publish_to_stdout: bool, -} - -impl EtwPublisher { - pub fn new() -> Result { - let provider = EtwTraceProvider::new(); - Ok(EtwPublisher { - provider, - publish_to_stdout: true, - }) - } -} - -fn log_level_to_etw(level: LogLevel) -> win_etw_provider::Level { - match level { - LogLevel::Error => win_etw_provider::Level::ERROR, - LogLevel::Warn => win_etw_provider::Level::WARN, - LogLevel::Info => win_etw_provider::Level::INFO, - LogLevel::Debug => win_etw_provider::Level::VERBOSE, - LogLevel::Trace => win_etw_provider::Level(6), - LogLevel::Unspecified => win_etw_provider::Level(6), - } -} - -fn i32_to_log_level(value: i32) -> LogLevel { - match value { - 0 => LogLevel::Unspecified, - 1 => LogLevel::Error, - 2 => LogLevel::Warn, - 3 => LogLevel::Info, - 4 => LogLevel::Debug, - 5 => LogLevel::Trace, - _ => LogLevel::Unspecified, - } -} - -impl MessagePublisher for EtwPublisher { - fn publish(&self, log_level: LogLevel, message: &str) { - let options = win_etw_provider::EventOptions { - level: Some(log_level_to_etw(log_level)), - ..Default::default() - }; - self.provider.write(Some(&options), message); - - if self.publish_to_stdout { - println!("{}", message); - } - } -} - -struct MessageProcessor { - sender: Mutex>, -} - -impl MessageProcessor { - pub fn start_processing() -> Self { - let (sender, receiver) = mpsc::channel::(); - thread::spawn(move || -> Result<(), LogError> { - for message in receiver { - // Process the received message - if let Some(indexlog) = message.index_construction_log { - let str = format!( - "Time for {}% of index build completed: {:.3} seconds, {:.3}B cycles", - indexlog.percentage_complete, - indexlog.time_spent_in_seconds, - indexlog.g_cycles_spent - ); - publish(i32_to_log_level(indexlog.log_level), &str)?; - } - - if let Some(disk_index_log) = message.disk_index_construction_log { - let str = format!( - "Time for disk index build [Checkpoint: {:?}] completed: {:.3} seconds, {:.3}B cycles", - DiskIndexConstructionCheckpoint::from_i32(disk_index_log.checkpoint).unwrap_or(DiskIndexConstructionCheckpoint::None), - disk_index_log.time_spent_in_seconds, - disk_index_log.g_cycles_spent - ); - publish(i32_to_log_level(disk_index_log.log_level), &str)?; - } - - if let Some(tracelog) = message.trace_log { - let str = format!("{}:{}", tracelog.log_level, tracelog.log_line); - publish(i32_to_log_level(tracelog.log_level), &str)?; - } - - if let Some(err) = message.error_log { - publish(i32_to_log_level(err.log_level), &err.error_message)?; - } - } - - Ok(()) - }); - - let sender = Mutex::new(sender); - MessageProcessor { sender } - } - - /// Log the message. - fn log(&self, message: Log) -> Result<(), LogError> { - Ok(self - .sender - .lock() - .map_err(|err| LogError::LockPoisonError { - err: err.to_string(), - })? - .send(message)?) - } -} - -lazy_static::lazy_static! { - /// Singleton logger. - static ref PROCESSOR: MessageProcessor = { - - MessageProcessor::start_processing() - }; -} - -lazy_static::lazy_static! { - /// Singleton publisher. - static ref PUBLISHER: Result = { - EtwPublisher::new() - }; -} - -/// Send a message to the logging system. -pub fn send_log(message: Log) -> Result<(), LogError> { - PROCESSOR.log(message) -} - -fn publish(log_level: LogLevel, message: &str) -> Result<(), LogError> { - match *PUBLISHER { - Ok(ref etw_publisher) => { - etw_publisher.publish(log_level, message); - Ok(()) - } - Err(ref err) => Err(LogError::ETWProviderError { err: err.clone() }), - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/trace_logger.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/trace_logger.rs deleted file mode 100644 index 96ef386..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/trace_logger.rs +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::logger::indexlog::{Log, TraceLog}; -use crate::message_handler::send_log; - -use log; - -pub struct TraceLogger {} - -fn level_to_i32(value: log::Level) -> i32 { - match value { - log::Level::Error => 1, - log::Level::Warn => 2, - log::Level::Info => 3, - log::Level::Debug => 4, - log::Level::Trace => 5, - } -} - -impl log::Log for TraceLogger { - fn enabled(&self, metadata: &log::Metadata) -> bool { - metadata.level() <= log::max_level() - } - - fn log(&self, record: &log::Record) { - let message = record.args().to_string(); - let metadata = record.metadata(); - let mut log = Log::default(); - let trace_log = TraceLog { - log_line: message, - log_level: level_to_i32(metadata.level()), - }; - log.trace_log = Some(trace_log); - let _ = send_log(log); - } - - fn flush(&self) {} -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/Cargo.toml deleted file mode 100644 index 057f9e8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "platform" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -log="0.4.18" -winapi = { version = "0.3.9", features = ["errhandlingapi", "fileapi", "ioapiset", "handleapi", "winnt", "minwindef", "basetsd", "winerror", "winbase"] } - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_handle.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_handle.rs deleted file mode 100644 index 23da879..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_handle.rs +++ /dev/null @@ -1,212 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::ffi::CString; -use std::{io, ptr}; - -use winapi::um::fileapi::OPEN_EXISTING; -use winapi::um::winbase::{FILE_FLAG_NO_BUFFERING, FILE_FLAG_OVERLAPPED, FILE_FLAG_RANDOM_ACCESS}; -use winapi::um::winnt::{FILE_SHARE_DELETE, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE}; - -use winapi::{ - shared::minwindef::DWORD, - um::{ - errhandlingapi::GetLastError, - fileapi::CreateFileA, - handleapi::{CloseHandle, INVALID_HANDLE_VALUE}, - winnt::HANDLE, - }, -}; - -pub const FILE_ATTRIBUTE_READONLY: DWORD = 0x00000001; - -/// `AccessMode` determines how a file can be accessed. -/// These modes are used when creating or opening a file to decide what operations are allowed -/// to be performed on the file. -/// -/// # Variants -/// -/// - `Read`: The file is opened in read-only mode. -/// -/// - `Write`: The file is opened in write-only mode. -/// -/// - `ReadWrite`: The file is opened for both reading and writing. -pub enum AccessMode { - Read, - Write, - ReadWrite, -} - -/// `ShareMode` determines how a file can be shared. -/// -/// These modes are used when creating or opening a file to decide what operations other -/// opening instances of the file can perform on it. -/// # Variants -/// - `None`: Prevents other processes from opening a file if they request delete, -/// read, or write access. -/// -/// - `Read`: Allows subsequent open operations on the same file to request read access. -/// -/// - `Write`: Allows subsequent open operations on the same file file to request write access. -/// -/// - `Delete`: Allows subsequent open operations on the same file file to request delete access. -pub enum ShareMode { - None, - Read, - Write, - Delete, -} - -/// # Windows File Handle Wrapper -/// -/// Introduces a Rust-friendly wrapper around the native Windows `HANDLE` object, `FileHandle`. -/// `FileHandle` provides safe creation and automatic cleanup of Windows file handles, leveraging Rust's ownership model. - -/// `FileHandle` struct that wraps a native Windows `HANDLE` object -#[cfg(target_os = "windows")] -pub struct FileHandle { - handle: HANDLE, -} - -impl FileHandle { - /// Creates a new `FileHandle` by opening an existing file with the given access and shared mode. - /// - /// This function is marked unsafe because it creates a raw pointer to the filename and try to create - /// a Windows `HANDLE` object without checking if you have sufficient permissions. - /// - /// # Safety - /// - /// Ensure that the file specified by `file_name` is valid and the calling process has - /// sufficient permissions to perform the specified `access_mode` and `share_mode` operations. - /// - /// # Parameters - /// - /// - `file_name`: The name of the file. - /// - `access_mode`: The access mode to be used for the file. - /// - `share_mode`: The share mode to be used for the file - /// - /// # Errors - /// This function will return an error if the `file_name` is invalid or if the file cannot - /// be opened with the specified `access_mode` and `share_mode`. - pub unsafe fn new( - file_name: &str, - access_mode: AccessMode, - share_mode: ShareMode, - ) -> io::Result { - let file_name_c = CString::new(file_name).map_err(|_| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid file name. {}", file_name), - ) - })?; - - let dw_desired_access = match access_mode { - AccessMode::Read => GENERIC_READ, - AccessMode::Write => GENERIC_WRITE, - AccessMode::ReadWrite => GENERIC_READ | GENERIC_WRITE, - }; - - let dw_share_mode = match share_mode { - ShareMode::None => 0, - ShareMode::Read => FILE_SHARE_READ, - ShareMode::Write => FILE_SHARE_WRITE, - ShareMode::Delete => FILE_SHARE_DELETE, - }; - - let dw_flags_and_attributes = FILE_ATTRIBUTE_READONLY - | FILE_FLAG_NO_BUFFERING - | FILE_FLAG_OVERLAPPED - | FILE_FLAG_RANDOM_ACCESS; - - let handle = unsafe { - CreateFileA( - file_name_c.as_ptr(), - dw_desired_access, - dw_share_mode, - ptr::null_mut(), - OPEN_EXISTING, - dw_flags_and_attributes, - ptr::null_mut(), - ) - }; - - if handle == INVALID_HANDLE_VALUE { - let error_code = unsafe { GetLastError() }; - Err(io::Error::from_raw_os_error(error_code as i32)) - } else { - Ok(Self { handle }) - } - } - - pub fn raw_handle(&self) -> HANDLE { - self.handle - } -} - -impl Drop for FileHandle { - /// Automatically closes the `FileHandle` when it goes out of scope. - /// Any errors in closing the handle are logged, as `Drop` does not support returning `Result`. - fn drop(&mut self) { - let result = unsafe { CloseHandle(self.handle) }; - if result == 0 { - let error_code = unsafe { GetLastError() }; - let error = io::Error::from_raw_os_error(error_code as i32); - - // Only log the error if dropping the handle fails, since Rust's Drop trait does not support returning Result types from the drop method, - // and panicking in the drop method is considered bad practice - log::warn!("Error when dropping IOCompletionPort: {:?}", error); - } - } -} - -/// Returns a `FileHandle` with an `INVALID_HANDLE_VALUE`. -impl Default for FileHandle { - fn default() -> Self { - Self { - handle: INVALID_HANDLE_VALUE, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs::File; - use std::path::Path; - - #[test] - fn test_create_file() { - // Create a dummy file - let dummy_file_path = "dummy_file.txt"; - { - let _file = File::create(dummy_file_path).expect("Failed to create dummy file."); - } - - let path = Path::new(dummy_file_path); - { - let file_handle = unsafe { - FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) - }; - - // Check that the file handle is valid - assert!(file_handle.is_ok()); - } - - // Try to delete the file. If the handle was correctly dropped, this should succeed. - match std::fs::remove_file(dummy_file_path) { - Ok(()) => (), // File was deleted successfully, which means the handle was closed. - Err(e) => panic!("Failed to delete file: {}", e), // Failed to delete the file, likely because the handle is still open. - } - } - - #[test] - fn test_file_not_found() { - let path = Path::new("non_existent_file.txt"); - let file_handle = - unsafe { FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) }; - - // Check that opening a non-existent file returns an error - assert!(file_handle.is_err()); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_io.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_io.rs deleted file mode 100644 index e5de247..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_io.rs +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -/// The module provides unsafe wrappers around two Windows API functions: `ReadFile` and `GetQueuedCompletionStatus`. -/// -/// These wrappers aim to simplify and abstract the use of these functions, providing easier error handling and a safer interface. -/// They return standard Rust `io::Result` types for convenience and consistency with the rest of the Rust standard library. -use std::io; -use std::ptr; - -use winapi::{ - ctypes::c_void, - shared::{ - basetsd::ULONG_PTR, - minwindef::{DWORD, FALSE}, - winerror::{ERROR_IO_PENDING, WAIT_TIMEOUT}, - }, - um::{ - errhandlingapi::GetLastError, fileapi::ReadFile, ioapiset::GetQueuedCompletionStatus, - minwinbase::OVERLAPPED, - }, -}; - -use crate::FileHandle; -use crate::IOCompletionPort; - -/// Asynchronously queue a read request from a file into a buffer slice. -/// -/// Wraps the unsafe Windows API function `ReadFile`, making it safe to call only when the overlapped buffer -/// remains valid and unchanged anywhere else during the entire async operation. -/// -/// Returns a boolean indicating whether the read operation completed synchronously or is pending. -/// -/// # Safety -/// -/// This function is marked as `unsafe` because it uses raw pointers and requires the caller to ensure -/// that the buffer slice and the overlapped buffer stay valid during the whole async operation. -pub unsafe fn read_file_to_slice( - file_handle: &FileHandle, - buffer_slice: &mut [T], - overlapped: *mut OVERLAPPED, - offset: u64, -) -> io::Result { - let num_bytes = std::mem::size_of_val(buffer_slice); - unsafe { - ptr::write(overlapped, std::mem::zeroed()); - (*overlapped).u.s_mut().Offset = offset as u32; - (*overlapped).u.s_mut().OffsetHigh = (offset >> 32) as u32; - } - - let result = unsafe { - ReadFile( - file_handle.raw_handle(), - buffer_slice.as_mut_ptr() as *mut c_void, - num_bytes as DWORD, - ptr::null_mut(), - overlapped, - ) - }; - - match result { - FALSE => { - let error = unsafe { GetLastError() }; - if error != ERROR_IO_PENDING { - Err(io::Error::from_raw_os_error(error as i32)) - } else { - Ok(false) - } - } - _ => Ok(true), - } -} - -/// Retrieves the results of an asynchronous I/O operation on an I/O completion port. -/// -/// Wraps the unsafe Windows API function `GetQueuedCompletionStatus`, making it safe to call only when the overlapped buffer -/// remains valid and unchanged anywhere else during the entire async operation. -/// -/// Returns a boolean indicating whether an I/O operation completed synchronously or is still pending. -/// -/// # Safety -/// -/// This function is marked as `unsafe` because it uses raw pointers and requires the caller to ensure -/// that the overlapped buffer stays valid during the whole async operation. -pub unsafe fn get_queued_completion_status( - completion_port: &IOCompletionPort, - lp_number_of_bytes: &mut DWORD, - lp_completion_key: &mut ULONG_PTR, - lp_overlapped: *mut *mut OVERLAPPED, - dw_milliseconds: DWORD, -) -> io::Result { - let result = unsafe { - GetQueuedCompletionStatus( - completion_port.raw_handle(), - lp_number_of_bytes, - lp_completion_key, - lp_overlapped, - dw_milliseconds, - ) - }; - - match result { - 0 => { - let error = unsafe { GetLastError() }; - if error == WAIT_TIMEOUT { - Ok(false) - } else { - Err(io::Error::from_raw_os_error(error as i32)) - } - } - _ => Ok(true), - } -} - -#[cfg(test)] -mod tests { - use crate::file_handle::{AccessMode, ShareMode}; - - use super::*; - use std::fs::File; - use std::io::Write; - use std::path::Path; - - #[test] - fn test_read_file_to_slice() { - // Create a temporary file and write some data into it - let path = Path::new("temp.txt"); - { - let mut file = File::create(path).unwrap(); - file.write_all(b"Hello, world!").unwrap(); - } - - let mut buffer: [u8; 512] = [0; 512]; - let mut overlapped = unsafe { std::mem::zeroed::() }; - { - let file_handle = unsafe { - FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) - } - .unwrap(); - - // Call the function under test - let result = - unsafe { read_file_to_slice(&file_handle, &mut buffer, &mut overlapped, 0) }; - - assert!(result.is_ok()); - let result_str = std::str::from_utf8(&buffer[.."Hello, world!".len()]).unwrap(); - assert_eq!(result_str, "Hello, world!"); - } - - // Clean up - std::fs::remove_file("temp.txt").unwrap(); - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/io_completion_port.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/io_completion_port.rs deleted file mode 100644 index 5bb3322..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/io_completion_port.rs +++ /dev/null @@ -1,142 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::io; - -use winapi::{ - ctypes::c_void, - shared::{basetsd::ULONG_PTR, minwindef::DWORD}, - um::{ - errhandlingapi::GetLastError, - handleapi::{CloseHandle, INVALID_HANDLE_VALUE}, - ioapiset::CreateIoCompletionPort, - winnt::HANDLE, - }, -}; - -use crate::FileHandle; - -/// This module provides a safe and idiomatic Rust interface over the IOCompletionPort handle and associated Windows API functions. -/// This struct represents an I/O completion port, which is an object used in asynchronous I/O operations on Windows. -pub struct IOCompletionPort { - io_completion_port: HANDLE, -} - -impl IOCompletionPort { - /// Create a new IOCompletionPort. - /// This function wraps the Windows CreateIoCompletionPort function, providing error handling and automatic resource management. - /// - /// # Arguments - /// - /// * `file_handle` - A reference to a FileHandle to associate with the IOCompletionPort. - /// * `existing_completion_port` - An optional reference to an existing IOCompletionPort. If provided, the new IOCompletionPort will be associated with it. - /// * `completion_key` - The completion key associated with the file handle. - /// * `number_of_concurrent_threads` - The maximum number of threads that the operating system can allow to concurrently process I/O completion packets for the I/O completion port. - /// - /// # Return - /// - /// Returns a Result with the new IOCompletionPort if successful, or an io::Error if the function fails. - pub fn new( - file_handle: &FileHandle, - existing_completion_port: Option<&IOCompletionPort>, - completion_key: ULONG_PTR, - number_of_concurrent_threads: DWORD, - ) -> io::Result { - let io_completion_port = unsafe { - CreateIoCompletionPort( - file_handle.raw_handle(), - existing_completion_port - .map_or(std::ptr::null_mut::(), |io_completion_port| { - io_completion_port.raw_handle() - }), - completion_key, - number_of_concurrent_threads, - ) - }; - - if io_completion_port == INVALID_HANDLE_VALUE { - let error_code = unsafe { GetLastError() }; - return Err(io::Error::from_raw_os_error(error_code as i32)); - } - - Ok(IOCompletionPort { io_completion_port }) - } - - pub fn raw_handle(&self) -> HANDLE { - self.io_completion_port - } -} - -impl Drop for IOCompletionPort { - /// Drop method for IOCompletionPort. - /// This wraps the Windows CloseHandle function, providing automatic resource cleanup when the IOCompletionPort is dropped. - /// If an error occurs while dropping, it is logged and the drop continues. This is because panicking in Drop can cause unwinding issues. - fn drop(&mut self) { - let result = unsafe { CloseHandle(self.io_completion_port) }; - if result == 0 { - let error_code = unsafe { GetLastError() }; - let error = io::Error::from_raw_os_error(error_code as i32); - - // Only log the error if dropping the handle fails, since Rust's Drop trait does not support returning Result types from the drop method, - // and panicking in the drop method is considered bad practice - log::warn!("Error when dropping IOCompletionPort: {:?}", error); - } - } -} - -impl Default for IOCompletionPort { - /// Create a default IOCompletionPort, whose handle is set to INVALID_HANDLE_VALUE. - /// Returns a new IOCompletionPort with handle set to INVALID_HANDLE_VALUE. - fn default() -> Self { - Self { - io_completion_port: INVALID_HANDLE_VALUE, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::file_handle::{AccessMode, ShareMode}; - - #[test] - fn create_io_completion_port() { - let file_name = "../diskann/tests/data/delete_set_50pts.bin"; - let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) } - .expect("Failed to create file handle."); - - let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0); - - assert!( - io_completion_port.is_ok(), - "Failed to create IOCompletionPort." - ); - } - - #[test] - fn drop_io_completion_port() { - let file_name = "../diskann/tests/data/delete_set_50pts.bin"; - let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) } - .expect("Failed to create file handle."); - - let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0) - .expect("Failed to create IOCompletionPort."); - - // After this line, io_completion_port goes out of scope and its Drop trait will be called. - let _ = io_completion_port; - // We have no easy way to test that the Drop trait works correctly, but if it doesn't, - // a resource leak or other problem may become apparent in later tests or in real use of the code. - } - - #[test] - fn default_io_completion_port() { - let io_completion_port = IOCompletionPort::default(); - assert_eq!( - io_completion_port.raw_handle(), - INVALID_HANDLE_VALUE, - "Default IOCompletionPort did not have INVALID_HANDLE_VALUE." - ); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/lib.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/lib.rs deleted file mode 100644 index e282570..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/lib.rs +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] - -pub mod perf; -pub use perf::{get_process_cycle_time, get_process_handle}; - -pub mod file_io; -pub use file_io::{get_queued_completion_status, read_file_to_slice}; - -pub mod file_handle; -pub use file_handle::FileHandle; - -pub mod io_completion_port; -pub use io_completion_port::IOCompletionPort; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/perf.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/perf.rs deleted file mode 100644 index 1ea146f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/perf.rs +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[cfg(target_os = "windows")] -#[link(name = "kernel32")] -extern "system" { - fn OpenProcess(dwDesiredAccess: u32, bInheritHandle: bool, dwProcessId: u32) -> usize; - fn QueryProcessCycleTime(hProcess: usize, lpCycleTime: *mut u64) -> bool; - fn GetCurrentProcessId() -> u32; -} - -/// Get current process handle. -pub fn get_process_handle() -> Option { - if cfg!(windows) { - const PROCESS_QUERY_INFORMATION: u32 = 0x0400; - const PROCESS_VM_READ: u32 = 0x0010; - - unsafe { - let current_process_id = GetCurrentProcessId(); - let handle = OpenProcess( - PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, - false, - current_process_id, - ); - if handle == 0 { - None - } else { - Some(handle) - } - } - } else { - None - } -} - -pub fn get_process_cycle_time(process_handle: Option) -> Option { - let mut cycle_time: u64 = 0; - if cfg!(windows) { - if let Some(handle) = process_handle { - let result = unsafe { QueryProcessCycleTime(handle, &mut cycle_time as *mut u64) }; - if result { - return Some(cycle_time); - } - } - } - - None -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/project.code-workspace b/packages/leann-backend-diskann/third_party/DiskANN/rust/project.code-workspace deleted file mode 100644 index 29bed00..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/project.code-workspace +++ /dev/null @@ -1,58 +0,0 @@ -{ - "folders": [ - { - "path": "." - } - ], - "settings": { - "search.exclude": { - "target": true, - }, - "files.exclude": { - "target": true, - }, - "rust-analyzer.linkedProjects": [ - ".\\vector\\Cargo.toml", - ".\\vector\\Cargo.toml", - ".\\vector\\Cargo.toml", - ".\\diskann\\Cargo.toml" - ], - "[rust]": { - "editor.defaultFormatter": "rust-lang.rust-analyzer", - "editor.formatOnSave": true, - } - }, - "launch": { - "version": "0.2.0", - "configurations": [ - { - "name": "Build memory index", - "type": "cppvsdbg", - "request": "launch", - "program": "${workspaceRoot}\\target\\debug\\build_memory_index.exe", - "args": [ - "--data_type", - "float", - "--dist_fn", - "l2", - "--data_path", - ".\\base1m.fbin", - "--index_path_prefix", - ".\\rust_index_sift_base_R32_L50_A1.2_T1", - "-R", - "64", - "-L", - "100", - "--alpha", - "1.2", - "-T", - "1" - ], - "stopAtEntry": false, - "cwd": "c:\\data", - "environment": [], - "externalConsole": true - }, - ] - } -} \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/readme.md b/packages/leann-backend-diskann/third_party/DiskANN/rust/readme.md deleted file mode 100644 index a6c5a1b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/readme.md +++ /dev/null @@ -1,25 +0,0 @@ - -# readme - -run commands under disnann_rust directory. - -build: -``` -cargo build // Debug - -cargo build -r // Release -``` - - -run: -``` -cargo run // Debug - -cargo run -r // Release -``` - - -test: -``` -cargo test -``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/rust-toolchain.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/rust-toolchain.toml deleted file mode 100644 index 183a72c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/rust-toolchain.toml +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[toolchain] -channel = "stable" diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/Cargo.toml deleted file mode 100644 index 709a290..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/Cargo.toml +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "vector" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -half = "2.2.1" -thiserror = "1.0.40" -bytemuck = "1.7.0" - -[build-dependencies] -cc = "1.0.79" - -[dev-dependencies] -base64 = "0.21.2" -bincode = "1.3.3" -serde = "1.0.163" -approx = "0.5.1" -rand = "0.8.5" - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/build.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/build.rs deleted file mode 100644 index 2d36c21..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/build.rs +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -fn main() { - println!("cargo:rerun-if-changed=distance.c"); - if cfg!(target_os = "macos") { - std::env::set_var("CFLAGS", "-mavx2 -mfma -Wno-error -MP -O2 -D NDEBUG -D MKL_ILP64 -D USE_AVX2 -D USE_ACCELERATED_PQ -D NOMINMAX -D _TARGET_ARM_APPLE_DARWIN"); - - cc::Build::new() - .file("distance.c") - .warnings_into_errors(true) - .debug(false) - .target("x86_64-apple-darwin") - .compile("nativefunctions.lib"); - } else { - std::env::set_var("CFLAGS", "/permissive- /MP /ifcOutput /GS- /W3 /Gy /Zi /Gm- /O2 /Ob2 /Zc:inline /fp:fast /D NDEBUG /D MKL_ILP64 /D USE_AVX2 /D USE_ACCELERATED_PQ /D NOMINMAX /fp:except- /errorReport:prompt /WX /openmp:experimental /Zc:forScope /GR /arch:AVX2 /Gd /Oy /Oi /MD /std:c++14 /FC /EHsc /nologo /Ot"); - // std::env::set_var("CFLAGS", "/permissive- /MP /ifcOutput /GS- /W3 /Gy /Zi /Gm- /Obd /Zc:inline /fp:fast /D DEBUG /D MKL_ILP64 /D USE_AVX2 /D USE_ACCELERATED_PQ /D NOMINMAX /fp:except- /errorReport:prompt /WX /openmp:experimental /Zc:forScope /GR /arch:AVX512 /Gd /Oy /Oi /MD /std:c++14 /FC /EHsc /nologo /Ot"); - - cc::Build::new() - .file("distance.c") - .warnings_into_errors(true) - .debug(false) - .compile("nativefunctions"); - - println!("cargo:rustc-link-arg=nativefunctions.lib"); - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/distance.c b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/distance.c deleted file mode 100644 index ee5333a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/distance.c +++ /dev/null @@ -1,35 +0,0 @@ -#include -#include - -inline __m256i load_128bit_to_256bit(const __m128i *ptr) -{ - __m128i value128 = _mm_loadu_si128(ptr); - __m256i value256 = _mm256_castsi128_si256(value128); - return _mm256_inserti128_si256(value256, _mm_setzero_si128(), 1); -} - -float distance_compare_avx512f_f16(const unsigned char *vec1, const unsigned char *vec2, size_t size) -{ - __m512 sum_squared_diff = _mm512_setzero_ps(); - - for (int i = 0; i < size / 16; i += 1) - { - __m512 v1 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(vec1 + i * 2 * 16))); - __m512 v2 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(vec2 + i * 2 * 16))); - - __m512 diff = _mm512_sub_ps(v1, v2); - sum_squared_diff = _mm512_fmadd_ps(diff, diff, sum_squared_diff); - } - - size_t i = (size / 16) * 16; - - if (i != size) - { - __m512 va = _mm512_cvtph_ps(load_128bit_to_256bit((const __m128i *)(vec1 + i * 2))); - __m512 vb = _mm512_cvtph_ps(load_128bit_to_256bit((const __m128i *)(vec2 + i * 2))); - __m512 diff512 = _mm512_sub_ps(va, vb); - sum_squared_diff = _mm512_fmadd_ps(diff512, diff512, sum_squared_diff); - } - - return _mm512_reduce_add_ps(sum_squared_diff); -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance.rs deleted file mode 100644 index 8ca6cb2..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance.rs +++ /dev/null @@ -1,442 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use crate::l2_float_distance::{distance_l2_vector_f16, distance_l2_vector_f32}; -use crate::{Half, Metric}; - -/// Distance contract for full-precision vertex -pub trait FullPrecisionDistance { - /// Get the distance between vertex a and vertex b - fn distance_compare(a: &[T; N], b: &[T; N], vec_type: Metric) -> f32; -} - -// reason = "Not supported Metric type Metric::Cosine" -#[allow(clippy::panic)] -impl FullPrecisionDistance for [f32; N] { - /// Calculate distance between two f32 Vertex - #[inline(always)] - fn distance_compare(a: &[f32; N], b: &[f32; N], metric: Metric) -> f32 { - match metric { - Metric::L2 => distance_l2_vector_f32::(a, b), - _ => panic!("Not supported Metric type {:?}", metric), - } - } -} - -// reason = "Not supported Metric type Metric::Cosine" -#[allow(clippy::panic)] -impl FullPrecisionDistance for [Half; N] { - fn distance_compare(a: &[Half; N], b: &[Half; N], metric: Metric) -> f32 { - match metric { - Metric::L2 => distance_l2_vector_f16::(a, b), - _ => panic!("Not supported Metric type {:?}", metric), - } - } -} - -// reason = "Not yet supported Vector i8" -#[allow(clippy::panic)] -impl FullPrecisionDistance for [i8; N] { - fn distance_compare(_a: &[i8; N], _b: &[i8; N], _metric: Metric) -> f32 { - panic!("Not supported VectorType i8") - } -} - -// reason = "Not yet supported Vector u8" -#[allow(clippy::panic)] -impl FullPrecisionDistance for [u8; N] { - fn distance_compare(_a: &[u8; N], _b: &[u8; N], _metric: Metric) -> f32 { - panic!("Not supported VectorType u8") - } -} - -#[cfg(test)] -mod distance_test { - use super::*; - - #[repr(C, align(32))] - pub struct F32Slice112([f32; 112]); - - #[repr(C, align(32))] - pub struct F16Slice112([Half; 112]); - - fn get_turing_test_data() -> (F32Slice112, F32Slice112) { - let a_slice: [f32; 112] = [ - 0.13961786, - -0.031577103, - -0.09567415, - 0.06695563, - -0.1588727, - 0.089852564, - -0.019837005, - 0.07497972, - 0.010418192, - -0.054594643, - 0.08613386, - -0.05103466, - 0.16568437, - -0.02703799, - 0.00728657, - -0.15313251, - 0.16462992, - -0.030570814, - 0.11635703, - 0.23938893, - 0.018022912, - -0.12646551, - 0.018048918, - -0.035986554, - 0.031986624, - -0.015286017, - 0.010117953, - -0.032691937, - 0.12163067, - -0.04746277, - 0.010213069, - -0.043672588, - -0.099362016, - 0.06599016, - -0.19397286, - -0.13285528, - -0.22040887, - 0.017690737, - -0.104262285, - -0.0044555613, - -0.07383778, - -0.108652934, - 0.13399786, - 0.054912474, - 0.20181285, - 0.1795591, - -0.05425621, - -0.10765217, - 0.1405377, - -0.14101997, - -0.12017701, - 0.011565498, - 0.06952187, - 0.060136646, - 0.0023214167, - 0.04204699, - 0.048470616, - 0.17398086, - 0.024218207, - -0.15626553, - -0.11291045, - -0.09688122, - 0.14393932, - -0.14713104, - -0.108876854, - 0.035279203, - -0.05440188, - 0.017205412, - 0.011413814, - 0.04009471, - 0.11070237, - -0.058998976, - 0.07260045, - -0.057893746, - -0.0036240944, - -0.0064988653, - -0.13842176, - -0.023219328, - 0.0035885905, - -0.0719257, - -0.21335067, - 0.11415403, - -0.0059823603, - 0.12091869, - 0.08136634, - -0.10769281, - 0.024518685, - 0.0009200326, - -0.11628049, - 0.07448965, - 0.13736208, - -0.04144517, - -0.16426727, - -0.06380103, - -0.21386267, - 0.022373492, - -0.05874115, - 0.017314062, - -0.040344074, - 0.01059176, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ]; - let b_slice: [f32; 112] = [ - -0.07209058, - -0.17755842, - -0.030627966, - 0.163028, - -0.2233766, - 0.057412963, - 0.0076995124, - -0.017121306, - -0.015759075, - -0.026947778, - -0.010282468, - -0.23968373, - -0.021486737, - -0.09903155, - 0.09361805, - 0.0042711576, - -0.08695552, - -0.042165346, - 0.064218745, - -0.06707651, - 0.07846054, - 0.12235762, - -0.060716823, - 0.18496591, - -0.13023394, - 0.022469055, - 0.056764495, - 0.07168404, - -0.08856144, - -0.15343173, - 0.099879816, - -0.033529017, - 0.0795304, - -0.009242254, - -0.10254546, - 0.13086525, - -0.101518914, - -0.1031299, - -0.056826904, - 0.033196196, - 0.044143833, - -0.049787212, - -0.018148342, - -0.11172959, - -0.06776237, - -0.09185828, - -0.24171598, - 0.05080982, - -0.0727684, - 0.045031235, - -0.11363879, - -0.063389264, - 0.105850354, - -0.19847773, - 0.08828623, - -0.087071925, - 0.033512704, - 0.16118294, - 0.14111553, - 0.020884402, - -0.088860825, - 0.018745849, - 0.047522716, - -0.03665169, - 0.15726231, - -0.09930561, - 0.057844743, - -0.10532736, - -0.091297254, - 0.067029804, - 0.04153976, - 0.06393326, - 0.054578528, - 0.0038539872, - 0.1023088, - -0.10653885, - -0.108500294, - -0.046606563, - 0.020439683, - -0.120957725, - -0.13334097, - -0.13425854, - -0.20481694, - 0.07009538, - 0.08660361, - -0.0096641015, - 0.095316306, - -0.002898167, - -0.19680002, - 0.08466311, - 0.04812689, - -0.028978813, - 0.04780206, - -0.2001506, - -0.036866356, - -0.023720587, - 0.10731964, - 0.05517358, - -0.09580819, - 0.14595725, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ]; - - (F32Slice112(a_slice), F32Slice112(b_slice)) - } - - fn get_turing_test_data_f16() -> (F16Slice112, F16Slice112) { - let (a_slice, b_slice) = get_turing_test_data(); - let a_data = a_slice.0.iter().map(|x| Half::from_f32(*x)); - let b_data = b_slice.0.iter().map(|x| Half::from_f32(*x)); - - ( - F16Slice112(a_data.collect::>().try_into().unwrap()), - F16Slice112(b_data.collect::>().try_into().unwrap()), - ) - } - - use crate::test_util::*; - use approx::assert_abs_diff_eq; - - #[test] - fn test_dist_l2_float_turing() { - // two vectors are allocated in the contiguous heap memory - let (a_slice, b_slice) = get_turing_test_data(); - let distance = <[f32; 112] as FullPrecisionDistance>::distance_compare( - &a_slice.0, - &b_slice.0, - Metric::L2, - ); - - assert_abs_diff_eq!( - distance, - no_vector_compare_f32(&a_slice.0, &b_slice.0), - epsilon = 1e-6 - ); - } - - #[test] - fn test_dist_l2_f16_turing() { - // two vectors are allocated in the contiguous heap memory - let (a_slice, b_slice) = get_turing_test_data_f16(); - let distance = <[Half; 112] as FullPrecisionDistance>::distance_compare( - &a_slice.0, - &b_slice.0, - Metric::L2, - ); - - // Note the variance between the full 32 bit precision and the 16 bit precision - assert_eq!(distance, no_vector_compare_f16(&a_slice.0, &b_slice.0)); - } - - #[test] - fn distance_test() { - #[repr(C, align(32))] - struct Vector32ByteAligned { - v: [f32; 512], - } - - // two vectors are allocated in the contiguous heap memory - let two_vec = Box::new(Vector32ByteAligned { - v: [ - 69.02492, 78.84786, 63.125072, 90.90581, 79.2592, 70.81731, 3.0829668, 33.33287, - 20.777142, 30.147898, 23.681915, 42.553043, 12.602162, 7.3808074, 19.157589, - 65.6791, 76.44677, 76.89124, 86.40756, 84.70118, 87.86142, 16.126896, 5.1277637, - 95.11038, 83.946945, 22.735607, 11.548555, 59.51482, 24.84603, 15.573776, 78.27185, - 71.13179, 38.574017, 80.0228, 13.175261, 62.887978, 15.205181, 18.89392, 96.13162, - 87.55455, 34.179806, 62.920044, 4.9305916, 54.349373, 21.731495, 14.982187, - 40.262867, 20.15214, 36.61963, 72.450806, 55.565, 95.5375, 93.73356, 95.36308, - 66.30762, 58.0397, 18.951357, 67.11702, 43.043316, 30.65622, 99.85361, 2.5889993, - 27.844774, 39.72441, 46.463238, 71.303764, 90.45308, 36.390602, 63.344395, - 26.427078, 35.99528, 82.35505, 32.529175, 23.165905, 74.73179, 9.856939, 59.38126, - 35.714924, 79.81213, 46.704124, 24.47884, 36.01743, 0.46678782, 29.528152, - 1.8980742, 24.68853, 75.58984, 98.72279, 68.62601, 11.890173, 49.49361, 55.45572, - 72.71067, 34.107483, 51.357758, 76.400635, 81.32725, 66.45081, 17.848074, - 62.398876, 94.20444, 2.10886, 17.416393, 64.88253, 29.000723, 62.434315, 53.907238, - 70.51412, 78.70744, 55.181683, 64.45116, 23.419212, 53.68544, 43.506958, 46.89598, - 35.905994, 64.51397, 91.95555, 20.322979, 74.80128, 97.548744, 58.312725, 78.81985, - 31.911612, 14.445949, 49.85094, 70.87396, 40.06766, 7.129991, 78.48008, 75.21636, - 93.623604, 95.95479, 29.571129, 22.721554, 26.73875, 52.075504, 56.783104, - 94.65493, 61.778534, 85.72401, 85.369514, 29.922367, 41.410553, 94.12884, - 80.276855, 55.604828, 54.70947, 74.07216, 44.61955, 31.38113, 68.48596, 34.56782, - 14.424729, 48.204506, 9.675444, 32.01946, 92.32695, 36.292683, 78.31955, 98.05327, - 14.343918, 46.017002, 95.90888, 82.63626, 16.873539, 3.698051, 7.8042626, - 64.194405, 96.71023, 67.93692, 21.618402, 51.92182, 22.834194, 61.56986, 19.749891, - 55.31206, 38.29552, 67.57593, 67.145836, 38.92673, 94.95708, 72.38746, 90.70901, - 69.43995, 9.394085, 31.646872, 88.20112, 9.134722, 99.98214, 5.423498, 41.51995, - 76.94409, 77.373276, 3.2966614, 9.611201, 57.231106, 30.747868, 76.10228, 91.98308, - 70.893585, 0.9067178, 43.96515, 16.321218, 27.734184, 83.271835, 88.23312, - 87.16445, 5.556643, 15.627432, 58.547127, 93.6459, 40.539192, 49.124157, 91.13276, - 57.485855, 8.827019, 4.9690843, 46.511234, 53.91469, 97.71925, 20.135271, - 23.353004, 70.92099, 93.38748, 87.520134, 51.684677, 29.89813, 9.110392, 65.809204, - 34.16554, 93.398605, 84.58669, 96.409645, 9.876037, 94.767784, 99.21523, 1.9330144, - 94.92429, 75.12728, 17.218828, 97.89164, 35.476578, 77.629456, 69.573746, - 40.200542, 42.117836, 5.861628, 75.45282, 82.73633, 0.98086596, 77.24894, - 11.248695, 61.070026, 52.692616, 80.5449, 80.76036, 29.270136, 67.60252, 48.782394, - 95.18851, 83.47162, 52.068756, 46.66002, 90.12216, 15.515327, 33.694042, 96.963036, - 73.49627, 62.805485, 44.715607, 59.98627, 3.8921833, 37.565327, 29.69184, - 39.429665, 83.46899, 44.286453, 21.54851, 56.096413, 18.169249, 5.214751, - 14.691341, 99.779335, 26.32643, 67.69903, 36.41243, 67.27333, 12.157213, 96.18984, - 2.438283, 78.14289, 0.14715195, 98.769, 53.649532, 21.615898, 39.657497, 95.45616, - 18.578386, 71.47976, 22.348118, 17.85519, 6.3717127, 62.176777, 22.033644, - 23.178005, 79.44858, 89.70233, 37.21273, 71.86182, 21.284317, 52.908623, 30.095518, - 63.64478, 77.55823, 80.04871, 15.133011, 30.439043, 70.16561, 4.4014096, 89.28944, - 26.29093, 46.827854, 11.764729, 61.887516, 47.774887, 57.19503, 59.444664, - 28.592825, 98.70386, 1.2497544, 82.28431, 46.76423, 83.746124, 53.032673, 86.53457, - 99.42168, 90.184, 92.27852, 9.059965, 71.75723, 70.45299, 10.924053, 68.329704, - 77.27232, 6.677854, 75.63629, 57.370533, 17.09031, 10.554659, 99.56178, 37.53221, - 72.311104, 75.7565, 65.2042, 36.096478, 64.69502, 38.88497, 64.33723, 84.87812, - 66.84958, 8.508932, 79.134, 83.431015, 66.72124, 61.801838, 64.30524, 37.194263, - 77.94725, 89.705185, 23.643505, 19.505919, 48.40264, 43.01083, 21.171177, - 18.717121, 10.805857, 69.66983, 77.85261, 57.323063, 3.28964, 38.758026, 5.349946, - 7.46572, 57.485138, 30.822384, 33.9411, 95.53746, 65.57723, 42.1077, 28.591347, - 11.917269, 5.031073, 31.835615, 19.34116, 85.71027, 87.4516, 1.3798475, 70.70583, - 51.988052, 45.217144, 14.308596, 54.557167, 86.18323, 79.13666, 76.866745, - 46.010685, 79.739235, 44.667603, 39.36416, 72.605896, 73.83187, 13.137412, - 6.7911267, 63.952374, 10.082436, 86.00318, 99.760376, 92.84948, 63.786434, - 3.4429908, 18.244314, 75.65299, 14.964747, 70.126366, 80.89449, 91.266655, - 96.58798, 46.439327, 38.253975, 87.31036, 21.093178, 37.19671, 58.28973, 9.75231, - 12.350321, 25.75115, 87.65073, 53.610504, 36.850048, 18.66356, 94.48941, 83.71898, - 44.49315, 44.186737, 19.360733, 84.365974, 46.76272, 44.924366, 50.279808, - 54.868866, 91.33004, 18.683397, 75.13282, 15.070831, 47.04839, 53.780903, - 26.911152, 74.65651, 57.659935, 25.604189, 37.235474, 65.39667, 53.952206, - 40.37131, 59.173275, 96.00756, 54.591274, 10.787476, 69.51549, 31.970142, - 25.408005, 55.972492, 85.01888, 97.48981, 91.006134, 28.98619, 97.151276, - 34.388496, 47.498177, 11.985874, 64.73775, 33.877014, 13.370312, 34.79146, - 86.19321, 15.019405, 94.07832, 93.50433, 60.168625, 50.95409, 38.27827, 47.458614, - 32.83715, 69.54998, 69.0361, 84.1418, 34.270298, 74.23852, 70.707466, 78.59845, - 9.651399, 24.186779, 58.255756, 53.72362, 92.46477, 97.75528, 20.257462, 30.122698, - 50.41517, 28.156603, 42.644154, - ], - }); - - let distance = compare::(256, Metric::L2, &two_vec.v); - - assert_eq!(distance, 429141.2); - } - - fn compare(dim: usize, metric: Metric, v: &[f32]) -> f32 - where - for<'a> [T; N]: FullPrecisionDistance, - { - let a_ptr = v.as_ptr(); - let b_ptr = unsafe { a_ptr.add(dim) }; - - let a_ref = - <&[f32; N]>::try_from(unsafe { std::slice::from_raw_parts(a_ptr, dim) }).unwrap(); - let b_ref = - <&[f32; N]>::try_from(unsafe { std::slice::from_raw_parts(b_ptr, dim) }).unwrap(); - - <[f32; N]>::distance_compare(a_ref, b_ref, metric) - } -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance_test.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance_test.rs deleted file mode 100644 index 0def026..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/distance_test.rs +++ /dev/null @@ -1,152 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[cfg(test)] -mod e2e_test { - - #[repr(C, align(32))] - pub struct F32Slice104([f32; 104]); - - #[repr(C, align(32))] - pub struct F16Slice104([Half; 104]); - - use approx::assert_abs_diff_eq; - - use crate::half::Half; - use crate::l2_float_distance::{distance_l2_vector_f16, distance_l2_vector_f32}; - - fn no_vector_compare_f32(a: &[f32], b: &[f32]) -> f32 { - let mut sum = 0.0; - for i in 0..a.len() { - let a_f32 = a[i]; - let b_f32 = b[i]; - let diff = a_f32 - b_f32; - sum += diff * diff; - } - sum - } - - fn no_vector_compare(a: &[Half], b: &[Half]) -> f32 { - let mut sum = 0.0; - for i in 0..a.len() { - let a_f32 = a[i].to_f32(); - let b_f32 = b[i].to_f32(); - let diff = a_f32 - b_f32; - sum += diff * diff; - } - sum - } - - #[test] - fn avx2_matches_novector() { - for i in 1..3 { - let (f1, f2) = get_test_data(0, i); - - let distance_f32x8 = distance_l2_vector_f32::<104>(&f1.0, &f2.0); - let distance = no_vector_compare_f32(&f1.0, &f2.0); - - assert_abs_diff_eq!(distance, distance_f32x8, epsilon = 1e-6); - } - } - - #[test] - fn avx2_matches_novector_random() { - let (f1, f2) = get_test_data_random(); - - let distance_f32x8 = distance_l2_vector_f32::<104>(&f1.0, &f2.0); - let distance = no_vector_compare_f32(&f1.0, &f2.0); - - assert_abs_diff_eq!(distance, distance_f32x8, epsilon = 1e-4); - } - - #[test] - fn avx_f16_matches_novector() { - for i in 1..3 { - let (f1, f2) = get_test_data_f16(0, i); - let _a_slice = f1.0.map(|x| x.to_f32().to_string()).join(", "); - let _b_slice = f2.0.map(|x| x.to_f32().to_string()).join(", "); - - let expected = no_vector_compare(f1.0[0..].as_ref(), f2.0[0..].as_ref()); - let distance_f16x8 = distance_l2_vector_f16::<104>(&f1.0, &f2.0); - - assert_abs_diff_eq!(distance_f16x8, expected, epsilon = 1e-4); - } - } - - #[test] - fn avx_f16_matches_novector_random() { - let (f1, f2) = get_test_data_f16_random(); - - let expected = no_vector_compare(f1.0[0..].as_ref(), f2.0[0..].as_ref()); - let distance_f16x8 = distance_l2_vector_f16::<104>(&f1.0, &f2.0); - - assert_abs_diff_eq!(distance_f16x8, expected, epsilon = 1e-4); - } - - fn get_test_data_f16(i1: usize, i2: usize) -> (F16Slice104, F16Slice104) { - let (a_slice, b_slice) = get_test_data(i1, i2); - let a_data = a_slice.0.iter().map(|x| Half::from_f32(*x)); - let b_data = b_slice.0.iter().map(|x| Half::from_f32(*x)); - - ( - F16Slice104(a_data.collect::>().try_into().unwrap()), - F16Slice104(b_data.collect::>().try_into().unwrap()), - ) - } - - fn get_test_data(i1: usize, i2: usize) -> (F32Slice104, F32Slice104) { - use base64::{engine::general_purpose, Engine as _}; - - let b64 = general_purpose::STANDARD.decode(TEST_DATA).unwrap(); - - let decoded: Vec> = bincode::deserialize(&b64).unwrap(); - debug_assert!(decoded.len() > i1); - debug_assert!(decoded.len() > i2); - - let mut f1 = F32Slice104([0.0; 104]); - let v1 = &decoded[i1]; - debug_assert!(v1.len() == 104); - f1.0.copy_from_slice(v1); - - let mut f2 = F32Slice104([0.0; 104]); - let v2 = &decoded[i2]; - debug_assert!(v2.len() == 104); - f2.0.copy_from_slice(v2); - - (f1, f2) - } - - fn get_test_data_f16_random() -> (F16Slice104, F16Slice104) { - let (a_slice, b_slice) = get_test_data_random(); - let a_data = a_slice.0.iter().map(|x| Half::from_f32(*x)); - let b_data = b_slice.0.iter().map(|x| Half::from_f32(*x)); - - ( - F16Slice104(a_data.collect::>().try_into().unwrap()), - F16Slice104(b_data.collect::>().try_into().unwrap()), - ) - } - - fn get_test_data_random() -> (F32Slice104, F32Slice104) { - use rand::Rng; - - let mut rng = rand::thread_rng(); - let mut f1 = F32Slice104([0.0; 104]); - - for i in 0..104 { - f1.0[i] = rng.gen_range(-1.0..1.0); - } - - let mut f2 = F32Slice104([0.0; 104]); - - for i in 0..104 { - f2.0[i] = rng.gen_range(-1.0..1.0); - } - - (f1, f2) - } - - const TEST_DATA: &str = "BQAAAAAAAABoAAAAAAAAAPz3Dj7+VgG9z/DDvQkgiT2GryK+nwS4PTeBorz4jpk9ELEqPKKeX73zZrA9uAlRvSqpKT7Gft28LsTuO8XOHL6/lCg+pW/6vJhM7j1fInU+yaSTPC2AAb5T25M8o2YTvWgEAz00cnq8xcUlPPvnBb2AGfk9UmhCvbdUJzwH4jK9UH7Lvdklhz3SoEa+NwsIvt2yYb4q7JA8d4fVvfX/kbtDOJe9boXevbw2CT7n62A9B6hOPlfeNz7CO169vnjcvR3pDz6KZxC+XR/2vTd9PTx7YY492FF2PekiGDt3OSw9IIlGPQooMj5DZcY8EgQgvpg9572paca91GQTPoWpFr7U+t697YAQPYHUXr1d8ow8AQE7PFo6JD3tt+I96ahxvYuvlD3+IW29N4Jtu2/01Ltvvg2+dja+vI8uazvITZO9mXhavpfJ6T2tB8S7OKT3PWWjpj0Mjty9advIPFgucTp3JO69CI6YPaWoDD5pwim9rjUovh2qgr3R/lq+nUi3PI+acL041o081D8lvRCJLTwAAAAAAAAAAAAAAAAAAAAAaAAAAAAAAAA6pJO94NE1voDn+rzQ8CY+1rxkvtspaz0xTPw7+0GMvC0ZgbyWwdy8zHcovKdvdb70BLC8DtHKvdK6vz0R9Ys7vBWyvZK1LL0ehYM9aV+JveuvoD2ilvo9NLJ4vbRnPT4MXAW+BhG4POOBaD0Vz5I9s1+1vTUdHb7Kjcw9uVUJvdbgoj3TbBe8WwPSvYoBBj4m6c+9xTXTvVTDaL28+Ac9KtA0Pa3tS73Vq5S8fNLkvf/Gir0yILy9ZYR3vvUdUD2ZB5W9rHI4PXS76L070oG9EsjYPb89S75pz7Q9xFKyvZ5ECT0kDSU+l4AQPsQVqzyq/LW95ZCZPC6nQj0VIBa9XwkhPr1gy72c7mw937XXvQ76ur3sRok9mCUqPXHvgj28jV89LZN8O0eH0T0KMdq9ZzXevYbmPr0fcac8r7j3vYmKCL4Sewm+iLtRviuOjz08XbE9LlYevDI1wz0s7z278oVJvtpjrT20IEU9+mTtvBjMQz1H9Ey+LQEXva1Rwrxmyts9sf1hPRY3xL3RdRU+AAAAAAAAAAAAAAAAAAAAAGgAAAAAAAAARqSTvbYJpLx1x869cW67PeeJhb7/cBu9m0eFPQO3oL0I+L49YQDavTYSez3SmTg96hBGPuh4oL2x2ow6WdCUO6XUSz4xcU88GReAvVfekj0Ph3Y9z43hvBzT5z1I2my9UVy3vAj8jL08Gtm9CfJcPRihTr1+8Yu9TiP+PNrJa77Dfa09IhpEPesJNr0XzFU8yye3PZKFyz3uzJ09FLRUvYq3l73X4X07DDUzvq9VXjwWtg8+JrzYPcFCkr0jDCg9T9zlvZbZjz4Y8pM89xo8PgAcfbvYSnY8XoFKvO05/L36yzE8J+5yPqfe5r2AZFq8ULRDvnkTgrw+S7q9qGYLvQDZYL1T8d09bFikvZw3+jsYLdO8H3GVveHBYT4gnsE8ZBIJPpzOEj7OSDC+ZYu+vFc1Erzko4M9GqLtPBHH5TwpeRs+miC4PBHH5Tw9Z9k9VUsUPjnppj0oC5C9mcqDvY7y1rxdvZU8PdFAPov9lz0bOmq94kdyPBBokTxtOj89fu4avSsazj1P7iE+x8YkPAAAAAAAAAAAAAAAAAAAAABoAAAAAAAAAHEruT3mgKM8JnEvvAsfHL63906+ifhgvldl1r14OeO9waUyuw3yUzx+PDW9UbDhPQP4Lb4KRRk+Oky2vaLfaT30mrA9YMeZPfzPMz4h42M+XfCHva4AGr6MOSM+iBOzvdsaE7xFxgI+gJGXvVMzE75kHY+8oAWNvVqNK7yOx589fU3lvVVPg730Cwk+DKkEPWYtxjqQ2MK9H0T+vTnGQj2yq5w8L49BvrEJrzyB4Yo9AXV7PYGCLr3MxsG9oWM7PTyu8TzEOhW+dyWrvUTxHD2nL+c9+VKFPcthhLsc0PM8FdyPPeLj/z1WAHS8ZvW2PGg4Cb5u3IU9g4CovSHW+L2CWoG++nZnPAi2ST3HmUC9P5rJuxQbU765lwU+7FLBPUPTfL0uGgk+yKy2PYwXaT1I4I+9AU6VPQ5QaDx9mdE8Qg8zPfGCUjzD/io9rr+BvTNDqT0MFNi9mHatvS1iJD0nVrK78WmIPE0QsL3PAQq9cMRgPWXmmr3yTcw9UcXrPccwa76+cBq+5iVOvUg9c70AAAAAAAAAAAAAAAAAAAAAaAAAAAAAAAB/K7k9hCsnPUJXJr2Wg4a9MEtXve33Sj0VJZ89pciEvWLqwLzUgyu8ADTGPAVenL2UZ/c96YtMved+Wr3LUro9H8a7vGTSA77C5n69Lf3pPQj4KD5cFKq9fZ0uvvYQCT7b23G9XGMCPrGuy736Z9A9kZzFPSuCSD7/9/07Y4/6POxLir3/JBS9qFKMvkSzjryPgVY+ugq8PC9yhbsXaiq+O6WfPcvFK7vZXAy+goAQvXpHHj5jwPI87eokvrySET5QoOm8h8ixOhXzKb5s8+A9sjcJPjiLAz598yQ9yCYSPq6eGz4rvjE82lvGvWuIOLx23zK9hHg8vTWOv70/Tse81fA6Pr2wNz34Eza+2Uj3PZ3trr0aXAI9PCkKPiybe721P9U9QkNLO927jT3LpRA+mpJUvUeU6rwC/Qa+lr4Cvgrpnj1pQ/i9TxhSvJqYr72RS6y8aQLTPQzPiz3vSRY94NfrPJl6LL2adjO8iYfPuhRzZz2f7R8+iVskPcUeXr12ZiI+nd3xvIYv8bwqYlg+AAAAAAAAAAAAAAAAAAAAAA=="; -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/half.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/half.rs deleted file mode 100644 index 87d7df6..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/half.rs +++ /dev/null @@ -1,82 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use bytemuck::{Pod, Zeroable}; -use half::f16; -use std::convert::AsRef; -use std::fmt; - -// Define the Half type as a new type over f16. -// the memory layout of the Half struct will be the same as the memory layout of the f16 type itself. -// The Half struct serves as a simple wrapper around the f16 type and does not introduce any additional memory overhead. -// Test function: -// use half::f16; -// pub struct Half(f16); -// fn main() { -// let size_of_half = std::mem::size_of::(); -// let alignment_of_half = std::mem::align_of::(); -// println!("Size of Half: {} bytes", size_of_half); -// println!("Alignment of Half: {} bytes", alignment_of_half); -// } -// Output: -// Size of Half: 2 bytes -// Alignment of Half: 2 bytes -pub struct Half(f16); - -unsafe impl Pod for Half {} -unsafe impl Zeroable for Half {} - -// Implement From for Half -impl From for f32 { - fn from(val: Half) -> Self { - val.0.to_f32() - } -} - -// Implement AsRef for Half so that it can be used in distance_compare. -impl AsRef for Half { - fn as_ref(&self) -> &f16 { - &self.0 - } -} - -// Implement From for Half. -impl Half { - pub fn from_f32(value: f32) -> Self { - Self(f16::from_f32(value)) - } -} - -// Implement Default for Half. -impl Default for Half { - fn default() -> Self { - Self(f16::from_f32(Default::default())) - } -} - -// Implement Clone for Half. -impl Clone for Half { - fn clone(&self) -> Self { - Half(self.0) - } -} - -// Implement PartialEq for Half. -impl fmt::Debug for Half { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Half({:?})", self.0) - } -} - -impl Copy for Half {} - -impl Half { - pub fn to_f32(&self) -> f32 { - self.0.to_f32() - } -} - -unsafe impl Send for Half {} -unsafe impl Sync for Half {} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/l2_float_distance.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/l2_float_distance.rs deleted file mode 100644 index b818899..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/l2_float_distance.rs +++ /dev/null @@ -1,78 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] - -//! Distance calculation for L2 Metric - -#[cfg(not(target_feature = "avx2"))] -compile_error!("Library must be compiled with -C target-feature=+avx2"); - -use std::arch::x86_64::*; - -use crate::Half; - -/// Calculate the distance by vector arithmetic -#[inline(never)] -pub fn distance_l2_vector_f16(a: &[Half; N], b: &[Half; N]) -> f32 { - debug_assert_eq!(N % 8, 0); - - // make sure the addresses are bytes aligned - debug_assert_eq!(a.as_ptr().align_offset(32), 0); - debug_assert_eq!(b.as_ptr().align_offset(32), 0); - - unsafe { - let mut sum = _mm256_setzero_ps(); - let a_ptr = a.as_ptr() as *const __m128i; - let b_ptr = b.as_ptr() as *const __m128i; - - // Iterate over the elements in steps of 8 - for i in (0..N).step_by(8) { - let a_vec = _mm256_cvtph_ps(_mm_load_si128(a_ptr.add(i / 8))); - let b_vec = _mm256_cvtph_ps(_mm_load_si128(b_ptr.add(i / 8))); - - let diff = _mm256_sub_ps(a_vec, b_vec); - sum = _mm256_fmadd_ps(diff, diff, sum); - } - - let x128: __m128 = _mm_add_ps(_mm256_extractf128_ps(sum, 1), _mm256_castps256_ps128(sum)); - /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */ - let x64: __m128 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); - /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */ - let x32: __m128 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); - /* Conversion to float is a no-op on x86-64 */ - _mm_cvtss_f32(x32) - } -} - -/// Calculate the distance by vector arithmetic -#[inline(never)] -pub fn distance_l2_vector_f32(a: &[f32; N], b: &[f32; N]) -> f32 { - debug_assert_eq!(N % 8, 0); - - // make sure the addresses are bytes aligned - debug_assert_eq!(a.as_ptr().align_offset(32), 0); - debug_assert_eq!(b.as_ptr().align_offset(32), 0); - - unsafe { - let mut sum = _mm256_setzero_ps(); - - // Iterate over the elements in steps of 8 - for i in (0..N).step_by(8) { - let a_vec = _mm256_load_ps(&a[i]); - let b_vec = _mm256_load_ps(&b[i]); - let diff = _mm256_sub_ps(a_vec, b_vec); - sum = _mm256_fmadd_ps(diff, diff, sum); - } - - let x128: __m128 = _mm_add_ps(_mm256_extractf128_ps(sum, 1), _mm256_castps256_ps128(sum)); - /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */ - let x64: __m128 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); - /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */ - let x32: __m128 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); - /* Conversion to float is a no-op on x86-64 */ - _mm_cvtss_f32(x32) - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/lib.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/lib.rs deleted file mode 100644 index d221070..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/lib.rs +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![cfg_attr( - not(test), - warn(clippy::panic, clippy::unwrap_used, clippy::expect_used) -)] - -// #![feature(stdsimd)] -// mod f32x16; -// Uncomment above 2 to experiment with f32x16 -mod distance; -mod half; -mod l2_float_distance; -mod metric; -mod utils; - -pub use crate::half::Half; -pub use distance::FullPrecisionDistance; -pub use metric::Metric; -pub use utils::prefetch_vector; - -#[cfg(test)] -mod distance_test; -mod test_util; diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/metric.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/metric.rs deleted file mode 100644 index c60ef29..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/metric.rs +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#![warn(missing_debug_implementations, missing_docs)] -use std::str::FromStr; - -/// Distance metric -#[derive(Debug, PartialEq, Eq, Clone, Copy)] -pub enum Metric { - /// Squared Euclidean (L2-Squared) - L2, - - /// Cosine similarity - /// TODO: T should be float for Cosine distance - Cosine, -} - -#[derive(thiserror::Error, Debug)] -pub enum ParseMetricError { - #[error("Invalid format for Metric: {0}")] - InvalidFormat(String), -} - -impl FromStr for Metric { - type Err = ParseMetricError; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "l2" => Ok(Metric::L2), - "cosine" => Ok(Metric::Cosine), - _ => Err(ParseMetricError::InvalidFormat(String::from(s))), - } - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/test_util.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/test_util.rs deleted file mode 100644 index 7cfc929..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/test_util.rs +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -#[cfg(test)] -use crate::Half; - -#[cfg(test)] -pub fn no_vector_compare_f16(a: &[Half], b: &[Half]) -> f32 { - let mut sum = 0.0; - debug_assert_eq!(a.len(), b.len()); - - for i in 0..a.len() { - sum += (a[i].to_f32() - b[i].to_f32()).powi(2); - } - sum -} - -#[cfg(test)] -pub fn no_vector_compare_f32(a: &[f32], b: &[f32]) -> f32 { - let mut sum = 0.0; - debug_assert_eq!(a.len(), b.len()); - - for i in 0..a.len() { - sum += (a[i] - b[i]).powi(2); - } - sum -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/utils.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/utils.rs deleted file mode 100644 index a61c99a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector/src/utils.rs +++ /dev/null @@ -1,21 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::arch::x86_64::{_mm_prefetch, _MM_HINT_T0}; - -/// Prefetch the given vector in chunks of 64 bytes, which is a cache line size -/// NOTE: good efficiency when total_vec_size is integral multiple of 64 -#[inline] -pub fn prefetch_vector(vec: &[T]) { - let vec_ptr = vec.as_ptr() as *const i8; - let vecsize = std::mem::size_of_val(vec); - let max_prefetch_size = (vecsize / 64) * 64; - - for d in (0..max_prefetch_size).step_by(64) { - unsafe { - _mm_prefetch(vec_ptr.add(d), _MM_HINT_T0); - } - } -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/Cargo.toml b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/Cargo.toml deleted file mode 100644 index 6f50ad9..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/Cargo.toml +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. -[package] -name = "vector_base64" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -base64 = "0.21.2" -bincode = "1.3.3" -half = "2.2.1" -serde = "1.0.163" - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/src/main.rs b/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/src/main.rs deleted file mode 100644 index 2867436..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/rust/vector_base64/src/main.rs +++ /dev/null @@ -1,82 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT license. - */ -use std::fs::File; -use std::io::{self, BufReader, Read}; -use std::{env, vec}; - -fn main() -> io::Result<()> { - // Retrieve command-line arguments - let args: Vec = env::args().collect(); - - // Check if the correct number of arguments is provided - if args.len() != 4 { - print_usage(); - return Ok(()); - } - - // Retrieve the input and output file paths from the arguments - let input_file_path = &args[1]; - let item_count: usize = args[2].parse::().unwrap(); - let return_dimension: usize = args[3].parse::().unwrap(); - - // Open the input file for reading - let mut input_file = BufReader::new(File::open(input_file_path)?); - - // Read the first 8 bytes as metadata - let mut metadata = [0; 8]; - input_file.read_exact(&mut metadata)?; - - // Extract the number of points and dimension from the metadata - let _ = i32::from_le_bytes(metadata[..4].try_into().unwrap()); - let mut dimension: usize = (i32::from_le_bytes(metadata[4..].try_into().unwrap())) as usize; - if return_dimension < dimension { - dimension = return_dimension; - } - - let mut float_array = Vec::>::with_capacity(item_count); - - // Process each data point - for _ in 0..item_count { - // Read one data point from the input file - let mut buffer = vec![0; dimension * std::mem::size_of::()]; - match input_file.read_exact(&mut buffer) { - Ok(()) => { - let mut float_data = buffer - .chunks_exact(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect::>(); - - let mut i = return_dimension; - while i > dimension { - float_data.push(0.0); - i -= 1; - } - - float_array.push(float_data); - } - Err(err) => { - println!("Error: {}", err); - break; - } - } - } - - use base64::{engine::general_purpose, Engine as _}; - - let encoded: Vec = bincode::serialize(&float_array).unwrap(); - let b64 = general_purpose::STANDARD.encode(encoded); - println!("Float {}", b64); - - Ok(()) -} - -/// Prints the usage information -fn print_usage() { - println!("Usage: program_name input_file "); - println!( - "Itemcount is the number of items to convert. Expand to dimension if provided is smaller" - ); -} - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/setup.py b/packages/leann-backend-diskann/third_party/DiskANN/setup.py deleted file mode 100644 index 01184f8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/setup.py +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -import os -import re -import shutil -import subprocess -import sys -from pathlib import Path - -from setuptools import Extension, setup -from setuptools.command.build_ext import build_ext -from setuptools.command.install_lib import install_lib - -# Convert distutils Windows platform specifiers to CMake -A arguments -PLAT_TO_CMAKE = { - "win-amd64": "x64" -} - - -class CMakeExtension(Extension): - def __init__(self, name: str, sourcedir: str = "") -> None: - super().__init__(name, sources=[]) - self.sourcedir = os.fspath(Path(sourcedir).resolve()) - - -class CMakeBuild(build_ext): - def build_extension(self, ext: CMakeExtension) -> None: - # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ - ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] - extdir = ext_fullpath.parent.resolve() - # Using this requires trailing slash for auto-detection & inclusion of - # auxiliary "native" libs - - debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug - cfg = "Debug" if debug else "Release" - - # CMake lets you override the generator - we need to check this. - # Can be set with Conda-Build, for example. - cmake_generator = os.environ.get("CMAKE_GENERATOR", "") - - # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON - # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code - # from Python. - cmake_args = [ - f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}", - f"-DPYTHON_EXECUTABLE={sys.executable}", - f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm - f"-DVERSION_INFO={self.distribution.get_version()}" # commented out, we want this set in the CMake file - ] - build_args = [] - # Adding CMake arguments set as environment variable - # (needed e.g. to build for ARM OSx on conda-forge) - if "CMAKE_ARGS" in os.environ: - cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] - - # In this example, we pass in the version to C++. You might not need to. - # cmake_args += [f"-DVERSION_INFO={self.distribution.get_version()}"] # type: ignore[attr-defined] - - if self.compiler.compiler_type != "msvc": - # Using Ninja-build since it a) is available as a wheel and b) - # multithreads automatically. MSVC would require all variables be - # exported for Ninja to pick it up, which is a little tricky to do. - # Users can override the generator with CMAKE_GENERATOR in CMake - # 3.15+. - if not cmake_generator or cmake_generator == "Ninja": - try: - import ninja # noqa: F401 - - ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" - cmake_args += [ - "-GNinja", - f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", - ] - except ImportError: - pass - - else: - - # Single config generators are handled "normally" - single_config = any(x in cmake_generator for x in {"NMake", "Ninja"}) - - # CMake allows an arch-in-generator style for backward compatibility - contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"}) - - # Specify the arch if using MSVC generator, but only if it doesn't - # contain a backward-compatibility arch spec already in the - # generator name. - if not single_config and not contains_arch: - cmake_args += ["-A", PLAT_TO_CMAKE[self.plat_name]] - - # Multi-config generators have a different way to specify configs - if not single_config: - cmake_args += [ - f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{cfg.upper()}={extdir}" - ] - build_args += ["--config", cfg] - - if sys.platform.startswith("darwin"): - # Cross-compile support for macOS - respect ARCHFLAGS if set - archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", "")) - if archs: - cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))] - - # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level - # across all generators. - if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ: - # self.parallel is a Python 3 only way to set parallel jobs by hand - # using -j in the build_ext call, not supported by pip or PyPA-build. - if hasattr(self, "parallel") and self.parallel: - # CMake 3.12+ only. - build_args += [f"-j{self.parallel}"] - - build_temp = Path(self.build_temp) / ext.name - if not build_temp.exists(): - build_temp.mkdir(parents=True) - - # this next line is problematic. we tell it to use the ext.sourcedir but, when - # using `python -m build`, we actually have a copy of everything made and pushed - # into a venv isolation area - if os.environ.get("USE_CONDA", "") == '1' and os.environ.get("CONDA_PREFIX", "") != "": - subprocess.run( - ["cmake", "-DPYBIND=True", "-DCMAKE_PREFIX_PATH=" + os.environ.get("CONDA_PREFIX", ""), - "-DProtobuf_DIR=" + os.path.join(os.environ.get("CONDA_PREFIX", ""), "lib/cmake/protobuf"), - ext.sourcedir] + cmake_args, cwd=build_temp, check=True - ) - else: - subprocess.run( - ["cmake", "-DPYBIND=True", ext.sourcedir] + cmake_args, cwd=build_temp, check=True - ) - - subprocess.run( - ["cmake", "--build", "."] + build_args, cwd=build_temp, check=True - ) - - -class InstallCMakeLibs(install_lib): - def run(self): - """ - Windows only copy from the x64/Release directory and place them in the package - """ - - self.announce("Moving library files", level=3) - - self.skip_build = True - - # we only need to move the windows build output - windows_build_output_dir = Path('.') / 'x64' / 'Release' - - if windows_build_output_dir.exists(): - libs = [ - os.path.join(windows_build_output_dir, _lib) for _lib in - os.listdir(windows_build_output_dir) if - os.path.isfile(os.path.join(windows_build_output_dir, _lib)) and - os.path.splitext(_lib)[1] in [".dll", '.lib', '.pyd', '.exp'] - ] - - for lib in libs: - shutil.move( - lib, - os.path.join(self.build_dir, 'diskannpy', os.path.basename(lib)) - ) - - super().run() - - -setup( - ext_modules=[CMakeExtension("diskannpy._diskannpy", ".")], - cmdclass={ - "build_ext": CMakeBuild, - 'install_lib': InstallCMakeLibs - }, - zip_safe=False, - package_dir={"diskannpy": "python/src"}, - exclude_package_data={"diskannpy": ["diskann_bindings.cpp"]} -) diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/src/CMakeLists.txt deleted file mode 100644 index 97b00c7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/CMakeLists.txt +++ /dev/null @@ -1,26 +0,0 @@ -#Copyright(c) Microsoft Corporation.All rights reserved. -#Licensed under the MIT license. - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_COMPILE_WARNING_AS_ERROR ON) - -if(MSVC) - add_subdirectory(dll) -else() - #file(GLOB CPP_SOURCES *.cpp) - set(CPP_SOURCES abstract_data_store.cpp ann_exception.cpp apple_aligned_file_reader.cpp disk_utils.cpp - distance.cpp index.cpp in_mem_graph_store.cpp in_mem_data_store.cpp - linux_aligned_file_reader.cpp math_utils.cpp natural_number_map.cpp - in_mem_data_store.cpp in_mem_graph_store.cpp - natural_number_set.cpp memory_mapper.cpp partition.cpp pq.cpp - pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp filter_utils.cpp index_factory.cpp abstract_index.cpp pq_l2_distance.cpp pq_data_store.cpp) - if (RESTAPI) - list(APPEND CPP_SOURCES restapi/search_wrapper.cpp restapi/server.cpp) - endif() - add_library(${PROJECT_NAME} ${CPP_SOURCES}) - add_library(${PROJECT_NAME}_s STATIC ${CPP_SOURCES}) -endif() - -if (NOT MSVC) - install(TARGETS ${PROJECT_NAME} LIBRARY) -endif() diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_data_store.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_data_store.cpp deleted file mode 100644 index 0cff015..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_data_store.cpp +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "abstract_data_store.h" - -namespace diskann -{ - -template -AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) - : _capacity(capacity), _dim(dim) -{ -} - -template location_t AbstractDataStore::capacity() const -{ - return _capacity; -} - -template size_t AbstractDataStore::get_dims() const -{ - return _dim; -} - -template location_t AbstractDataStore::resize(const location_t new_num_points) -{ - if (new_num_points > _capacity) - { - return expand(new_num_points); - } - else if (new_num_points < _capacity) - { - return shrink(new_num_points); - } - else - { - return _capacity; - } -} - -template DISKANN_DLLEXPORT class AbstractDataStore; -template DISKANN_DLLEXPORT class AbstractDataStore; -template DISKANN_DLLEXPORT class AbstractDataStore; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_index.cpp deleted file mode 100644 index 9266582..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/abstract_index.cpp +++ /dev/null @@ -1,334 +0,0 @@ -#include "common_includes.h" -#include "windows_customizations.h" -#include "abstract_index.h" - -namespace diskann -{ - -template -void AbstractIndex::build(const data_type *data, const size_t num_points_to_load, const std::vector &tags) -{ - auto any_data = std::any(data); - auto any_tags_vec = TagVector(tags); - this->_build(any_data, num_points_to_load, any_tags_vec); -} - -template -std::pair AbstractIndex::search(const data_type *query, const size_t K, const uint32_t L, - IDType *indices, float *distances) -{ - auto any_indices = std::any(indices); - auto any_query = std::any(query); - return _search(any_query, K, L, any_indices, distances); -} - -template -size_t AbstractIndex::search_with_tags(const data_type *query, const uint64_t K, const uint32_t L, tag_type *tags, - float *distances, std::vector &res_vectors, bool use_filters, - const std::string filter_label) -{ - auto any_query = std::any(query); - auto any_tags = std::any(tags); - auto any_res_vectors = DataVector(res_vectors); - return this->_search_with_tags(any_query, K, L, any_tags, distances, any_res_vectors, use_filters, filter_label); -} - -template -std::pair AbstractIndex::search_with_filters(const DataType &query, const std::string &raw_label, - const size_t K, const uint32_t L, IndexType *indices, - float *distances) -{ - auto any_indices = std::any(indices); - return _search_with_filters(query, raw_label, K, L, any_indices, distances); -} - -template -void AbstractIndex::search_with_optimized_layout(const data_type *query, size_t K, size_t L, uint32_t *indices) -{ - auto any_query = std::any(query); - this->_search_with_optimized_layout(any_query, K, L, indices); -} - -template -int AbstractIndex::insert_point(const data_type *point, const tag_type tag) -{ - auto any_point = std::any(point); - auto any_tag = std::any(tag); - return this->_insert_point(any_point, any_tag); -} - -template -int AbstractIndex::insert_point(const data_type *point, const tag_type tag, const std::vector &labels) -{ - auto any_point = std::any(point); - auto any_tag = std::any(tag); - auto any_labels = Labelvector(labels); - return this->_insert_point(any_point, any_tag, any_labels); -} - -template int AbstractIndex::lazy_delete(const tag_type &tag) -{ - auto any_tag = std::any(tag); - return this->_lazy_delete(any_tag); -} - -template -void AbstractIndex::lazy_delete(const std::vector &tags, std::vector &failed_tags) -{ - auto any_tags = TagVector(tags); - auto any_failed_tags = TagVector(failed_tags); - this->_lazy_delete(any_tags, any_failed_tags); -} - -template void AbstractIndex::get_active_tags(tsl::robin_set &active_tags) -{ - auto any_active_tags = TagRobinSet(active_tags); - this->_get_active_tags(any_active_tags); -} - -template void AbstractIndex::set_start_points_at_random(data_type radius, uint32_t random_seed) -{ - auto any_radius = std::any(radius); - this->_set_start_points_at_random(any_radius, random_seed); -} - -template int AbstractIndex::get_vector_by_tag(tag_type &tag, data_type *vec) -{ - auto any_tag = std::any(tag); - auto any_data_ptr = std::any(vec); - return this->_get_vector_by_tag(any_tag, any_data_ptr); -} - -template void AbstractIndex::set_universal_label(const label_type universal_label) -{ - auto any_label = std::any(universal_label); - this->_set_universal_label(any_label); -} - -// exports -template DISKANN_DLLEXPORT void AbstractIndex::build(const float *data, const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const int8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const uint8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const float *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const int8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const uint8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const float *data, const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const int8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const uint8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const float *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const int8_t *data, - const size_t num_points_to_load, - const std::vector &tags); -template DISKANN_DLLEXPORT void AbstractIndex::build(const uint8_t *data, - const size_t num_points_to_load, - const std::vector &tags); - -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair AbstractIndex::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair AbstractIndex::search_with_filters( - const DataType &query, const std::string &raw_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); - -template DISKANN_DLLEXPORT std::pair AbstractIndex::search_with_filters( - const DataType &query, const std::string &raw_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const float *query, const uint64_t K, const uint32_t L, int32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const uint8_t *query, const uint64_t K, const uint32_t L, int32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const int8_t *query, const uint64_t K, const uint32_t L, int32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const float *query, const uint64_t K, const uint32_t L, uint32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const uint8_t *query, const uint64_t K, const uint32_t L, uint32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const int8_t *query, const uint64_t K, const uint32_t L, uint32_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const float *query, const uint64_t K, const uint32_t L, int64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const uint8_t *query, const uint64_t K, const uint32_t L, int64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const int8_t *query, const uint64_t K, const uint32_t L, int64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const float *query, const uint64_t K, const uint32_t L, uint64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const uint8_t *query, const uint64_t K, const uint32_t L, uint64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT size_t AbstractIndex::search_with_tags( - const int8_t *query, const uint64_t K, const uint32_t L, uint64_t *tags, float *distances, - std::vector &res_vectors, bool use_filters, const std::string filter_label); - -template DISKANN_DLLEXPORT void AbstractIndex::search_with_optimized_layout(const float *query, size_t K, - size_t L, uint32_t *indices); -template DISKANN_DLLEXPORT void AbstractIndex::search_with_optimized_layout(const uint8_t *query, size_t K, - size_t L, uint32_t *indices); -template DISKANN_DLLEXPORT void AbstractIndex::search_with_optimized_layout(const int8_t *query, size_t K, - size_t L, uint32_t *indices); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const float *point, const int32_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const uint8_t *point, const int32_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const int8_t *point, const int32_t tag); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const float *point, const uint32_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const uint8_t *point, const uint32_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const int8_t *point, const uint32_t tag); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const float *point, const int64_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const uint8_t *point, const int64_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const int8_t *point, const int64_t tag); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const float *point, const uint64_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const uint8_t *point, const uint64_t tag); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point(const int8_t *point, const uint64_t tag); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const int32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const int32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const int32_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const uint32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const uint32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const uint32_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const int64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const int64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const int64_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const uint64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const uint64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const uint64_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const int32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const int32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const int32_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const uint32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const uint32_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const uint32_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const int64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const int64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const int64_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const float *point, const uint64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const uint8_t *point, const uint64_t tag, const std::vector &labels); -template DISKANN_DLLEXPORT int AbstractIndex::insert_point( - const int8_t *point, const uint64_t tag, const std::vector &labels); - -template DISKANN_DLLEXPORT int AbstractIndex::lazy_delete(const int32_t &tag); -template DISKANN_DLLEXPORT int AbstractIndex::lazy_delete(const uint32_t &tag); -template DISKANN_DLLEXPORT int AbstractIndex::lazy_delete(const int64_t &tag); -template DISKANN_DLLEXPORT int AbstractIndex::lazy_delete(const uint64_t &tag); - -template DISKANN_DLLEXPORT void AbstractIndex::lazy_delete(const std::vector &tags, - std::vector &failed_tags); -template DISKANN_DLLEXPORT void AbstractIndex::lazy_delete(const std::vector &tags, - std::vector &failed_tags); -template DISKANN_DLLEXPORT void AbstractIndex::lazy_delete(const std::vector &tags, - std::vector &failed_tags); -template DISKANN_DLLEXPORT void AbstractIndex::lazy_delete(const std::vector &tags, - std::vector &failed_tags); - -template DISKANN_DLLEXPORT void AbstractIndex::get_active_tags(tsl::robin_set &active_tags); -template DISKANN_DLLEXPORT void AbstractIndex::get_active_tags(tsl::robin_set &active_tags); -template DISKANN_DLLEXPORT void AbstractIndex::get_active_tags(tsl::robin_set &active_tags); -template DISKANN_DLLEXPORT void AbstractIndex::get_active_tags(tsl::robin_set &active_tags); - -template DISKANN_DLLEXPORT void AbstractIndex::set_start_points_at_random(float radius, uint32_t random_seed); -template DISKANN_DLLEXPORT void AbstractIndex::set_start_points_at_random(uint8_t radius, - uint32_t random_seed); -template DISKANN_DLLEXPORT void AbstractIndex::set_start_points_at_random(int8_t radius, uint32_t random_seed); - -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int32_t &tag, float *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int32_t &tag, uint8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int32_t &tag, int8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint32_t &tag, float *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint32_t &tag, uint8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint32_t &tag, int8_t *vec); - -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int64_t &tag, float *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int64_t &tag, uint8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(int64_t &tag, int8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint64_t &tag, float *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint64_t &tag, uint8_t *vec); -template DISKANN_DLLEXPORT int AbstractIndex::get_vector_by_tag(uint64_t &tag, int8_t *vec); - -template DISKANN_DLLEXPORT void AbstractIndex::set_universal_label(const uint16_t label); -template DISKANN_DLLEXPORT void AbstractIndex::set_universal_label(const uint32_t label); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/ann_exception.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/ann_exception.cpp deleted file mode 100644 index ba55e36..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/ann_exception.cpp +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "ann_exception.h" -#include -#include - -namespace diskann -{ -ANNException::ANNException(const std::string &message, int errorCode) - : std::runtime_error(message), _errorCode(errorCode) -{ -} - -std::string package_string(const std::string &item_name, const std::string &item_val) -{ - return std::string("[") + item_name + ": " + std::string(item_val) + std::string("]"); -} - -ANNException::ANNException(const std::string &message, int errorCode, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum) - : ANNException(package_string(std::string("FUNC"), funcSig) + package_string(std::string("FILE"), fileName) + - package_string(std::string("LINE"), std::to_string(lineNum)) + " " + message, - errorCode) -{ -} - -FileException::FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum) - : ANNException(std::string(" While opening file \'") + filename + std::string("\', error code: ") + - std::to_string(e.code().value()) + " " + e.code().message(), - e.code().value(), funcSig, fileName, lineNum) -{ -} - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/apple_aligned_file_reader.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/apple_aligned_file_reader.cpp deleted file mode 100644 index 4ef1c22..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/apple_aligned_file_reader.cpp +++ /dev/null @@ -1,383 +0,0 @@ -#include "aligned_file_reader.h" -#ifdef __APPLE__ - -#include "apple_aligned_file_reader.h" -#include "utils.h" - -#define SECTOR_LEN 4096 - -AppleAlignedFileReader::AppleAlignedFileReader() -{ - this->file_desc = -1; - diskann::cout << "AppleAlignedFileReader created, this=" << this << std::endl; -} - -AppleAlignedFileReader::~AppleAlignedFileReader() -{ - diskann::cout << "AppleAlignedFileReader destructor called, this=" << this << std::endl; - - // 先解注册所有线程 - deregister_all_threads(); - - // 关闭文件描述符 - if (this->file_desc >= 0) - { - diskann::cout << "Closing file in destructor, fd=" << this->file_desc << std::endl; - ::close(this->file_desc); - this->file_desc = -1; - } -} - -IOContext &AppleAlignedFileReader::get_ctx() -{ - auto thread_id = std::this_thread::get_id(); - - // 创建一个静态空上下文用于错误情况 - static IOContext empty_ctx; - static bool initialized = false; - - if (!initialized) - { - empty_ctx.queue = nullptr; - empty_ctx.grp = nullptr; - empty_ctx.channel = nullptr; - initialized = true; - } - - std::unique_lock lk(this->ctx_mut); - - // 如果线程未注册，自动注册它 - if (ctx_map.find(thread_id) == ctx_map.end()) - { - lk.unlock(); - diskann::cerr << "Thread " << thread_id << " not registered, auto-registering" << std::endl; - - // 自动注册线程 - if (this->file_desc >= 0) - { - this->register_thread(); - - // 再次检查是否注册成功 - lk.lock(); - if (ctx_map.find(thread_id) != ctx_map.end()) - { - return ctx_map[thread_id]; - } - lk.unlock(); - } - - return empty_ctx; - } - - // 如果已注册，直接返回上下文 - IOContext &ctx = ctx_map[thread_id]; - lk.unlock(); - return ctx; -} - -void AppleAlignedFileReader::register_thread() -{ - auto current_id = std::this_thread::get_id(); - diskann::cout << "register_thread called from thread " << current_id << " on instance " << this << std::endl; - - // 检查文件描述符是否有效 - if (this->file_desc < 0) - { - diskann::cerr << "Thread " << current_id << " - register_thread called with invalid file descriptor" - << std::endl; - return; - } - - // 检查线程是否已注册 - { - std::lock_guard ctx_lock(this->ctx_mut); - if (ctx_map.find(current_id) != ctx_map.end()) - { - diskann::cout << "Thread " << current_id << " already registered" << std::endl; - return; - } - } - - // 创建线程上下文 - IOContext ctx; - ctx.queue = nullptr; - ctx.grp = nullptr; - ctx.channel = nullptr; - - std::string queue_name = - "diskann_io_" + std::to_string(*static_cast(static_cast(¤t_id))); - ctx.queue = dispatch_queue_create(queue_name.c_str(), DISPATCH_QUEUE_SERIAL); - if (!ctx.queue) - { - diskann::cerr << "Failed to create queue for thread " << current_id << std::endl; - return; - } - - ctx.grp = dispatch_group_create(); - if (!ctx.grp) - { - diskann::cerr << "Failed to create group for thread " << current_id << std::endl; - dispatch_release(ctx.queue); - return; - } - - // 复制文件描述符 - int dup_fd = ::dup(this->file_desc); - if (dup_fd == -1) - { - diskann::cerr << "Failed to duplicate file descriptor: " << this->file_desc << ", errno=" << errno << std::endl; - dispatch_release(ctx.grp); - dispatch_release(ctx.queue); - return; - } - - // 创建IO通道 - ctx.channel = dispatch_io_create(DISPATCH_IO_RANDOM, dup_fd, ctx.queue, ^(int error) { - ::close(dup_fd); - diskann::cout << "IO channel cleanup called, closed fd=" << dup_fd << std::endl; - }); - - if (!ctx.channel) - { - diskann::cerr << "Failed to create IO channel for thread " << current_id << ", fd=" << dup_fd - << ", errno=" << errno << std::endl; - ::close(dup_fd); - dispatch_release(ctx.grp); - dispatch_release(ctx.queue); - return; - } - - // 设置IO通道参数 - dispatch_io_set_low_water(ctx.channel, SECTOR_LEN); - dispatch_io_set_high_water(ctx.channel, SECTOR_LEN * 16); - - // 添加到线程映射 - { - std::lock_guard ctx_lock(this->ctx_mut); - ctx_map[current_id] = ctx; - } - - diskann::cout << "Thread " << current_id << " successfully registered with fd=" << dup_fd << std::endl; -} - -void AppleAlignedFileReader::deregister_thread() -{ - auto my_id = std::this_thread::get_id(); - diskann::cout << "deregister_thread called from thread " << my_id << " on instance " << this << std::endl; - - IOContext ctx; - bool found = false; - - { - std::lock_guard ctx_lock(this->ctx_mut); - if (ctx_map.find(my_id) != ctx_map.end()) - { - ctx = ctx_map[my_id]; - ctx_map.erase(my_id); - found = true; - } - } - - if (!found) - { - diskann::cerr << "Thread " << my_id << " not registered, cannot deregister" << std::endl; - return; - } - - if (ctx.channel) - { - dispatch_io_close(ctx.channel, DISPATCH_IO_STOP); - dispatch_release(ctx.channel); - } - - if (ctx.grp) - { - dispatch_release(ctx.grp); - } - - if (ctx.queue) - { - dispatch_release(ctx.queue); - } - - diskann::cout << "Thread " << my_id << " deregistered" << std::endl; -} - -void AppleAlignedFileReader::deregister_all_threads() -{ - diskann::cout << "deregister_all_threads called on instance " << this << std::endl; - - std::vector contexts; - - { - std::lock_guard ctx_lock(this->ctx_mut); - diskann::cout << "Deregistering " << ctx_map.size() << " threads" << std::endl; - for (auto &pair : ctx_map) - { - contexts.push_back(pair.second); - } - ctx_map.clear(); - } - - for (auto &ctx : contexts) - { - if (ctx.channel) - { - dispatch_io_close(ctx.channel, DISPATCH_IO_STOP); - dispatch_release(ctx.channel); - } - - if (ctx.grp) - { - dispatch_release(ctx.grp); - } - - if (ctx.queue) - { - dispatch_release(ctx.queue); - } - } - - diskann::cout << "All threads deregistered" << std::endl; -} - -void AppleAlignedFileReader::open(const std::string &fname) -{ - diskann::cout << "open called for file: " << fname << " on instance " << this << std::endl; - - // 关闭已存在的文件 - if (this->file_desc >= 0) - { - diskann::cout << "Closing existing file descriptor: " << this->file_desc << std::endl; - ::close(this->file_desc); - this->file_desc = -1; - } - - // 清空所有线程上下文 - deregister_all_threads(); - - // 打开新文件 - this->file_desc = ::open(fname.c_str(), O_RDONLY); - if (this->file_desc == -1) - { - diskann::cerr << "Failed to open file: " << fname << ", errno=" << errno << std::endl; - throw std::runtime_error("Failed to open file"); // 文件打开失败是致命错误 - } - - // 获取文件信息 - struct stat file_info; - if (::fstat(this->file_desc, &file_info) == 0) - { - diskann::cout << "File opened successfully: " << fname << ", size: " << file_info.st_size - << " bytes, fd=" << this->file_desc << std::endl; - } - else - { - diskann::cout << "File opened but couldn't get file info, fd=" << this->file_desc << std::endl; - } -} - -void AppleAlignedFileReader::close() -{ - diskann::cout << "close called on instance " << this << std::endl; - - // 先清理线程上下文 - deregister_all_threads(); - - // 关闭文件描述符 - if (this->file_desc >= 0) - { - diskann::cout << "Closing file descriptor: " << this->file_desc << std::endl; - ::close(this->file_desc); - this->file_desc = -1; - } -} - -void AppleAlignedFileReader::read(std::vector &read_reqs, IOContext &ctx, bool async) -{ - auto thread_id = std::this_thread::get_id(); - - // 如果通道无效，自动尝试注册线程 - if (!ctx.channel && this->file_desc >= 0) - { - diskann::cout << "Auto-registering thread " << thread_id << " during read" << std::endl; - this->register_thread(); - // 获取新的上下文 - ctx = this->get_ctx(); - } - - // 安全检查 - if (!ctx.channel || !ctx.queue || !ctx.grp) - { - diskann::cerr << "Invalid IO context in thread " << thread_id << std::endl; - return; - } - - dispatch_io_t channel = ctx.channel; - dispatch_queue_t q = ctx.queue; - dispatch_group_t group = ctx.grp; - - // 处理所有读取请求 - uint64_t n_reqs = read_reqs.size(); - for (uint64_t i = 0; i < n_reqs; i++) - { - AlignedRead &req = read_reqs[i]; - - // 检查对齐 - if (!IS_ALIGNED(req.buf, SECTOR_LEN) || !IS_ALIGNED(req.offset, SECTOR_LEN) || !IS_ALIGNED(req.len, SECTOR_LEN)) - { - diskann::cerr << "Thread " << thread_id << " - alignment error for request " << i << std::endl; - continue; - } - - dispatch_group_enter(group); - - dispatch_io_read(channel, req.offset, req.len, q, ^(bool done, dispatch_data_t data, int error) { - if (error) - { - diskann::cerr << "Thread " << thread_id << " read error: " << error << " when reading at offset " - << req.offset << std::endl; - if (done) - dispatch_group_leave(group); - return; - } - - if (data) - { - size_t actual_size = dispatch_data_get_size(data); - if (actual_size > 0) - { - __block size_t total_copied = 0; - dispatch_data_apply(data, - ^(dispatch_data_t region, size_t region_offset, const void *buffer, size_t size) { - if (region_offset + size <= req.len) - { - memcpy((char *)req.buf + region_offset, buffer, size); - total_copied += size; - return (bool)true; - } - diskann::cerr << "Buffer overflow: region_offset=" << region_offset - << ", size=" << size << ", req.len=" << req.len << std::endl; - return (bool)false; - }); - - if (total_copied != req.len && done) - { - diskann::cerr << "Warning: Only copied " << total_copied << " of " << req.len - << " requested bytes" << std::endl; - } - } - } - - // 仅在完成时离开组 - if (done) - { - dispatch_group_leave(group); - } - }); - } - - dispatch_group_wait(group, DISPATCH_TIME_FOREVER); -} - -#endif \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/disk_utils.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/disk_utils.cpp deleted file mode 100644 index a17d126..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/disk_utils.cpp +++ /dev/null @@ -1,1544 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "common_includes.h" - -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) -#include "gperftools/malloc_extension.h" -#endif - -#ifdef __APPLE__ -#include -#else -#include "mkl.h" -#endif - -#include "logger.h" -#include "disk_utils.h" -#include "cached_io.h" -#include "index.h" -#include "omp.h" -#include "percentile_stats.h" -#include "partition.h" -#include "pq_flash_index.h" -#include "timer.h" -#include "tsl/robin_set.h" - -namespace diskann -{ - -void add_new_file_to_single_index(std::string index_file, std::string new_file) -{ - std::unique_ptr metadata; - size_t nr, nc; - diskann::load_bin(index_file, metadata, nr, nc); - if (nc != 1) - { - std::stringstream stream; - stream << "Error, index file specified does not have correct metadata. " << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - size_t index_ending_offset = metadata[nr - 1]; - size_t read_blk_size = 64 * 1024 * 1024; - cached_ofstream writer(index_file, read_blk_size); - size_t check_file_size = get_file_size(index_file); - if (check_file_size != index_ending_offset) - { - std::stringstream stream; - stream << "Error, index file specified does not have correct metadata " - "(last entry must match the filesize). " - << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - - cached_ifstream reader(new_file, read_blk_size); - size_t fsize = reader.get_file_size(); - if (fsize == 0) - { - std::stringstream stream; - stream << "Error, new file specified is empty. Not appending. " << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - - size_t num_blocks = DIV_ROUND_UP(fsize, read_blk_size); - char *dump = new char[read_blk_size]; - for (uint64_t i = 0; i < num_blocks; i++) - { - size_t cur_block_size = - read_blk_size > fsize - (i * read_blk_size) ? fsize - (i * read_blk_size) : read_blk_size; - reader.read(dump, cur_block_size); - writer.write(dump, cur_block_size); - } - // reader.close(); - // writer.close(); - - delete[] dump; - std::vector new_meta; - for (uint64_t i = 0; i < nr; i++) - new_meta.push_back(metadata[i]); - new_meta.push_back(metadata[nr - 1] + fsize); - - diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); -} - -double get_memory_budget(double search_ram_budget) -{ - double final_index_ram_limit = search_ram_budget; - if (search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB > THRESHOLD_FOR_CACHING_IN_GB) - { // slack for space used by cached - // nodes - final_index_ram_limit = search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB; - } - return final_index_ram_limit * 1024 * 1024 * 1024; -} - -double get_memory_budget(const std::string &mem_budget_str) -{ - double search_ram_budget = atof(mem_budget_str.c_str()); - return get_memory_budget(search_ram_budget); -} - -size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim, - const std::vector ¶m_list) -{ - size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(final_index_ram_limit / (double)points_num)); - diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; - if (param_list.size() >= 6) - { - float compress_ratio = (float)atof(param_list[5].c_str()); - if (compress_ratio > 0 && compress_ratio <= 1) - { - size_t chunks_by_cr = (size_t)(std::floor)(compress_ratio * dim); - - if (chunks_by_cr > 0 && chunks_by_cr < num_pq_chunks) - { - diskann::cout << "Compress ratio:" << compress_ratio << " new #pq_chunks:" << chunks_by_cr << std::endl; - num_pq_chunks = chunks_by_cr; - } - else - { - diskann::cout << "Compress ratio: " << compress_ratio << " #new pq_chunks: " << chunks_by_cr - << " is either zero or greater than num_pq_chunks: " << num_pq_chunks - << ". num_pq_chunks is unchanged. " << std::endl; - } - } - else - { - diskann::cerr << "Compression ratio: " << compress_ratio << " should be in (0,1]" << std::endl; - } - } - - num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; - num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - num_pq_chunks = num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; - - diskann::cout << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector." - << std::endl; - return num_pq_chunks; -} - -template T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - warmup_num = 100000; - diskann::cout << "Generating random warmup file with dim " << warmup_dim << " and aligned dim " - << warmup_aligned_dim << std::flush; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } - diskann::cout << "..done" << std::endl; - return warmup; -} - -#ifdef EXEC_ENV_OLS -template -T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - uint64_t file_dim, file_aligned_dim; - - if (files.fileExists(cache_warmup_file)) - { - diskann::load_aligned_bin(files, cache_warmup_file, warmup, warmup_num, file_dim, file_aligned_dim); - diskann::cout << "In the warmup file: " << cache_warmup_file << " File dim: " << file_dim - << " File aligned dim: " << file_aligned_dim << " Expected dim: " << warmup_dim - << " Expected aligned dim: " << warmup_aligned_dim << std::endl; - - if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) - { - std::stringstream stream; - stream << "Mismatched dimensions in sample file. file_dim = " << file_dim - << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - diskann::cerr << stream.str(); - throw diskann::ANNException(stream.str(), -1); - } - } - else - { - warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); - } - return warmup; -} -#endif - -template -T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - size_t file_dim, file_aligned_dim; - - if (file_exists(cache_warmup_file)) - { - diskann::load_aligned_bin(cache_warmup_file, warmup, (size_t &)warmup_num, file_dim, file_aligned_dim); - if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) - { - std::stringstream stream; - stream << "Mismatched dimensions in sample file. file_dim = " << file_dim - << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - } - else - { - warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); - } - return warmup; -} - -/*************************************************** - Support for Merging Many Vamana Indices - ***************************************************/ - -void read_idmap(const std::string &fname, std::vector &ivecs) -{ - uint32_t npts32, dim; - size_t actual_file_size = get_file_size(fname); - std::ifstream reader(fname.c_str(), std::ios::binary); - reader.read((char *)&npts32, sizeof(uint32_t)); - reader.read((char *)&dim, sizeof(uint32_t)); - if (dim != 1 || actual_file_size != ((size_t)npts32) * sizeof(uint32_t) + 2 * sizeof(uint32_t)) - { - std::stringstream stream; - stream << "Error reading idmap file. Check if the file is bin file with " - "1 dimensional data. Actual: " - << actual_file_size << ", expected: " << (size_t)npts32 + 2 * sizeof(uint32_t) << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - ivecs.resize(npts32); - reader.read((char *)ivecs.data(), ((size_t)npts32) * sizeof(uint32_t)); - reader.close(); -} - -int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, const std::string &idmaps_prefix, - const std::string &idmaps_suffix, const uint64_t nshards, uint32_t max_degree, - const std::string &output_vamana, const std::string &medoids_file, bool use_filters, - const std::string &labels_to_medoids_file) -{ - // Read ID maps - std::vector vamana_names(nshards); - std::vector> idmaps(nshards); - for (uint64_t shard = 0; shard < nshards; shard++) - { - vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; - read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, idmaps[shard]); - } - - // find max node id - size_t nnodes = 0; - size_t nelems = 0; - for (auto &idmap : idmaps) - { - for (auto &id : idmap) - { - nnodes = std::max(nnodes, (size_t)id); - } - nelems += idmap.size(); - } - nnodes++; - diskann::cout << "# nodes: " << nnodes << ", max. degree: " << max_degree << std::endl; - - // compute inverse map: node -> shards - std::vector> node_shard; - node_shard.reserve(nelems); - for (size_t shard = 0; shard < nshards; shard++) - { - diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; - for (size_t idx = 0; idx < idmaps[shard].size(); idx++) - { - size_t node_id = idmaps[shard][idx]; - node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)shard)); - } - } - std::sort(node_shard.begin(), node_shard.end(), [](const auto &left, const auto &right) { - return left.first < right.first || (left.first == right.first && left.second < right.second); - }); - diskann::cout << "Finished computing node -> shards map" << std::endl; - - // will merge all the labels to medoids files of each shard into one - // combined file - if (use_filters) - { - std::unordered_map> global_label_to_medoids; - - for (size_t i = 0; i < nshards; i++) - { - std::ifstream mapping_reader; - std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; - mapping_reader.open(map_file); - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(mapping_reader, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - uint32_t label = 0; - while (std::getline(iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - - uint32_t token_as_num = std::stoul(token); - - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; - } - global_label_to_medoids[label].push_back(idmaps[i][medoid]); - line_cnt++; - } - mapping_reader.close(); - } - - std::ofstream mapping_writer(labels_to_medoids_file); - assert(mapping_writer.is_open()); - for (auto iter : global_label_to_medoids) - { - mapping_writer << iter.first << ", "; - auto &vec = iter.second; - for (uint32_t idx = 0; idx < vec.size() - 1; idx++) - { - mapping_writer << vec[idx] << ", "; - } - mapping_writer << vec[vec.size() - 1] << std::endl; - } - mapping_writer.close(); - } - - // create cached vamana readers - std::vector vamana_readers(nshards); - for (size_t i = 0; i < nshards; i++) - { - vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); - size_t expected_file_size; - vamana_readers[i].read((char *)&expected_file_size, sizeof(uint64_t)); - } - - size_t vamana_metadata_size = - sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint64_t); // expected file size + max degree + - // medoid_id + frozen_point info - - // create cached vamana writers - cached_ofstream merged_vamana_writer(output_vamana, BUFFER_SIZE_FOR_CACHED_IO); - - size_t merged_index_size = vamana_metadata_size; // we initialize the size of the merged index to - // the metadata size - size_t merged_index_frozen = 0; - merged_vamana_writer.write((char *)&merged_index_size, - sizeof(uint64_t)); // we will overwrite the index size at the end - - uint32_t output_width = max_degree; - uint32_t max_input_width = 0; - // read width from each vamana to advance buffer by sizeof(uint32_t) bytes - for (auto &reader : vamana_readers) - { - uint32_t input_width; - reader.read((char *)&input_width, sizeof(uint32_t)); - max_input_width = input_width > max_input_width ? input_width : max_input_width; - } - - diskann::cout << "Max input width: " << max_input_width << ", output width: " << output_width << std::endl; - - merged_vamana_writer.write((char *)&output_width, sizeof(uint32_t)); - std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); - uint32_t nshards_u32 = (uint32_t)nshards; - uint32_t one_val = 1; - medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); - medoid_writer.write((char *)&one_val, sizeof(uint32_t)); - - uint64_t vamana_index_frozen = 0; // as of now the functionality to merge many overlapping vamana - // indices is supported only for bulk indices without frozen point. - // Hence the final index will also not have any frozen points. - for (uint64_t shard = 0; shard < nshards; shard++) - { - uint32_t medoid; - // read medoid - vamana_readers[shard].read((char *)&medoid, sizeof(uint32_t)); - vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); - assert(vamana_index_frozen == false); - // rename medoid - medoid = idmaps[shard][medoid]; - - medoid_writer.write((char *)&medoid, sizeof(uint32_t)); - // write renamed medoid - if (shard == (nshards - 1)) //--> uncomment if running hierarchical - merged_vamana_writer.write((char *)&medoid, sizeof(uint32_t)); - } - merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); - medoid_writer.close(); - - diskann::cout << "Starting merge" << std::endl; - - // Gopal. random_shuffle() is deprecated. - std::random_device rng; - std::mt19937 urng(rng()); - - std::vector nhood_set(nnodes, 0); - std::vector final_nhood; - - uint32_t nnbrs = 0, shard_nnbrs = 0; - uint32_t cur_id = 0; - for (const auto &id_shard : node_shard) - { - uint32_t node_id = id_shard.first; - uint32_t shard_id = id_shard.second; - if (cur_id < node_id) - { - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (size_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - if (cur_id % 499999 == 1) - { - diskann::cout << "." << std::flush; - } - cur_id = node_id; - nnbrs = 0; - for (auto &p : final_nhood) - nhood_set[p] = 0; - final_nhood.clear(); - } - // read from shard_id ifstream - vamana_readers[shard_id].read((char *)&shard_nnbrs, sizeof(uint32_t)); - - if (shard_nnbrs == 0) - { - diskann::cout << "WARNING: shard #" << shard_id << ", node_id " << node_id << " has 0 nbrs" << std::endl; - } - - std::vector shard_nhood(shard_nnbrs); - if (shard_nnbrs > 0) - vamana_readers[shard_id].read((char *)shard_nhood.data(), shard_nnbrs * sizeof(uint32_t)); - // rename nodes - for (uint64_t j = 0; j < shard_nnbrs; j++) - { - if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) - { - nhood_set[idmaps[shard_id][shard_nhood[j]]] = 1; - final_nhood.emplace_back(idmaps[shard_id][shard_nhood[j]]); - } - } - } - - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (size_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - if (nnbrs > 0) - { - merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); - } - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - for (auto &p : final_nhood) - nhood_set[p] = 0; - final_nhood.clear(); - - diskann::cout << "Expected size: " << merged_index_size << std::endl; - - merged_vamana_writer.reset(); - merged_vamana_writer.write((char *)&merged_index_size, sizeof(uint64_t)); - - diskann::cout << "Finished merge" << std::endl; - return 0; -} - -// TODO: Make this a streaming implementation to avoid exceeding the memory -// budget -/* If the number of filters per point N exceeds the graph degree R, - then it is difficult to have edges to all labels from this point. - This function break up such dense points to have only a threshold of maximum - T labels per point It divides one graph nodes to multiple nodes and append - the new nodes at the end. The dummy map contains the real graph id of the - new nodes added to the graph */ -template -void breakup_dense_points(const std::string data_file, const std::string labels_file, uint32_t density, - const std::string out_data_file, const std::string out_labels_file, - const std::string out_metadata_file) -{ - std::string token, line; - std::ifstream labels_stream(labels_file); - T *data; - size_t npts, ndims; - diskann::load_bin(data_file, data, npts, ndims); - - std::unordered_map dummy_pt_ids; - uint32_t next_dummy_id = (uint32_t)npts; - - uint32_t point_cnt = 0; - - std::vector> labels_per_point; - labels_per_point.resize(npts); - - uint32_t dense_pts = 0; - if (labels_stream.is_open()) - { - while (getline(labels_stream, line)) - { - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - uint32_t label_host = point_cnt; - while (getline(iss, token, ',')) - { - if (lbl_cnt == density) - { - if (label_host == point_cnt) - dense_pts++; - label_host = next_dummy_id; - labels_per_point.resize(next_dummy_id + 1); - dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; - next_dummy_id++; - lbl_cnt = 0; - } - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - uint32_t token_as_num = std::stoul(token); - labels_per_point[label_host].push_back(token_as_num); - lbl_cnt++; - } - point_cnt++; - } - } - diskann::cout << "fraction of dense points with >= " << density << " labels = " << (float)dense_pts / (float)npts - << std::endl; - - if (labels_per_point.size() != 0) - { - diskann::cout << labels_per_point.size() << " is the new number of points" << std::endl; - std::ofstream label_writer(out_labels_file); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < labels_per_point.size(); i++) - { - for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) - { - label_writer << labels_per_point[i][j] << ","; - } - if (labels_per_point[i].size() != 0) - label_writer << labels_per_point[i][labels_per_point[i].size() - 1]; - label_writer << std::endl; - } - label_writer.close(); - } - - if (dummy_pt_ids.size() != 0) - { - diskann::cout << dummy_pt_ids.size() << " is the number of dummy points created" << std::endl; - - T *ptr = (T *)std::realloc((void *)data, labels_per_point.size() * ndims * sizeof(T)); - if (ptr == nullptr) - { - diskann::cerr << "Realloc failed while creating dummy points" << std::endl; - free(data); - data = nullptr; - throw new diskann::ANNException("Realloc failed while expanding data.", -1, __FUNCTION__, __FILE__, - __LINE__); - } - else - { - data = ptr; - } - - std::ofstream dummy_writer(out_metadata_file); - assert(dummy_writer.is_open()); - for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) - { - dummy_writer << i->first << "," << i->second << std::endl; - std::memcpy(data + i->first * ndims, data + i->second * ndims, ndims * sizeof(T)); - } - dummy_writer.close(); - } - - diskann::save_bin(out_data_file, data, labels_per_point.size(), ndims); -} - -void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file) -{ // assumes ith row is for ith - // point in labels file - diskann::cout << "Extracting labels for shard" << std::endl; - - uint32_t *ids = nullptr; - size_t num_ids, tmp_dim; - diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); - - uint32_t counter = 0, shard_counter = 0; - std::string cur_line; - - std::ifstream label_reader(in_label_file); - std::ofstream label_writer(shard_label_file); - assert(label_reader.is_open()); - assert(label_reader.is_open()); - if (label_reader && label_writer) - { - while (std::getline(label_reader, cur_line)) - { - if (shard_counter >= num_ids) - { - break; - } - if (counter == ids[shard_counter]) - { - label_writer << cur_line << "\n"; - shard_counter++; - } - counter++; - } - } - if (ids != nullptr) - delete[] ids; -} - -template -int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, - double sampling_rate, double ram_budget, std::string mem_index_path, - std::string medoids_file, std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, - const uint32_t Lf) -{ - size_t base_num, base_dim; - diskann::get_bin_metadata(base_file, base_num, base_dim); - - double full_index_ram = estimate_ram_usage(base_num, (uint32_t)base_dim, sizeof(T), R); - - // TODO: Make this honest when there is filter support - if (full_index_ram < ram_budget * 1024 * 1024 * 1024) - { - diskann::cout << "Full index fits in RAM budget, should consume at most " - << full_index_ram / (1024 * 1024 * 1024) << "GiBs, so building in one shot" << std::endl; - - diskann::IndexWriteParameters paras = diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_saturate_graph(!use_filters) - .with_num_threads(num_threads) - .build(); - using TagT = uint32_t; - diskann::Index _index(compareMetric, base_dim, base_num, - std::make_shared(paras), nullptr, - defaults::NUM_FROZEN_POINTS_STATIC, false, false, false, - build_pq_bytes > 0, build_pq_bytes, use_opq, use_filters); - if (!use_filters) - _index.build(base_file.c_str(), base_num); - else - { - if (universal_label != "") - { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(base_file.c_str(), label_file, base_num); - } - _index.save(mem_index_path.c_str()); - - if (use_filters) - { - // need to copy the labels_to_medoids file to the specified input - // file - std::remove(labels_to_medoids_file.c_str()); - std::string mem_labels_to_medoid_file = mem_index_path + "_labels_to_medoids.txt"; - copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); - std::remove(mem_labels_to_medoid_file.c_str()); - } - - std::remove(medoids_file.c_str()); - std::remove(centroids_file.c_str()); - return 0; - } - - diskann::cout << "Full index does not fit in RAM budget, building in multiple shots" << std::endl; - - // where the universal label is to be saved in the final graph - std::string final_index_universal_label_file = mem_index_path + "_universal_label.txt"; - - std::string merged_index_prefix = mem_index_path + "_tempFiles"; - - Timer timer; - int num_parts = - partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - diskann::cout << timer.elapsed_seconds_for_step("partitioning data ") << std::endl; - - std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; - std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); - - timer.reset(); - for (int p = 0; p < num_parts; p++) - { -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); -#endif - - std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - - std::string shard_ids_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; - - std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - - retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); - - std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; - - diskann::IndexWriteParameters low_degree_params = diskann::IndexWriteParametersBuilder(L, 2 * R / 3) - .with_filter_list_size(Lf) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - - size_t shard_base_dim, shard_base_pts; - get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); - - diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, - std::make_shared(low_degree_params), nullptr, - defaults::NUM_FROZEN_POINTS_STATIC, false, false, false, build_pq_bytes > 0, - build_pq_bytes, use_opq); - if (!use_filters) - { - _index.build(shard_base_file.c_str(), shard_base_pts); - } - else - { - diskann::extract_shard_labels(label_file, shard_ids_file, shard_labels_file); - if (universal_label != "") - { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(shard_base_file.c_str(), shard_labels_file, shard_base_pts); - } - - // cal deg stats - size_t max_deg = 0, min_deg = SIZE_MAX, avg_deg = 0, cnt_deg = 0; - _index.get_degree_stats(max_deg, min_deg, avg_deg, cnt_deg); - std::cout << "! For shard " << p << " Degree stats: " << max_deg << ", " << min_deg << ", " << avg_deg << ", " - << cnt_deg << std::endl; - std::string shard_degree_stats_file = shard_index_file + "_degree_stats.txt"; - _index.dump_degree_stats(shard_degree_stats_file); - - _index.save(shard_index_file.c_str()); - // copy universal label file from first shard to the final destination - // index, since all shards anyway share the universal label - if (p == 0) - { - std::string shard_universal_label_file = shard_index_file + "_universal_label.txt"; - if (universal_label != "") - { - copy_file(shard_universal_label_file, final_index_universal_label_file); - } - } - - std::remove(shard_base_file.c_str()); - } - diskann::cout << timer.elapsed_seconds_for_step("building indices on shards") << std::endl; - - timer.reset(); - diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", merged_index_prefix + "_subshard-", - "_ids_uint32.bin", num_parts, R, mem_index_path, medoids_file, use_filters, - labels_to_medoids_file); - diskann::cout << timer.elapsed_seconds_for_step("merging indices") << std::endl; - - // delete tempFiles - for (int p = 0; p < num_parts; p++) - { - std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - std::string shard_id_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; - std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; - std::string shard_index_file_data = shard_index_file + ".data"; - - // std::remove(shard_base_file.c_str()); - // std::remove(shard_id_file.c_str()); - // std::remove(shard_index_file.c_str()); - // std::remove(shard_index_file_data.c_str()); - if (use_filters) - { - std::string shard_index_label_file = shard_index_file + "_labels.txt"; - std::string shard_index_univ_label_file = shard_index_file + "_universal_label.txt"; - std::string shard_index_label_map_file = shard_index_file + "_labels_to_medoids.txt"; - std::remove(shard_labels_file.c_str()); - std::remove(shard_index_label_file.c_str()); - std::remove(shard_index_label_map_file.c_str()); - std::remove(shard_index_univ_label_file.c_str()); - } - } - return 0; -} - -// General purpose support for DiskANN interface - -// optimizes the beamwidth to maximize QPS for a given L_search subject to -// 99.9 latency not blowing up -template -uint32_t optimize_beamwidth(std::unique_ptr> &pFlashIndex, T *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, - uint32_t nthreads, uint32_t start_bw) -{ - uint32_t cur_bw = start_bw; - double max_qps = 0; - uint32_t best_bw = start_bw; - bool stop_flag = false; - - while (!stop_flag) - { - std::vector tuning_sample_result_ids_64(tuning_sample_num, 0); - std::vector tuning_sample_result_dists(tuning_sample_num, 0); - diskann::QueryStats *stats = new diskann::QueryStats[tuning_sample_num]; - - auto s = std::chrono::high_resolution_clock::now(); -#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (int64_t i = 0; i < (int64_t)tuning_sample_num; i++) - { - pFlashIndex->cached_beam_search(tuning_sample + (i * tuning_sample_aligned_dim), 1, L, - tuning_sample_result_ids_64.data() + (i * 1), - tuning_sample_result_dists.data() + (i * 1), cur_bw, false, stats + i); - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - double qps = (1.0f * (float)tuning_sample_num) / (1.0f * (float)diff.count()); - - double lat_999 = diskann::get_percentile_stats( - stats, tuning_sample_num, 0.999f, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - double mean_latency = diskann::get_mean_stats( - stats, tuning_sample_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - if (qps > max_qps && lat_999 < (15000) + mean_latency * 2) - { - max_qps = qps; - best_bw = cur_bw; - cur_bw = (uint32_t)(std::ceil)((float)cur_bw * 1.1f); - } - else - { - stop_flag = true; - } - if (cur_bw > 64) - stop_flag = true; - - delete[] stats; - } - return best_bw; -} - -template -void create_disk_layout(const std::string base_file, const std::string mem_index_file, const std::string output_file, - const std::string reorder_data_file) -{ - uint32_t npts, ndims; - - // amount to read or write in one shot - size_t read_blk_size = 64 * 1024 * 1024; - size_t write_blk_size = read_blk_size; - cached_ifstream base_reader(base_file, read_blk_size); - base_reader.read((char *)&npts, sizeof(uint32_t)); - base_reader.read((char *)&ndims, sizeof(uint32_t)); - - size_t npts_64, ndims_64; - npts_64 = npts; - ndims_64 = ndims; - - // Check if we need to append data for re-ordering - bool append_reorder_data = false; - std::ifstream reorder_data_reader; - - uint32_t npts_reorder_file = 0, ndims_reorder_file = 0; - if (reorder_data_file != std::string("")) - { - append_reorder_data = true; - size_t reorder_data_file_size = get_file_size(reorder_data_file); - reorder_data_reader.exceptions(std::ofstream::failbit | std::ofstream::badbit); - - try - { - reorder_data_reader.open(reorder_data_file, std::ios::binary); - reorder_data_reader.read((char *)&npts_reorder_file, sizeof(uint32_t)); - reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(uint32_t)); - if (npts_reorder_file != npts) - throw ANNException("Mismatch in num_points between reorder " - "data file and base file", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (reorder_data_file_size != 8 + sizeof(float) * (size_t)npts_reorder_file * (size_t)ndims_reorder_file) - throw ANNException("Discrepancy in reorder data file size ", -1, __FUNCSIG__, __FILE__, __LINE__); - } - catch (std::system_error &e) - { - throw FileException(reorder_data_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - - // create cached reader + writer - size_t actual_file_size = get_file_size(mem_index_file); - diskann::cout << "Vamana index file size=" << actual_file_size << std::endl; - std::ifstream vamana_reader(mem_index_file, std::ios::binary); - cached_ofstream diskann_writer(output_file, write_blk_size); - - // metadata: width, medoid - uint32_t width_u32, medoid_u32; - size_t index_file_size; - - vamana_reader.read((char *)&index_file_size, sizeof(uint64_t)); - if (index_file_size != actual_file_size) - { - std::stringstream stream; - stream << "Vamana Index file size does not match expected size per " - "meta-data." - << " file size from file: " << index_file_size << " actual file size: " << actual_file_size << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - uint64_t vamana_frozen_num = false, vamana_frozen_loc = 0; - - vamana_reader.read((char *)&width_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&medoid_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); - // compute - uint64_t medoid, max_node_len, nnodes_per_sector; - npts_64 = (uint64_t)npts; - medoid = (uint64_t)medoid_u32; - if (vamana_frozen_num == 1) - vamana_frozen_loc = medoid; - max_node_len = (((uint64_t)width_u32 + 1) * sizeof(uint32_t)) + (ndims_64 * sizeof(T)); - nnodes_per_sector = defaults::SECTOR_LEN / max_node_len; // 0 if max_node_len > SECTOR_LEN - - diskann::cout << "medoid: " << medoid << "B" << std::endl; - diskann::cout << "max_node_len: " << max_node_len << "B" << std::endl; - diskann::cout << "nnodes_per_sector: " << nnodes_per_sector << "B" << std::endl; - - // defaults::SECTOR_LEN buffer for each sector - std::unique_ptr sector_buf = std::make_unique(defaults::SECTOR_LEN); - std::unique_ptr multisector_buf = std::make_unique(ROUND_UP(max_node_len, defaults::SECTOR_LEN)); - std::unique_ptr node_buf = std::make_unique(max_node_len); - uint32_t &nnbrs = *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)); - uint32_t *nhood_buf = (uint32_t *)(node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(uint32_t)); - - // number of sectors (1 for meta data) - uint64_t n_sectors = nnodes_per_sector > 0 ? ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector - : npts_64 * DIV_ROUND_UP(max_node_len, defaults::SECTOR_LEN); - uint64_t n_reorder_sectors = 0; - uint64_t n_data_nodes_per_sector = 0; - - if (append_reorder_data) - { - n_data_nodes_per_sector = defaults::SECTOR_LEN / (ndims_reorder_file * sizeof(float)); - n_reorder_sectors = ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; - } - uint64_t disk_index_file_size = (n_sectors + n_reorder_sectors + 1) * defaults::SECTOR_LEN; - - std::vector output_file_meta; - output_file_meta.push_back(npts_64); - output_file_meta.push_back(ndims_64); - output_file_meta.push_back(medoid); - output_file_meta.push_back(max_node_len); - output_file_meta.push_back(nnodes_per_sector); - output_file_meta.push_back(vamana_frozen_num); - output_file_meta.push_back(vamana_frozen_loc); - output_file_meta.push_back((uint64_t)append_reorder_data); - if (append_reorder_data) - { - output_file_meta.push_back(n_sectors + 1); - output_file_meta.push_back(ndims_reorder_file); - output_file_meta.push_back(n_data_nodes_per_sector); - } - output_file_meta.push_back(disk_index_file_size); - - diskann_writer.write(sector_buf.get(), defaults::SECTOR_LEN); - - std::unique_ptr cur_node_coords = std::make_unique(ndims_64); - diskann::cout << "# sectors: " << n_sectors << std::endl; - uint64_t cur_node_id = 0; - - if (nnodes_per_sector > 0) - { // Write multiple nodes per sector - for (uint64_t sector = 0; sector < n_sectors; sector++) - { - if (sector % 100000 == 0) - { - diskann::cout << "Sector #" << sector << "written" << std::endl; - } - memset(sector_buf.get(), 0, defaults::SECTOR_LEN); - for (uint64_t sector_node_id = 0; sector_node_id < nnodes_per_sector && cur_node_id < npts_64; - sector_node_id++) - { - memset(node_buf.get(), 0, max_node_len); - // read cur node's nnbrs - vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); - - // sanity checks on nnbrs - assert(nnbrs > 0); - assert(nnbrs <= width_u32); - - // read node's nhood - vamana_reader.read((char *)nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - if (nnbrs > width_u32) - { - vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), vamana_reader.cur); - } - - // write coords of node first - // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); - base_reader.read((char *)cur_node_coords.get(), sizeof(T) * ndims_64); - memcpy(node_buf.get(), cur_node_coords.get(), ndims_64 * sizeof(T)); - - // write nnbrs - *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)) = (std::min)(nnbrs, width_u32); - - // write nhood next - memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), nhood_buf, - (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - - // get offset into sector_buf - char *sector_node_buf = sector_buf.get() + (sector_node_id * max_node_len); - - // copy node buf into sector_node_buf - memcpy(sector_node_buf, node_buf.get(), max_node_len); - cur_node_id++; - } - // flush sector to disk - diskann_writer.write(sector_buf.get(), defaults::SECTOR_LEN); - } - } - else - { // Write multi-sector nodes - uint64_t nsectors_per_node = DIV_ROUND_UP(max_node_len, defaults::SECTOR_LEN); - for (uint64_t i = 0; i < npts_64; i++) - { - if ((i * nsectors_per_node) % 100000 == 0) - { - diskann::cout << "Sector #" << i * nsectors_per_node << "written" << std::endl; - } - memset(multisector_buf.get(), 0, nsectors_per_node * defaults::SECTOR_LEN); - - memset(node_buf.get(), 0, max_node_len); - // read cur node's nnbrs - vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); - - // sanity checks on nnbrs - assert(nnbrs > 0); - assert(nnbrs <= width_u32); - - // read node's nhood - vamana_reader.read((char *)nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - if (nnbrs > width_u32) - { - vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), vamana_reader.cur); - } - - // write coords of node first - // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); - base_reader.read((char *)cur_node_coords.get(), sizeof(T) * ndims_64); - memcpy(multisector_buf.get(), cur_node_coords.get(), ndims_64 * sizeof(T)); - - // write nnbrs - *(uint32_t *)(multisector_buf.get() + ndims_64 * sizeof(T)) = (std::min)(nnbrs, width_u32); - - // write nhood next - memcpy(multisector_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), nhood_buf, - (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - - // flush sector to disk - diskann_writer.write(multisector_buf.get(), nsectors_per_node * defaults::SECTOR_LEN); - } - } - - if (append_reorder_data) - { - diskann::cout << "Index written. Appending reorder data..." << std::endl; - - auto vec_len = ndims_reorder_file * sizeof(float); - std::unique_ptr vec_buf = std::make_unique(vec_len); - - for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) - { - if (sector % 100000 == 0) - { - diskann::cout << "Reorder data Sector #" << sector << "written" << std::endl; - } - - memset(sector_buf.get(), 0, defaults::SECTOR_LEN); - - for (uint64_t sector_node_id = 0; sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; - sector_node_id++) - { - memset(vec_buf.get(), 0, vec_len); - reorder_data_reader.read(vec_buf.get(), vec_len); - - // copy node buf into sector_node_buf - memcpy(sector_buf.get() + (sector_node_id * vec_len), vec_buf.get(), vec_len); - } - // flush sector to disk - diskann_writer.write(sector_buf.get(), defaults::SECTOR_LEN); - } - } - diskann_writer.close(); - diskann::save_bin(output_file, output_file_meta.data(), output_file_meta.size(), 1, 0); - diskann::cout << "Output disk index file written to " << output_file << std::endl; -} - -template -int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, const uint32_t filter_threshold, - const uint32_t Lf) -{ - std::stringstream parser; - parser << std::string(indexBuildParameters); - std::string cur_param; - std::vector param_list; - while (parser >> cur_param) - { - param_list.push_back(cur_param); - } - if (param_list.size() < 5 || param_list.size() > 9) - { - diskann::cout << "Correct usage of parameters is R (max degree)\n" - "L (indexing list size, better if >= R)\n" - "B (RAM limit of final index in GB)\n" // search - "M (memory limit while indexing)\n" // build - "T (number of threads for indexing)\n" - "B' (PQ bytes for disk index: optional parameter for " - "very large dimensional data)\n" - "reorder (set true to include full precision in data file" - ": optional paramter, use only when using disk PQ\n" - "build_PQ_byte (number of PQ bytes for inde build; set 0 to use " - "full precision vectors)\n" - "QD Quantized Dimension to overwrite the derived dim from B " - << std::endl; - return -1; - } - - if (!std::is_same::value && - (compareMetric == diskann::Metric::INNER_PRODUCT || compareMetric == diskann::Metric::COSINE)) - { - std::stringstream stream; - stream << "Disk-index build currently only supports floating point data for Max " - "Inner Product Search/ cosine similarity. " - << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - - size_t disk_pq_dims = 0; - bool use_disk_pq = false; - size_t build_pq_bytes = 0; - - // if there is a 6th parameter, it means we compress the disk index - // vectors also using PQ data (for very large dimensionality data). If the - // provided parameter is 0, it means we store full vectors. - if (param_list.size() > 5) - { - disk_pq_dims = atoi(param_list[5].c_str()); - use_disk_pq = true; - if (disk_pq_dims == 0) - use_disk_pq = false; - } - - bool reorder_data = false; - if (param_list.size() >= 7) - { - if (1 == atoi(param_list[6].c_str())) - { - reorder_data = true; - } - } - - if (param_list.size() >= 8) - { - build_pq_bytes = atoi(param_list[7].c_str()); - } - - std::string base_file(dataFilePath); - std::string data_file_to_use = base_file; - std::string labels_file_original = label_file; - std::string index_prefix_path(indexFilePath); - std::string labels_file_to_use = index_prefix_path + "_label_formatted.txt"; - std::string pq_pivots_path_base = codebook_prefix; - std::string pq_pivots_path = file_exists(pq_pivots_path_base) ? pq_pivots_path_base + "_pq_pivots.bin" - : index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; - std::string mem_index_path = index_prefix_path + "_mem.index"; - std::string disk_index_path = index_prefix_path + "_disk.index"; - std::string medoids_path = disk_index_path + "_medoids.bin"; - std::string centroids_path = disk_index_path + "_centroids.bin"; - - std::string labels_to_medoids_path = disk_index_path + "_labels_to_medoids.txt"; - std::string mem_labels_file = mem_index_path + "_labels.txt"; - std::string disk_labels_file = disk_index_path + "_labels.txt"; - std::string mem_univ_label_file = mem_index_path + "_universal_label.txt"; - std::string disk_univ_label_file = disk_index_path + "_universal_label.txt"; - std::string disk_labels_int_map_file = disk_index_path + "_labels_map.txt"; - std::string dummy_remap_file = disk_index_path + "_dummy_map.txt"; // remap will be used if we break-up points of - // high label-density to create copies - - std::string sample_base_prefix = index_prefix_path + "_sample"; - // optional, used if disk index file must store pq data - std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; - // optional, used if disk index must store pq data - std::string disk_pq_compressed_vectors_path = index_prefix_path + "_disk.index_pq_compressed.bin"; - std::string prepped_base = - index_prefix_path + - "_prepped_base.bin"; // temp file for storing pre-processed base file for cosine/ mips metrics - bool created_temp_file_for_processed_data = false; - - // output a new base file which contains extra dimension with sqrt(1 - - // ||x||^2/M^2) for every x, M is max norm of all points. Extra space on - // disk needed! - if (compareMetric == diskann::Metric::INNER_PRODUCT) - { - Timer timer; - std::cout << "Using Inner Product search, so need to pre-process base " - "data into temp file. Please ensure there is additional " - "(n*(d+1)*4) bytes for storing pre-processed base vectors, " - "apart from the interim indices created by DiskANN and the final index." - << std::endl; - data_file_to_use = prepped_base; - float max_norm_of_base = diskann::prepare_base_for_inner_products(base_file, prepped_base); - std::string norm_file = disk_index_path + "_max_base_norm.bin"; - - diskann::save_bin(norm_file, &max_norm_of_base, 1, 1); - diskann::cout << timer.elapsed_seconds_for_step("preprocessing data for inner product") << std::endl; - created_temp_file_for_processed_data = true; - - diskann::cout << "Reading max_norm_of_base from " << norm_file << std::endl; - float *max_norm_of_base_ptr; - size_t npts, ndims; - diskann::load_bin(norm_file, max_norm_of_base_ptr, npts, ndims); - if (max_norm_of_base != *max_norm_of_base_ptr) - { - diskann::cout << "max_norm_of_base mismatch: " << max_norm_of_base << " != " << *max_norm_of_base_ptr - << std::endl; - assert(false); - } - diskann::cout << "max_norm_of_base: " << max_norm_of_base << std::endl; - diskann::cout << "! Using prepped_base file at " << prepped_base << std::endl; - if (!file_exists(prepped_base)) - { - diskann::cout << "! prepped_base file does not exist, please check the file path" << std::endl; - assert(false); - } - } - else if (compareMetric == diskann::Metric::COSINE) - { - Timer timer; - std::cout << "Normalizing data for cosine to temporary file, please ensure there is additional " - "(n*d*4) bytes for storing normalized base vectors, " - "apart from the interim indices created by DiskANN and the final index." - << std::endl; - data_file_to_use = prepped_base; - diskann::normalize_data_file(base_file, prepped_base); - diskann::cout << timer.elapsed_seconds_for_step("preprocessing data for cosine") << std::endl; - created_temp_file_for_processed_data = true; - } - - uint32_t R = (uint32_t)atoi(param_list[0].c_str()); - uint32_t L = (uint32_t)atoi(param_list[1].c_str()); - - double final_index_ram_limit = get_memory_budget(param_list[2]); - if (final_index_ram_limit <= 0) - { - std::cerr << "Insufficient memory budget (or string was not in right " - "format). Should be > 0." - << std::endl; - return -1; - } - double indexing_ram_budget = (float)atof(param_list[3].c_str()); - if (indexing_ram_budget <= 0) - { - std::cerr << "Not building index. Please provide more RAM budget" << std::endl; - return -1; - } - uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); - - if (num_threads != 0) - { - omp_set_num_threads(num_threads); -#ifdef __x86_64__ - mkl_set_num_threads(num_threads); -#endif - } - - diskann::cout << "Starting index build: R=" << R << " L=" << L << " Query RAM budget: " << final_index_ram_limit - << " Indexing ram budget: " << indexing_ram_budget << " T: " << num_threads << std::endl; - - auto s = std::chrono::high_resolution_clock::now(); - - // If there is filter support, we break-up points which have too many labels - // into replica dummy points which evenly distribute the filters. The rest - // of index build happens on the augmented base and labels - std::string augmented_data_file, augmented_labels_file; - if (use_filters) - { - convert_labels_string_to_int(labels_file_original, labels_file_to_use, disk_labels_int_map_file, - universal_label); - augmented_data_file = index_prefix_path + "_augmented_data.bin"; - augmented_labels_file = index_prefix_path + "_augmented_labels.txt"; - if (filter_threshold != 0) - { - breakup_dense_points(data_file_to_use, labels_file_to_use, filter_threshold, augmented_data_file, - augmented_labels_file, - dummy_remap_file); // RKNOTE: This has large memory footprint, - // need to make this streaming - data_file_to_use = augmented_data_file; - labels_file_to_use = augmented_labels_file; - } - } - - size_t points_num, dim; - - diskann::cout << "getting bin metadata" << std::endl; - Timer timer; - diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); - diskann::cout << timer.elapsed_seconds_for_step("getting bin metadata") << std::endl; - const double p_val = ((double)MAX_PQ_TRAINING_SET_SIZE / (double)points_num); - - if (use_disk_pq) - { - generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, disk_pq_compressed_vectors_path, - compareMetric, p_val, disk_pq_dims); - } - size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(final_index_ram_limit / points_num)); - - num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; - num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - num_pq_chunks = num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; - - if (param_list.size() >= 9 && atoi(param_list[8].c_str()) <= MAX_PQ_CHUNKS && atoi(param_list[8].c_str()) > 0) - { - std::cout << "Use quantized dimension (QD) to overwrite derived quantized " - "dimension from search_DRAM_budget (B)" - << std::endl; - num_pq_chunks = atoi(param_list[8].c_str()); - } - - diskann::cout << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector." - << std::endl; - - generate_quantized_data(data_file_to_use, pq_pivots_path, pq_compressed_vectors_path, compareMetric, p_val, - num_pq_chunks, use_opq, codebook_prefix); - diskann::cout << timer.elapsed_seconds_for_step("generating quantized data") << std::endl; - -// Gopal. Splitting diskann_dll into separate DLLs for search and build. -// This code should only be available in the "build" DLL. -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); -#endif - // Whether it is cosine or inner product, we still L2 metric due to the pre-processing. - timer.reset(); - diskann::build_merged_vamana_index(data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, - indexing_ram_budget, mem_index_path, medoids_path, centroids_path, - build_pq_bytes, use_opq, num_threads, use_filters, labels_file_to_use, - labels_to_medoids_path, universal_label, Lf); - diskann::cout << timer.elapsed_seconds_for_step("building merged vamana index") << std::endl; - - timer.reset(); - if (!use_disk_pq) - { - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); - } - else - { - if (!reorder_data) - diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); - else - diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path, - data_file_to_use.c_str()); - } - diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") << std::endl; - - double ten_percent_points = std::ceil(points_num * 0.1); - double num_sample_points = - ten_percent_points > MAX_SAMPLE_POINTS_FOR_WARMUP ? MAX_SAMPLE_POINTS_FOR_WARMUP : ten_percent_points; - double sample_sampling_rate = num_sample_points / points_num; - gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); - if (use_filters) - { - copy_file(labels_file_to_use, disk_labels_file); - std::remove(mem_labels_file.c_str()); - if (universal_label != "") - { - copy_file(mem_univ_label_file, disk_univ_label_file); - std::remove(mem_univ_label_file.c_str()); - } - std::remove(augmented_data_file.c_str()); - std::remove(augmented_labels_file.c_str()); - std::remove(labels_file_to_use.c_str()); - } - if (created_temp_file_for_processed_data) - std::remove(prepped_base.c_str()); - std::remove(mem_index_path.c_str()); - std::remove((mem_index_path + ".data").c_str()); - std::remove((mem_index_path + ".tags").c_str()); - if (use_disk_pq) - std::remove(disk_pq_compressed_vectors_path.c_str()); - - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - diskann::cout << "Indexing time: " << diff.count() << std::endl; - - return 0; -} - -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); - -template DISKANN_DLLEXPORT int8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); - -#ifdef EXEC_ENV_OLS -template DISKANN_DLLEXPORT int8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -#endif - -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, int8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, uint8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, float *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, int8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, uint8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); -template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, float *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -// LabelT = uint16 -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); - -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -// Label=16_t -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -}; // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/distance.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/distance.cpp deleted file mode 100644 index 2fa4c7a..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/distance.cpp +++ /dev/null @@ -1,743 +0,0 @@ -// TODO -// CHECK COSINE ON LINUX - -#ifdef _WINDOWS -#include -#include -#include -#include -#include -#include "simd_utils.h" -#elif __APPLE__ -#include -#else -#include -#include "simd_utils.h" -#endif - -#include - -#include "distance.h" -#include "utils.h" -#include "logger.h" -#include "ann_exception.h" - -namespace diskann -{ - -// -// Base Class Implementatons -// -template -float Distance::compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const -{ - throw std::logic_error("This function is not implemented."); -} - -template uint32_t Distance::post_normalization_dimension(uint32_t orig_dimension) const -{ - return orig_dimension; -} - -template diskann::Metric Distance::get_metric() const -{ - return _distance_metric; -} - -template bool Distance::preprocessing_required() const -{ - return false; -} - -template -void Distance::preprocess_base_points(T *original_data, const size_t orig_dim, const size_t num_points) -{ -} - -template void Distance::preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query) -{ - std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); -} - -template size_t Distance::get_required_alignment() const -{ - return _alignment_factor; -} - -// -// Cosine distance functions. -// - -float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, uint32_t length) const -{ -#ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); -#else - int magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) - { - magA += ((int32_t)a[i]) * ((int32_t)a[i]); - magB += ((int32_t)b[i]) * ((int32_t)b[i]); - scalarProduct += ((int32_t)a[i]) * ((int32_t)b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); -#endif -} - -float DistanceCosineFloat::compare(const float *a, const float *b, uint32_t length) const -{ -#ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); -#else - float magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) - { - magA += (a[i]) * (a[i]); - magB += (b[i]) * (b[i]); - scalarProduct += (a[i]) * (b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (scalarProduct / (sqrt(magA) * sqrt(magB))); -#endif -} - -float SlowDistanceCosineUInt8::compare(const uint8_t *a, const uint8_t *b, uint32_t length) const -{ - int magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) - { - magA += ((uint32_t)a[i]) * ((uint32_t)a[i]); - magB += ((uint32_t)b[i]) * ((uint32_t)b[i]); - scalarProduct += ((uint32_t)a[i]) * ((uint32_t)b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); -} - -// -// L2 distance functions. -// - -float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t size) const -{ -#ifdef _WINDOWS -#ifdef USE_AVX2 - __m256 r = _mm256_setzero_ps(); - char *pX = (char *)a, *pY = (char *)b; - while (size >= 32) - { - __m256i r1 = _mm256_subs_epi8(_mm256_loadu_si256((__m256i *)pX), _mm256_loadu_si256((__m256i *)pY)); - r = _mm256_add_ps(r, _mm256_mul_epi8(r1, r1)); - pX += 32; - pY += 32; - size -= 32; - } - while (size > 0) - { - __m128i r2 = _mm_subs_epi8(_mm_loadu_si128((__m128i *)pX), _mm_loadu_si128((__m128i *)pY)); - r = _mm256_add_ps(r, _mm256_mul32_pi8(r2, r2)); - pX += 4; - pY += 4; - size -= 4; - } - r = _mm256_hadd_ps(_mm256_hadd_ps(r, r), r); - return r.m256_f32[0] + r.m256_f32[4]; -#else - int32_t result = 0; -#pragma omp simd reduction(+ : result) aligned(a, b : 8) - for (int32_t i = 0; i < (int32_t)size; i++) - { - result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i])); - } - return (float)result; -#endif -#else - int32_t result = 0; -#pragma omp simd reduction(+ : result) aligned(a, b : 8) - for (int32_t i = 0; i < (int32_t)size; i++) - { - result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i])); - } - return (float)result; -#endif -} - -float DistanceL2UInt8::compare(const uint8_t *a, const uint8_t *b, uint32_t size) const -{ - uint32_t result = 0; -#ifndef _WINDOWS -#pragma omp simd reduction(+ : result) aligned(a, b : 8) -#endif - for (int32_t i = 0; i < (int32_t)size; i++) - { - result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i])); - } - return (float)result; -} - -#ifndef _WINDOWS -float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) const -{ - a = (const float *)__builtin_assume_aligned(a, 32); - b = (const float *)__builtin_assume_aligned(b, 32); -#else -float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) const -{ -#endif - - float result = 0; -#ifdef USE_AVX2 - // assume size is divisible by 8 - uint16_t niters = (uint16_t)(size / 8); - __m256 sum = _mm256_setzero_ps(); - for (uint16_t j = 0; j < niters; j++) - { - // scope is a[8j:8j+7], b[8j:8j+7] - // load a_vec - if (j < (niters - 1)) - { - _mm_prefetch((char *)(a + 8 * (j + 1)), _MM_HINT_T0); - _mm_prefetch((char *)(b + 8 * (j + 1)), _MM_HINT_T0); - } - __m256 a_vec = _mm256_load_ps(a + 8 * j); - // load b_vec - __m256 b_vec = _mm256_load_ps(b + 8 * j); - // a_vec - b_vec - __m256 tmp_vec = _mm256_sub_ps(a_vec, b_vec); - - sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum); - } - - // horizontal add sum - result = _mm256_reduce_add_ps(sum); -#else -#ifndef _WINDOWS -#pragma omp simd reduction(+ : result) aligned(a, b : 32) -#endif - for (int32_t i = 0; i < (int32_t)size; i++) - { - result += (a[i] - b[i]) * (a[i] - b[i]); - } -#endif - return result; -} - -template float SlowDistanceL2::compare(const T *a, const T *b, uint32_t length) const -{ - float result = 0.0f; - for (uint32_t i = 0; i < length; i++) - { - result += ((float)(a[i] - b[i])) * (a[i] - b[i]); - } - return result; -} - -#ifdef _WINDOWS -float AVXDistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t length) const -{ - __m128 r = _mm_setzero_ps(); - __m128i r1; - while (length >= 16) - { - r1 = _mm_subs_epi8(_mm_load_si128((__m128i *)a), _mm_load_si128((__m128i *)b)); - r = _mm_add_ps(r, _mm_mul_epi8(r1)); - a += 16; - b += 16; - length -= 16; - } - r = _mm_hadd_ps(_mm_hadd_ps(r, r), r); - float res = r.m128_f32[0]; - - if (length >= 8) - { - __m128 r2 = _mm_setzero_ps(); - __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *)(a - 8)), _mm_load_si128((__m128i *)(b - 8))); - r2 = _mm_add_ps(r2, _mm_mulhi_epi8(r3)); - a += 8; - b += 8; - length -= 8; - r2 = _mm_hadd_ps(_mm_hadd_ps(r2, r2), r2); - res += r2.m128_f32[0]; - } - - if (length >= 4) - { - __m128 r2 = _mm_setzero_ps(); - __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *)(a - 12)), _mm_load_si128((__m128i *)(b - 12))); - r2 = _mm_add_ps(r2, _mm_mulhi_epi8_shift32(r3)); - res += r2.m128_f32[0] + r2.m128_f32[1]; - } - - return res; -} - -float AVXDistanceL2Float::compare(const float *a, const float *b, uint32_t length) const -{ - __m128 diff, v1, v2; - __m128 sum = _mm_set1_ps(0); - - while (length >= 4) - { - v1 = _mm_loadu_ps(a); - a += 4; - v2 = _mm_loadu_ps(b); - b += 4; - diff = _mm_sub_ps(v1, v2); - sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); - length -= 4; - } - - return sum.m128_f32[0] + sum.m128_f32[1] + sum.m128_f32[2] + sum.m128_f32[3]; -} -#else -float AVXDistanceL2Int8::compare(const int8_t *, const int8_t *, uint32_t) const -{ - return 0; -} -float AVXDistanceL2Float::compare(const float *, const float *, uint32_t) const -{ - return 0; -} -#endif - -template float DistanceInnerProduct::inner_product(const T *a, const T *b, uint32_t size) const -{ - if (!std::is_floating_point::value) - { - diskann::cerr << "ERROR: Inner Product only defined for float currently." << std::endl; - throw diskann::ANNException("ERROR: Inner Product only defined for float currently.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - float result = 0; - -#ifdef __GNUC__ -#ifdef USE_AVX2 -#define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm256_loadu_ps(addr1); \ - tmp2 = _mm256_loadu_ps(addr2); \ - tmp1 = _mm256_mul_ps(tmp1, tmp2); \ - dest = _mm256_add_ps(dest, tmp1); - - __m256 sum; - __m256 l0, l1; - __m256 r0, r1; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t DD = D - DR; - const float *l = (float *)a; - const float *r = (float *)b; - const float *e_l = l + DD; - const float *e_r = r + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; - - sum = _mm256_loadu_ps(unpack); - if (DR) - { - AVX_DOT(e_l, e_r, sum, l0, r0); - } - - for (uint32_t i = 0; i < DD; i += 16, l += 16, r += 16) - { - AVX_DOT(l, r, sum, l0, r0); - AVX_DOT(l + 8, r + 8, sum, l1, r1); - } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + unpack[5] + unpack[6] + unpack[7]; - -#else -#ifdef __SSE2__ -#define SSE_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm128_loadu_ps(addr1); \ - tmp2 = _mm128_loadu_ps(addr2); \ - tmp1 = _mm128_mul_ps(tmp1, tmp2); \ - dest = _mm128_add_ps(dest, tmp1); - __m128 sum; - __m128 l0, l1, l2, l3; - __m128 r0, r1, r2, r3; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t DD = D - DR; - const float *l = a; - const float *r = b; - const float *e_l = l + DD; - const float *e_r = r + DD; - float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; - - sum = _mm_load_ps(unpack); - switch (DR) - { - case 12: - SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); - case 8: - SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); - case 4: - SSE_DOT(e_l, e_r, sum, l0, r0); - default: - break; - } - for (uint32_t i = 0; i < DD; i += 16, l += 16, r += 16) - { - SSE_DOT(l, r, sum, l0, r0); - SSE_DOT(l + 4, r + 4, sum, l1, r1); - SSE_DOT(l + 8, r + 8, sum, l2, r2); - SSE_DOT(l + 12, r + 12, sum, l3, r3); - } - _mm_storeu_ps(unpack, sum); - result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; -#elif __APPLE__ - vDSP_dotpr((float *)a, (vDSP_Stride)1, (float *)b, (vDSP_Stride)1, &result, size); -#else - - float dot0, dot1, dot2, dot3; - const float *last = a + size; - const float *unroll_group = last - 3; - - /* Process 4 items with each loop for efficiency. */ - while (a < unroll_group) - { - dot0 = a[0] * b[0]; - dot1 = a[1] * b[1]; - dot2 = a[2] * b[2]; - dot3 = a[3] * b[3]; - result += dot0 + dot1 + dot2 + dot3; - a += 4; - b += 4; - } - /* Process last 0-3 pixels. Not needed for standard vector lengths. */ - while (a < last) - { - result += *a++ * *b++; - } -#endif -#endif -#endif - return result; -} - -template float DistanceFastL2::compare(const T *a, const T *b, float norm, uint32_t size) const -{ - float result = -2 * DistanceInnerProduct::inner_product(a, b, size); - result += norm; - return result; -} - -template float DistanceFastL2::norm(const T *a, uint32_t size) const -{ - if (!std::is_floating_point::value) - { - diskann::cerr << "ERROR: FastL2 only defined for float currently." << std::endl; - throw diskann::ANNException("ERROR: FastL2 only defined for float currently.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - float result = 0; -#ifdef __GNUC__ -#ifdef __AVX__ -#define AVX_L2NORM(addr, dest, tmp) \ - tmp = _mm256_loadu_ps(addr); \ - tmp = _mm256_mul_ps(tmp, tmp); \ - dest = _mm256_add_ps(dest, tmp); - - __m256 sum; - __m256 l0, l1; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t DD = D - DR; - const float *l = (float *)a; - const float *e_l = l + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; - - sum = _mm256_loadu_ps(unpack); - if (DR) - { - AVX_L2NORM(e_l, sum, l0); - } - for (uint32_t i = 0; i < DD; i += 16, l += 16) - { - AVX_L2NORM(l, sum, l0); - AVX_L2NORM(l + 8, sum, l1); - } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + unpack[5] + unpack[6] + unpack[7]; -#else -#ifdef __SSE2__ -#define SSE_L2NORM(addr, dest, tmp) \ - tmp = _mm128_loadu_ps(addr); \ - tmp = _mm128_mul_ps(tmp, tmp); \ - dest = _mm128_add_ps(dest, tmp); - - __m128 sum; - __m128 l0, l1, l2, l3; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t DD = D - DR; - const float *l = a; - const float *e_l = l + DD; - float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; - - sum = _mm_load_ps(unpack); - switch (DR) - { - case 12: - SSE_L2NORM(e_l + 8, sum, l2); - case 8: - SSE_L2NORM(e_l + 4, sum, l1); - case 4: - SSE_L2NORM(e_l, sum, l0); - default: - break; - } - for (uint32_t i = 0; i < DD; i += 16, l += 16) - { - SSE_L2NORM(l, sum, l0); - SSE_L2NORM(l + 4, sum, l1); - SSE_L2NORM(l + 8, sum, l2); - SSE_L2NORM(l + 12, sum, l3); - } - _mm_storeu_ps(unpack, sum); - result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; -#elif __APPLE__ - vDSP_dotpr((float *)a, 1, (float *)a, 1, &result, size); -#else - float dot0, dot1, dot2, dot3; - const float *last = a + size; - const float *unroll_group = last - 3; - - /* Process 4 items with each loop for efficiency. */ - while (a < unroll_group) - { - dot0 = a[0] * a[0]; - dot1 = a[1] * a[1]; - dot2 = a[2] * a[2]; - dot3 = a[3] * a[3]; - result += dot0 + dot1 + dot2 + dot3; - a += 4; - } - /* Process last 0-3 pixels. Not needed for standard vector lengths. */ - while (a < last) - { - result += (*a) * (*a); - a++; - } -#endif -#endif -#endif - return result; -} - -float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint32_t size) const -{ - float result = 0.0f; -#ifdef __APPLE__ - vDSP_dotpr(a, (vDSP_Stride)1, b, (vDSP_Stride)1, &result, size); -#else -#define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm256_loadu_ps(addr1); \ - tmp2 = _mm256_loadu_ps(addr2); \ - tmp1 = _mm256_mul_ps(tmp1, tmp2); \ - dest = _mm256_add_ps(dest, tmp1); - - __m256 sum; - __m256 l0, l1; - __m256 r0, r1; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t DD = D - DR; - const float *l = (float *)a; - const float *r = (float *)b; - const float *e_l = l + DD; - const float *e_r = r + DD; -#ifndef _WINDOWS - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; -#else - __declspec(align(32)) float unpack[8] = {0, 0, 0, 0, 0, 0, 0, 0}; -#endif - - sum = _mm256_loadu_ps(unpack); - if (DR) - { - AVX_DOT(e_l, e_r, sum, l0, r0); - } - - for (uint32_t i = 0; i < DD; i += 16, l += 16, r += 16) - { - AVX_DOT(l, r, sum, l0, r0); - AVX_DOT(l + 8, r + 8, sum, l1, r1); - } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + unpack[5] + unpack[6] + unpack[7]; -#endif - return -result; -} - -uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension(uint32_t orig_dimension) const -{ - return orig_dimension; -} -bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const -{ - return true; -} -void AVXNormalizedCosineDistanceFloat::preprocess_base_points(float *original_data, const size_t orig_dim, - const size_t num_points) -{ - for (uint32_t i = 0; i < num_points; i++) - { - normalize((float *)(original_data + i * orig_dim), orig_dim); - } -} - -void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, const size_t query_dim, - float *query_scratch) -{ - normalize_and_copy(query_vec, (uint32_t)query_dim, query_scratch); -} - -void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, - float *query_target) const -{ - float norm = get_norm(query_vec, query_dim); - - for (uint32_t i = 0; i < query_dim; i++) - { - query_target[i] = query_vec[i] / norm; - } -} - -// Get the right distance function for the given metric. -template <> diskann::Distance *get_distance_function(diskann::Metric m) -{ - if (m == diskann::Metric::L2) - { - if (Avx2SupportedCPU) - { - diskann::cout << "L2: Using AVX2 distance computation DistanceL2Float" << std::endl; - return new diskann::DistanceL2Float(); - } - else if (AvxSupportedCPU) - { - diskann::cout << "L2: AVX2 not supported. Using AVX distance computation" << std::endl; - return new diskann::AVXDistanceL2Float(); - } - else - { - diskann::cout << "L2: Older CPU. Using slow distance computation" << std::endl; - return new diskann::SlowDistanceL2(); - } - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "Cosine: Using either AVX or AVX2 implementation" << std::endl; - return new diskann::DistanceCosineFloat(); - } - else if (m == diskann::Metric::INNER_PRODUCT) - { - diskann::cout << "Inner product: Using AVX2 implementation " - "AVXDistanceInnerProductFloat" - << std::endl; - return new diskann::AVXDistanceInnerProductFloat(); - } - else if (m == diskann::Metric::FAST_L2) - { - diskann::cout << "Fast_L2: Using AVX2 implementation with norm " - "memoization DistanceFastL2" - << std::endl; - return new diskann::DistanceFastL2(); - } - else - { - std::stringstream stream; - stream << "Only L2, cosine, and inner product supported for floating " - "point vectors as of now." - << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template <> diskann::Distance *get_distance_function(diskann::Metric m) -{ - if (m == diskann::Metric::L2) - { - if (Avx2SupportedCPU) - { - diskann::cout << "Using AVX2 distance computation DistanceL2Int8." << std::endl; - return new diskann::DistanceL2Int8(); - } - else if (AvxSupportedCPU) - { - diskann::cout << "AVX2 not supported. Using AVX distance computation" << std::endl; - return new diskann::AVXDistanceL2Int8(); - } - else - { - diskann::cout << "Older CPU. Using slow distance computation " - "SlowDistanceL2Int." - << std::endl; - return new diskann::SlowDistanceL2(); - } - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "Using either AVX or AVX2 for Cosine similarity " - "DistanceCosineInt8." - << std::endl; - return new diskann::DistanceCosineInt8(); - } - else - { - std::stringstream stream; - stream << "Only L2 and cosine supported for signed byte vectors." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template <> diskann::Distance *get_distance_function(diskann::Metric m) -{ - if (m == diskann::Metric::L2) - { -#ifdef _WINDOWS - diskann::cout << "WARNING: AVX/AVX2 distance function not defined for Uint8. " - "Using " - "slow version. " - "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." - << std::endl; -#endif - return new diskann::DistanceL2UInt8(); - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "AVX/AVX2 distance function not defined for Uint8. Using " - "slow version SlowDistanceCosineUint8() " - "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." - << std::endl; - return new diskann::SlowDistanceCosineUInt8(); - } - else - { - std::stringstream stream; - stream << "Only L2 and cosine supported for uint32_t byte vectors." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template DISKANN_DLLEXPORT class DistanceInnerProduct; -template DISKANN_DLLEXPORT class DistanceInnerProduct; -template DISKANN_DLLEXPORT class DistanceInnerProduct; - -template DISKANN_DLLEXPORT class DistanceFastL2; -template DISKANN_DLLEXPORT class DistanceFastL2; -template DISKANN_DLLEXPORT class DistanceFastL2; - -template DISKANN_DLLEXPORT class SlowDistanceL2; -template DISKANN_DLLEXPORT class SlowDistanceL2; -template DISKANN_DLLEXPORT class SlowDistanceL2; - -// template DISKANN_DLLEXPORT Distance *get_distance_function(Metric m); -// template DISKANN_DLLEXPORT Distance *get_distance_function(Metric m); -// template DISKANN_DLLEXPORT Distance *get_distance_function(Metric m); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/dll/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/src/dll/CMakeLists.txt deleted file mode 100644 index 096d1b7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/dll/CMakeLists.txt +++ /dev/null @@ -1,35 +0,0 @@ -#Copyright(c) Microsoft Corporation.All rights reserved. -#Licensed under the MIT license. - -add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../distance.cpp ../pq_l2_distance.cpp ../memory_mapper.cpp ../index.cpp - ../in_mem_data_store.cpp ../pq_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp - ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp) - -set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") - -set(DISKANN_DLL_IMPLIB "${TARGET_DIR}/${PROJECT_NAME}.lib") - -if (NOT PYBIND) - target_compile_definitions(${PROJECT_NAME} PRIVATE DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS DISKANN_BUILD) -endif() -target_compile_definitions(${PROJECT_NAME} PRIVATE _USRDLL _WINDLL) -target_compile_options(${PROJECT_NAME} PRIVATE /GL) -target_include_directories(${PROJECT_NAME} PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES}) - -target_link_options(${PROJECT_NAME} PRIVATE /DLL /IMPLIB:${DISKANN_DLL_IMPLIB} /LTCG) -target_link_libraries(${PROJECT_NAME} PRIVATE ${DISKANN_MKL_LINK_LIBRARIES}) -target_link_libraries(${PROJECT_NAME} PRIVATE synchronization.lib) - -if (DISKANN_DLL_TCMALLOC_LINK_OPTIONS) - target_link_libraries(${PROJECT_NAME} PUBLIC ${DISKANN_DLL_TCMALLOC_LINK_OPTIONS}) -endif() - -# Copy OpenMP DLL and PDB. -set(RUNTIME_FILES_TO_COPY ${OPENMP_WINDOWS_RUNTIME_FILES} ${TCMALLOC_WINDOWS_RUNTIME_FILES}) - -foreach(RUNTIME_FILE ${RUNTIME_FILES_TO_COPY}) - add_custom_command(TARGET ${PROJECT_NAME} - POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy "${RUNTIME_FILE}" "${TARGET_DIR}") -endforeach() \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/dll/dllmain.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/dll/dllmain.cpp deleted file mode 100644 index 9f5ce44..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/dll/dllmain.cpp +++ /dev/null @@ -1,15 +0,0 @@ -// dllmain.cpp : Defines the entry point for the DLL application. -#include - -BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) -{ - switch (ul_reason_for_call) - { - case DLL_PROCESS_ATTACH: - case DLL_THREAD_ATTACH: - case DLL_THREAD_DETACH: - case DLL_PROCESS_DETACH: - break; - } - return TRUE; -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/filter_utils.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/filter_utils.cpp deleted file mode 100644 index 09d740e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/filter_utils.cpp +++ /dev/null @@ -1,355 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include - -#include -#include "filter_utils.h" -#include "index.h" -#include "parameters.h" -#include "utils.h" - -namespace diskann -{ -/* - * Using passed in parameters and files generated from step 3, - * builds a vanilla diskANN index for each label. - * - * Each index is saved under the following path: - * final_index_path_prefix + "_" + label - */ -template -void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, uint32_t R, - uint32_t L, float alpha, uint32_t num_threads) -{ - diskann::IndexWriteParameters label_index_build_parameters = diskann::IndexWriteParametersBuilder(L, R) - .with_saturate_graph(false) - .with_alpha(alpha) - .with_num_threads(num_threads) - .build(); - - std::cout << "Generating indices per label..." << std::endl; - // for each label, build an index on resp. points - double total_indexing_time = 0.0, indexing_percentage = 0.0; - std::cout.setstate(std::ios_base::failbit); - diskann::cout.setstate(std::ios_base::failbit); - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - - size_t number_of_label_points, dimension; - diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension); - - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, - std::make_shared(label_index_build_parameters), nullptr, - 0, false, false, false, false, 0, false); - - auto index_build_timer = std::chrono::high_resolution_clock::now(); - index.build(curr_label_input_data_path.c_str(), number_of_label_points); - std::chrono::duration current_indexing_time = - std::chrono::high_resolution_clock::now() - index_build_timer; - - total_indexing_time += current_indexing_time.count(); - indexing_percentage += (1 / (double)all_labels.size()); - print_progress(indexing_percentage); - - index.save(curr_label_index_path.c_str()); - } - std::cout.clear(); - diskann::cout.clear(); - - std::cout << "\nDone. Generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl; -} - -// for use on systems without writev (i.e. Windows) -template -tsl::robin_map> generate_label_specific_vector_files_compat( - path input_data_path, tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) -{ - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - std::ifstream input_data_stream(input_data_path); - - uint32_t number_of_points, dimension; - input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); - input_data_stream.read((char *)&dimension, sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - if (number_of_points != point_ids_to_labels.size()) - { - std::cerr << "Error: number of points in labels file and data file differ." << std::endl; - throw; - } - - tsl::robin_map labels_to_vectors; - tsl::robin_map labels_to_curr_vector; - tsl::robin_map> label_id_to_orig_id; - - for (const auto &lbl : all_labels) - { - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); - if (vectors == nullptr) - { - throw; - } - labels_to_vectors[lbl] = vectors; - labels_to_curr_vector[lbl] = 0; - label_id_to_orig_id[lbl].reserve(number_of_label_pts); - } - - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) - { - char *curr_vector = (char *)malloc(VECTOR_SIZE); - input_data_stream.read(curr_vector, VECTOR_SIZE); - for (const auto &lbl : point_ids_to_labels[point_id]) - { - char *curr_label_vector_ptr = labels_to_vectors[lbl] + (labels_to_curr_vector[lbl] * VECTOR_SIZE); - memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE); - labels_to_curr_vector[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); - } - free(curr_vector); - } - - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - - std::ofstream label_file_stream; - label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); - label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); - label_file_stream.write((char *)&dimension, sizeof(uint32_t)); - label_file_stream.write((char *)labels_to_vectors[lbl], number_of_label_pts * VECTOR_SIZE); - - label_file_stream.close(); - free(labels_to_vectors[lbl]); - } - input_data_stream.close(); - - std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; - - return label_id_to_orig_id; -} - -/* - * Manually loads a graph index in from a given file. - * - * Returns both the graph index and the size of the file in bytes. - */ -load_label_index_return_values load_label_index(path label_index_path, uint32_t label_number_of_points) -{ - std::ifstream label_index_stream; - label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_index_stream.open(label_index_path, std::ios::binary); - - uint64_t index_file_size, index_num_frozen_points; - uint32_t index_max_observed_degree, index_entry_point; - const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); - label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); - label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); - label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); - size_t bytes_read = INDEX_METADATA; - - std::vector> label_index(label_number_of_points); - uint32_t nodes_read = 0; - while (bytes_read != index_file_size) - { - uint32_t current_node_num_neighbors; - label_index_stream.read((char *)¤t_node_num_neighbors, sizeof(uint32_t)); - nodes_read++; - - std::vector current_node_neighbors(current_node_num_neighbors); - label_index_stream.read((char *)current_node_neighbors.data(), current_node_num_neighbors * sizeof(uint32_t)); - label_index[nodes_read - 1].swap(current_node_neighbors); - bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); - } - - return std::make_tuple(label_index, index_file_size); -} - -/* - * Parses the label datafile, which has comma-separated labels on - * each line. Line i corresponds to point id i. - * - * Returns three objects via std::tuple: - * 1. map: key is point id, value is vector of labels said point has - * 2. map: key is label, value is number of points with the label - * 3. the label universe as a set - */ -parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label) -{ - std::ifstream label_data_stream(label_data_path); - std::string line, token; - uint32_t line_cnt = 0; - - // allows us to reserve space for the points_to_labels vector - while (std::getline(label_data_stream, line)) - line_cnt++; - label_data_stream.clear(); - label_data_stream.seekg(0, std::ios::beg); - - // values to return - std::vector point_ids_to_labels(line_cnt); - tsl::robin_map labels_to_number_of_points; - label_set all_labels; - - std::vector points_with_universal_label; - line_cnt = 0; - while (std::getline(label_data_stream, line)) - { - std::istringstream current_labels_comma_separated(line); - label_set current_labels; - - // get point id - uint32_t point_id = line_cnt; - - // parse comma separated labels - bool current_universal_label_check = false; - while (getline(current_labels_comma_separated, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - - // if token is empty, there's no labels for the point - if (token == universal_label) - { - points_with_universal_label.push_back(point_id); - current_universal_label_check = true; - } - else - { - all_labels.insert(token); - current_labels.insert(token); - labels_to_number_of_points[token]++; - } - } - - if (current_labels.size() <= 0 && !current_universal_label_check) - { - std::cerr << "Error: " << point_id << " has no labels." << std::endl; - exit(-1); - } - point_ids_to_labels[point_id] = current_labels; - line_cnt++; - } - - // for every point with universal label, set its label set to all labels - // also, increment the count for number of points a label has - for (const auto &point_id : points_with_universal_label) - { - point_ids_to_labels[point_id] = all_labels; - for (const auto &lbl : all_labels) - labels_to_number_of_points[lbl]++; - } - - std::cout << "Identified " << all_labels.size() << " distinct label(s) for " << point_ids_to_labels.size() - << " points\n" - << std::endl; - - return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels); -} - -/* - * A templated function to parse a file of labels that are already represented - * as either uint16_t or uint32_t - * - * Returns two objects via std::tuple: - * 1. a vector of vectors of labels, where the outer vector is indexed by point id - * 2. a set of all labels - */ -template -std::tuple>, tsl::robin_set> parse_formatted_label_file(std::string label_file) -{ - std::vector> pts_to_labels; - tsl::robin_set labels; - - // Format of Label txt file: filters with comma separators - std::ifstream infile(label_file); - if (infile.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); - } - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(infile, line)) - { - line_cnt++; - } - pts_to_labels.resize(line_cnt, std::vector()); - - infile.clear(); - infile.seekg(0, std::ios::beg); - line_cnt = 0; - - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = static_cast(std::stoul(token)); - lbls.push_back(token_as_num); - labels.insert(token_as_num); - } - if (lbls.size() <= 0) - { - diskann::cout << "No label found"; - exit(-1); - } - std::sort(lbls.begin(), lbls.end()); - pts_to_labels[line_cnt] = lbls; - line_cnt++; - } - diskann::cout << "Identified " << labels.size() << " distinct label(s)" << std::endl; - - return std::make_tuple(pts_to_labels, labels); -} - -template DISKANN_DLLEXPORT std::tuple>, tsl::robin_set> -parse_formatted_label_file(path label_file); - -template DISKANN_DLLEXPORT std::tuple>, tsl::robin_set> -parse_formatted_label_file(path label_file); - -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, float alpha, - uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, - float alpha, uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, - float alpha, uint32_t num_threads); - -template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); -template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); -template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_data_store.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_data_store.cpp deleted file mode 100644 index cc7acf6..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_data_store.cpp +++ /dev/null @@ -1,401 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include "abstract_scratch.h" -#include "in_mem_data_store.h" - -#include "utils.h" - -namespace diskann -{ - -template -InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, - std::unique_ptr> distance_fn) - : AbstractDataStore(num_points, dim), _distance_fn(std::move(distance_fn)) -{ - _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); - alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); -} - -template InMemDataStore::~InMemDataStore() -{ - if (_data != nullptr) - { - aligned_free(this->_data); - } -} - -template size_t InMemDataStore::get_aligned_dim() const -{ - return _aligned_dim; -} - -template size_t InMemDataStore::get_alignment_factor() const -{ - return _distance_fn->get_required_alignment(); -} - -template location_t InMemDataStore::load(const std::string &filename) -{ - return load_impl(filename); -} - -#ifdef EXEC_ENV_OLS -template location_t InMemDataStore::load_impl(AlignedFileReader &reader) -{ - size_t file_dim, file_num_points; - - diskann::get_bin_metadata(reader, file_num_points, file_dim); - - if (file_dim != this->_dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << this->_dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_num_points > this->capacity()) - { - this->resize((location_t)file_num_points); - } - copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); - - return (location_t)file_num_points; -} -#endif - -template location_t InMemDataStore::load_impl(const std::string &filename) -{ - size_t file_dim, file_num_points; - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); - - if (file_dim != this->_dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << this->_dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_num_points > this->capacity()) - { - this->resize((location_t)file_num_points); - } - - copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); - - return (location_t)file_num_points; -} - -template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) -{ - return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); -} - -template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) -{ - memset(_data, 0, _aligned_dim * sizeof(data_t) * num_pts); - for (location_t i = 0; i < num_pts; i++) - { - std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); - } - - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data, this->_aligned_dim, num_pts); - } -} - -template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) -{ - size_t npts, ndim; - copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); - - if ((location_t)npts > this->capacity()) - { - std::stringstream ss; - ss << "Number of points in the file: " << filename - << " is greater than the capacity of data store: " << this->capacity() - << ". Must invoke resize before calling populate_data()" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if ((location_t)ndim != this->get_dims()) - { - std::stringstream ss; - ss << "Number of dimensions of a point in the file: " << filename - << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data, this->_aligned_dim, this->capacity()); - } -} - -template -void InMemDataStore::extract_data_to_bin(const std::string &filename, const location_t num_points) -{ - save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); -} - -template void InMemDataStore::get_vector(const location_t i, data_t *dest) const -{ - // REFACTOR TODO: Should we denormalize and return values? - memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); -} - -template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) -{ - size_t offset_in_data = loc * _aligned_dim; - memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); - memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, 1); - } -} - -template void InMemDataStore::prefetch_vector(const location_t loc) -{ - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc * sizeof(data_t), - sizeof(data_t) * _aligned_dim); -} - -template -void InMemDataStore::preprocess_query(const data_t *query, AbstractScratch *query_scratch) const -{ - if (query_scratch != nullptr) - { - memcpy(query_scratch->aligned_query_T(), query, sizeof(data_t) * this->get_dims()); - } - else - { - std::stringstream ss; - ss << "In InMemDataStore::preprocess_query: Query scratch is null"; - diskann::cerr << ss.str() << std::endl; - throw diskann::ANNException(ss.str(), -1); - } -} - -template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const -{ - return _distance_fn->compare(query, _data + _aligned_dim * loc, (uint32_t)_aligned_dim); -} - -template -void InMemDataStore::get_distance(const data_t *query, const location_t *locations, - const uint32_t location_count, float *distances, - AbstractScratch *scratch_space) const -{ - for (location_t i = 0; i < location_count; i++) - { - distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, (uint32_t)this->_aligned_dim); - } -} - -template -float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const -{ - return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, - (uint32_t)this->_aligned_dim); -} - -template -void InMemDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const -{ - for (int i = 0; i < ids.size(); i++) - { - distances[i] = - _distance_fn->compare(preprocessed_query, _data + ids[i] * _aligned_dim, (uint32_t)this->_aligned_dim); - } -} - -template location_t InMemDataStore::expand(const location_t new_size) -{ - if (new_size == this->capacity()) - { - return this->capacity(); - } - else if (new_size < this->capacity()) - { - std::stringstream ss; - ss << "Cannot 'expand' datastore when new capacity (" << new_size << ") < existing capacity(" - << this->capacity() << ")" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } -#ifndef _WINDOWS - data_t *new_data; - alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - memcpy(new_data, _data, this->capacity() * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; -#else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); -#endif - this->_capacity = new_size; - return this->_capacity; -} - -template location_t InMemDataStore::shrink(const location_t new_size) -{ - if (new_size == this->capacity()) - { - return this->capacity(); - } - else if (new_size > this->capacity()) - { - std::stringstream ss; - ss << "Cannot 'shrink' datastore when new capacity (" << new_size << ") > existing capacity(" - << this->capacity() << ")" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } -#ifndef _WINDOWS - data_t *new_data; - alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - memcpy(new_data, _data, new_size * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; -#else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); -#endif - this->_capacity = new_size; - return this->_capacity; -} - -template -void InMemDataStore::move_vectors(const location_t old_location_start, const location_t new_location_start, - const location_t num_locations) -{ - if (num_locations == 0 || old_location_start == new_location_start) - { - return; - } - - /* // Update pointers to the moved nodes. Note: the computation is correct - even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - */ - // The [start, end) interval which will contain obsolete points to be - // cleared. - uint32_t mem_clear_loc_start = old_location_start; - uint32_t mem_clear_loc_end_limit = old_location_start + num_locations; - - if (new_location_start < old_location_start) - { - // If ranges are overlapping, make sure not to clear the newly copied - // data. - if (mem_clear_loc_start < new_location_start + num_locations) - { - // Clear only after the end of the new range. - mem_clear_loc_start = new_location_start + num_locations; - } - } - else - { - // If ranges are overlapping, make sure not to clear the newly copied - // data. - if (mem_clear_loc_end_limit > new_location_start) - { - // Clear only up to the beginning of the new range. - mem_clear_loc_end_limit = new_location_start; - } - } - - // Use memmove to handle overlapping ranges. - copy_vectors(old_location_start, new_location_start, num_locations); - memset(_data + _aligned_dim * mem_clear_loc_start, 0, - sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); -} - -template -void InMemDataStore::copy_vectors(const location_t from_loc, const location_t to_loc, - const location_t num_points) -{ - assert(from_loc < this->_capacity); - assert(to_loc < this->_capacity); - assert(num_points < this->_capacity); - memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); -} - -template location_t InMemDataStore::calculate_medoid() const -{ - // allocate and init centroid - float *center = new float[_aligned_dim]; - for (size_t j = 0; j < _aligned_dim; j++) - center[j] = 0; - - for (size_t i = 0; i < this->capacity(); i++) - for (size_t j = 0; j < _aligned_dim; j++) - center[j] += (float)_data[i * _aligned_dim + j]; - - for (size_t j = 0; j < _aligned_dim; j++) - center[j] /= (float)this->capacity(); - - // compute all to one distance - float *distances = new float[this->capacity()]; - - // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. - // Problem is that we need to pass num_threads here, it is not clear - // if data store must be aware of threads! - // #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)this->capacity(); i++) - { - // extract point and distance reference - float &dist = distances[i]; - const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); - dist = 0; - float diff = 0; - for (size_t j = 0; j < _aligned_dim; j++) - { - diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); - dist += diff; - } - } - // find imin - uint32_t min_idx = 0; - float min_dist = distances[0]; - for (uint32_t i = 1; i < this->capacity(); i++) - { - if (distances[i] < min_dist) - { - min_idx = i; - min_dist = distances[i]; - } - } - - delete[] distances; - delete[] center; - return min_idx; -} - -template Distance *InMemDataStore::get_dist_fn() const -{ - return this->_distance_fn.get(); -} - -template DISKANN_DLLEXPORT class InMemDataStore; -template DISKANN_DLLEXPORT class InMemDataStore; -template DISKANN_DLLEXPORT class InMemDataStore; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_graph_store.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_graph_store.cpp deleted file mode 100644 index c12b251..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/in_mem_graph_store.cpp +++ /dev/null @@ -1,242 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "in_mem_graph_store.h" -#include "utils.h" - -namespace diskann -{ -InMemGraphStore::InMemGraphStore(const size_t total_pts, const size_t reserve_graph_degree) - : AbstractGraphStore(total_pts, reserve_graph_degree) -{ - this->resize_graph(total_pts); - for (size_t i = 0; i < total_pts; i++) - { - _graph[i].reserve(reserve_graph_degree); - } -} - -std::tuple InMemGraphStore::load(const std::string &index_path_prefix, - const size_t num_points) -{ - return load_impl(index_path_prefix, num_points); -} -int InMemGraphStore::store(const std::string &index_path_prefix, const size_t num_points, - const size_t num_frozen_points, const uint32_t start) -{ - return save_graph(index_path_prefix, num_points, num_frozen_points, start); -} -const std::vector &InMemGraphStore::get_neighbours(const location_t i) const -{ - return _graph.at(i); -} - -void InMemGraphStore::add_neighbour(const location_t i, location_t neighbour_id) -{ - _graph[i].emplace_back(neighbour_id); - if (_max_observed_degree < _graph[i].size()) - { - _max_observed_degree = (uint32_t)(_graph[i].size()); - } -} - -void InMemGraphStore::clear_neighbours(const location_t i) -{ - _graph[i].clear(); -}; -void InMemGraphStore::swap_neighbours(const location_t a, location_t b) -{ - _graph[a].swap(_graph[b]); -}; - -void InMemGraphStore::set_neighbours(const location_t i, std::vector &neighbours) -{ - _graph[i].assign(neighbours.begin(), neighbours.end()); - if (_max_observed_degree < neighbours.size()) - { - _max_observed_degree = (uint32_t)(neighbours.size()); - } -} - -size_t InMemGraphStore::resize_graph(const size_t new_size) -{ - _graph.resize(new_size); - set_total_points(new_size); - return _graph.size(); -} - -void InMemGraphStore::clear_graph() -{ - _graph.clear(); -} - -#ifdef EXEC_ENV_OLS -std::tuple InMemGraphStore::load_impl(AlignedFileReader &reader, size_t expected_num_points) -{ - size_t expected_file_size; - size_t file_frozen_pts; - uint32_t start; - - auto max_points = get_max_points(); - int header_size = 2 * sizeof(size_t) + 2 * sizeof(uint32_t); - std::unique_ptr header = std::make_unique(header_size); - read_array(reader, header.get(), header_size); - - expected_file_size = *((size_t *)header.get()); - _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); - start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t))); - file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t))); - - diskann::cout << "From graph header, expected_file_size: " << expected_file_size - << ", _max_observed_degree: " << _max_observed_degree << ", _start: " << start - << ", file_frozen_pts: " << file_frozen_pts << std::endl; - - diskann::cout << "Loading vamana graph from reader..." << std::flush; - - // If user provides more points than max_points - // resize the _graph to the larger size. - if (get_total_points() < expected_num_points) - { - diskann::cout << "resizing graph to " << expected_num_points << std::endl; - this->resize_graph(expected_num_points); - } - - uint32_t nodes_read = 0; - size_t cc = 0; - size_t graph_offset = header_size; - while (nodes_read < expected_num_points) - { - uint32_t k; - read_value(reader, k, graph_offset); - graph_offset += sizeof(uint32_t); - std::vector tmp(k); - tmp.reserve(k); - read_array(reader, tmp.data(), k, graph_offset); - graph_offset += k * sizeof(uint32_t); - cc += k; - _graph[nodes_read].swap(tmp); - nodes_read++; - if (nodes_read % 1000000 == 0) - { - diskann::cout << "." << std::flush; - } - if (k > _max_range_of_graph) - { - _max_range_of_graph = k; - } - } - - diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc << " out-edges, _start is set to " << start - << std::endl; - return std::make_tuple(nodes_read, start, file_frozen_pts); -} -#endif - -std::tuple InMemGraphStore::load_impl(const std::string &filename, - size_t expected_num_points) -{ - size_t expected_file_size; - size_t file_frozen_pts; - uint32_t start; - size_t file_offset = 0; // will need this for single file format support - - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); - in.open(filename, std::ios::binary); - in.seekg(file_offset, in.beg); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&_max_observed_degree, sizeof(uint32_t)); - in.read((char *)&start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); - size_t vamana_metadata_size = sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); - - diskann::cout << "From graph header, expected_file_size: " << expected_file_size - << ", _max_observed_degree: " << _max_observed_degree << ", _start: " << start - << ", file_frozen_pts: " << file_frozen_pts << std::endl; - - diskann::cout << "Loading vamana graph " << filename << "..." << std::flush; - - // If user provides more points than max_points - // resize the _graph to the larger size. - if (get_total_points() < expected_num_points) - { - diskann::cout << "resizing graph to " << expected_num_points << std::endl; - this->resize_graph(expected_num_points); - } - - size_t bytes_read = vamana_metadata_size; - size_t cc = 0; - uint32_t nodes_read = 0; - while (bytes_read != expected_file_size) - { - uint32_t k; - in.read((char *)&k, sizeof(uint32_t)); - - if (k == 0) - { - diskann::cerr << "ERROR: Point found with no out-neighbours, point#" << nodes_read << std::endl; - } - - cc += k; - ++nodes_read; - std::vector tmp(k); - tmp.reserve(k); - in.read((char *)tmp.data(), k * sizeof(uint32_t)); - _graph[nodes_read - 1].swap(tmp); - bytes_read += sizeof(uint32_t) * ((size_t)k + 1); - if (nodes_read % 10000000 == 0) - diskann::cout << "." << std::flush; - if (k > _max_range_of_graph) - { - _max_range_of_graph = k; - } - } - - diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc << " out-edges, _start is set to " << start - << std::endl; - return std::make_tuple(nodes_read, start, file_frozen_pts); -} - -int InMemGraphStore::save_graph(const std::string &index_path_prefix, const size_t num_points, - const size_t num_frozen_points, const uint32_t start) -{ - std::ofstream out; - open_file_to_write(out, index_path_prefix); - - size_t file_offset = 0; - out.seekp(file_offset, out.beg); - size_t index_size = 24; - uint32_t max_degree = 0; - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&_max_observed_degree, sizeof(uint32_t)); - uint32_t ep_u32 = start; - out.write((char *)&ep_u32, sizeof(uint32_t)); - out.write((char *)&num_frozen_points, sizeof(size_t)); - - // Note: num_points = _nd + _num_frozen_points - for (uint32_t i = 0; i < num_points; i++) - { - uint32_t GK = (uint32_t)_graph[i].size(); - out.write((char *)&GK, sizeof(uint32_t)); - out.write((char *)_graph[i].data(), GK * sizeof(uint32_t)); - max_degree = _graph[i].size() > max_degree ? (uint32_t)_graph[i].size() : max_degree; - index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); - } - out.seekp(file_offset, out.beg); - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&max_degree, sizeof(uint32_t)); - out.close(); - return (int)index_size; -} - -size_t InMemGraphStore::get_max_range_of_graph() -{ - return _max_range_of_graph; -} - -uint32_t InMemGraphStore::get_max_observed_degree() -{ - return _max_observed_degree; -} - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/index.cpp deleted file mode 100644 index 7f26288..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/index.cpp +++ /dev/null @@ -1,3524 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include - -#include - -#include "boost/dynamic_bitset.hpp" -#include "index_factory.h" -#include "memory_mapper.h" -#include "timer.h" -#include "tsl/robin_map.h" -#include "tsl/robin_set.h" -#include "windows_customizations.h" -#include "tag_uint128.h" -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) -#include "gperftools/malloc_extension.h" -#endif - -#ifdef _WINDOWS -#include -#endif - -#include "index.h" - -#define MAX_POINTS_FOR_USING_BITSET 10000000 - -namespace diskann -{ -// Initialize an index with metric m, load the data of type T with filename -// (bin), and initialize max_points -template -Index::Index(const IndexConfig &index_config, std::shared_ptr> data_store, - std::unique_ptr graph_store, - std::shared_ptr> pq_data_store) - : _dist_metric(index_config.metric), _dim(index_config.dimension), _max_points(index_config.max_points), - _num_frozen_pts(index_config.num_frozen_pts), _dynamic_index(index_config.dynamic_index), - _enable_tags(index_config.enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), - _pq_dist(index_config.pq_dist_build), _use_opq(index_config.use_opq), - _filtered_index(index_config.filtered_index), _num_pq_chunks(index_config.num_pq_chunks), - _delete_set(new tsl::robin_set), _conc_consolidate(index_config.concurrent_consolidate) -{ - if (_dynamic_index && !_enable_tags) - { - throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_pq_dist) - { - if (_dynamic_index) - throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (_dist_metric == diskann::Metric::INNER_PRODUCT) - throw ANNException("ERROR: Inner product metrics not yet supported " - "with PQ distance " - "base index", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_dynamic_index && _num_frozen_pts == 0) - { - _num_frozen_pts = 1; - } - // Sanity check. While logically it is correct, max_points = 0 causes - // downstream problems. - if (_max_points == 0) - { - _max_points = 1; - } - const size_t total_internal_points = _max_points + _num_frozen_pts; - - _start = (uint32_t)_max_points; - - _data_store = data_store; - _pq_data_store = pq_data_store; - _graph_store = std::move(graph_store); - - _locks = std::vector(total_internal_points); - if (_enable_tags) - { - _location_to_tag.reserve(total_internal_points); - _tag_to_location.reserve(total_internal_points); - } - - if (_dynamic_index) - { - this->enable_delete(); // enable delete by default for dynamic index - if (_filtered_index) - { - _location_to_labels.resize(total_internal_points); - } - } - - if (index_config.index_write_params != nullptr) - { - _indexingQueueSize = index_config.index_write_params->search_list_size; - _indexingRange = index_config.index_write_params->max_degree; - _indexingMaxC = index_config.index_write_params->max_occlusion_size; - _indexingAlpha = index_config.index_write_params->alpha; - _filterIndexingQueueSize = index_config.index_write_params->filter_list_size; - _indexingThreads = index_config.index_write_params->num_threads; - _saturate_graph = index_config.index_write_params->saturate_graph; - - if (index_config.index_search_params != nullptr) - { - uint32_t num_scratch_spaces = index_config.index_search_params->num_search_threads + _indexingThreads; - initialize_query_scratch(num_scratch_spaces, index_config.index_search_params->initial_search_list_size, - _indexingQueueSize, _indexingRange, _indexingMaxC, _data_store->get_dims()); - } - } -} - -template -Index::Index(Metric m, const size_t dim, const size_t max_points, - const std::shared_ptr index_parameters, - const std::shared_ptr index_search_params, const size_t num_frozen_pts, - const bool dynamic_index, const bool enable_tags, const bool concurrent_consolidate, - const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq, - const bool filtered_index) - : Index( - IndexConfigBuilder() - .with_metric(m) - .with_dimension(dim) - .with_max_points(max_points) - .with_index_write_params(index_parameters) - .with_index_search_params(index_search_params) - .with_num_frozen_pts(num_frozen_pts) - .is_dynamic_index(dynamic_index) - .is_enable_tags(enable_tags) - .is_concurrent_consolidate(concurrent_consolidate) - .is_pq_dist_build(pq_dist_build) - .with_num_pq_chunks(num_pq_chunks) - .is_use_opq(use_opq) - .is_filtered(filtered_index) - .with_data_type(diskann_type_to_name()) - .build(), - IndexFactory::construct_datastore(DataStoreStrategy::MEMORY, - (max_points == 0 ? (size_t)1 : max_points) + - (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), - dim, m), - IndexFactory::construct_graphstore(GraphStoreStrategy::MEMORY, - (max_points == 0 ? (size_t)1 : max_points) + - (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), - (size_t)((index_parameters == nullptr ? 0 : index_parameters->max_degree) * - defaults::GRAPH_SLACK_FACTOR * 1.05))) -{ - if (_pq_dist) - { - _pq_data_store = IndexFactory::construct_pq_datastore(DataStoreStrategy::MEMORY, max_points + num_frozen_pts, - dim, m, num_pq_chunks, use_opq); - } - else - { - _pq_data_store = _data_store; - } -} - -template Index::~Index() -{ - // Ensure that no other activity is happening before dtor() - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - for (auto &lock : _locks) - { - LockGuard lg(lock); - } - - if (_opt_graph != nullptr) - { - delete[] _opt_graph; - } - - if (!_query_scratch.empty()) - { - ScratchStoreManager> manager(_query_scratch); - manager.destroy(); - } -} - -template -void Index::initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, - uint32_t r, uint32_t maxc, size_t dim) -{ - for (uint32_t i = 0; i < num_threads; i++) - { - auto scratch = new InMemQueryScratch(search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), - _data_store->get_alignment_factor(), _pq_dist); - _query_scratch.push(scratch); - } -} - -template size_t Index::save_tags(std::string tags_file) -{ - if (!_enable_tags) - { - diskann::cout << "Not saving tags as they are not enabled." << std::endl; - return 0; - } - - size_t tag_bytes_written; - TagT *tag_data = new TagT[_nd + _num_frozen_pts]; - for (uint32_t i = 0; i < _nd; i++) - { - TagT tag; - if (_location_to_tag.try_get(i, tag)) - { - tag_data[i] = tag; - } - else - { - // catering to future when tagT can be any type. - std::memset((char *)&tag_data[i], 0, sizeof(TagT)); - } - } - if (_num_frozen_pts > 0) - { - std::memset((char *)&tag_data[_start], 0, sizeof(TagT) * _num_frozen_pts); - } - try - { - tag_bytes_written = save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1); - } - catch (std::system_error &e) - { - throw FileException(tags_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - delete[] tag_data; - return tag_bytes_written; -} - -template size_t Index::save_data(std::string data_file) -{ - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_pts is the valid - // location limit. - return _data_store->save(data_file, (location_t)(_nd + _num_frozen_pts)); -} - -// save the graph index on a file as an adjacency list. For each point, -// first store the number of neighbors, and then the neighbor list (each as -// 4 byte uint32_t) -template size_t Index::save_graph(std::string graph_file) -{ - return _graph_store->store(graph_file, _nd + _num_frozen_pts, _num_frozen_pts, _start); -} - -template -size_t Index::save_delete_list(const std::string &filename) -{ - if (_delete_set->size() == 0) - { - return 0; - } - std::unique_ptr delete_list = std::make_unique(_delete_set->size()); - uint32_t i = 0; - for (auto &del : *_delete_set) - { - delete_list[i++] = del; - } - return save_bin(filename, delete_list.get(), _delete_set->size(), 1); -} - -template -void Index::save(const char *filename, bool compact_before_save) -{ - diskann::Timer timer; - - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - if (compact_before_save) - { - compact_data(); - compact_frozen_point(); - } - else - { - if (!_data_compacted) - { - throw ANNException("Index save for non-compacted index is not yet implemented", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - } - - if (!_save_as_one_file) - { - if (_filtered_index) - { - if (_label_to_start_id.size() > 0) - { - std::ofstream medoid_writer(std::string(filename) + "_labels_to_medoids.txt"); - if (medoid_writer.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); - } - for (auto iter : _label_to_start_id) - { - medoid_writer << iter.first << ", " << iter.second << std::endl; - } - medoid_writer.close(); - } - - if (_use_universal_label) - { - std::ofstream universal_label_writer(std::string(filename) + "_universal_label.txt"); - assert(universal_label_writer.is_open()); - universal_label_writer << _universal_label << std::endl; - universal_label_writer.close(); - } - - if (_location_to_labels.size() > 0) - { - std::ofstream label_writer(std::string(filename) + "_labels.txt"); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) - { - for (uint32_t j = 0; j + 1 < _location_to_labels[i].size(); j++) - { - label_writer << _location_to_labels[i][j] << ","; - } - if (_location_to_labels[i].size() != 0) - label_writer << _location_to_labels[i][_location_to_labels[i].size() - 1]; - - label_writer << std::endl; - } - label_writer.close(); - - // write compacted raw_labels if data hence _location_to_labels was also compacted - if (compact_before_save && _dynamic_index) - { - _label_map = load_label_map(std::string(filename) + "_labels_map.txt"); - std::unordered_map mapped_to_raw_labels; - // invert label map - for (const auto &[key, value] : _label_map) - { - mapped_to_raw_labels.insert({value, key}); - } - - // write updated labels - std::ofstream raw_label_writer(std::string(filename) + "_raw_labels.txt"); - assert(raw_label_writer.is_open()); - for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) - { - for (uint32_t j = 0; j + 1 < _location_to_labels[i].size(); j++) - { - raw_label_writer << mapped_to_raw_labels[_location_to_labels[i][j]] << ","; - } - if (_location_to_labels[i].size() != 0) - raw_label_writer - << mapped_to_raw_labels[_location_to_labels[i][_location_to_labels[i].size() - 1]]; - - raw_label_writer << std::endl; - } - raw_label_writer.close(); - } - } - } - - std::string graph_file = std::string(filename); - std::string tags_file = std::string(filename) + ".tags"; - std::string data_file = std::string(filename) + ".data"; - std::string delete_list_file = std::string(filename) + ".del"; - - // Because the save_* functions use append mode, ensure that - // the files are deleted before save. Ideally, we should check - // the error code for delete_file, but will ignore now because - // delete should succeed if save will succeed. - delete_file(graph_file); - save_graph(graph_file); - delete_file(data_file); - save_data(data_file); - delete_file(tags_file); - save_tags(tags_file); - delete_file(delete_list_file); - save_delete_list(delete_list_file); - } - else - { - diskann::cout << "Save index in a single file currently not supported. " - "Not saving the index." - << std::endl; - } - - // If frozen points were temporarily compacted to _nd, move back to - // _max_points. - reposition_frozen_point_to_end(); - - diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 << "s." << std::endl; -} - -#ifdef EXEC_ENV_OLS -template -size_t Index::load_tags(AlignedFileReader &reader) -{ -#else -template -size_t Index::load_tags(const std::string tag_filename) -{ - if (_enable_tags && !file_exists(tag_filename)) - { - diskann::cerr << "Tag file " << tag_filename << " does not exist!" << std::endl; - throw diskann::ANNException("Tag file " + tag_filename + " does not exist!", -1, __FUNCSIG__, __FILE__, - __LINE__); - } -#endif - if (!_enable_tags) - { - diskann::cout << "Tags not loaded as tags not enabled." << std::endl; - return 0; - } - - size_t file_dim, file_num_points; - TagT *tag_data; -#ifdef EXEC_ENV_OLS - load_bin(reader, tag_data, file_num_points, file_dim); -#else - load_bin(std::string(tag_filename), tag_data, file_num_points, file_dim); -#endif - - if (file_dim != 1) - { - std::stringstream stream; - stream << "ERROR: Found " << file_dim << " dimensions for tags," - << "but tag file must have 1 dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - delete[] tag_data; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - const size_t num_data_points = file_num_points - _num_frozen_pts; - _location_to_tag.reserve(num_data_points); - _tag_to_location.reserve(num_data_points); - for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) - { - TagT tag = *(tag_data + i); - if (_delete_set->find(i) == _delete_set->end()) - { - _location_to_tag.set(i, tag); - _tag_to_location[tag] = i; - } - } - diskann::cout << "Tags loaded." << std::endl; - delete[] tag_data; - return file_num_points; -} - -template -#ifdef EXEC_ENV_OLS -size_t Index::load_data(AlignedFileReader &reader) -{ -#else -size_t Index::load_data(std::string filename) -{ -#endif - size_t file_dim, file_num_points; -#ifdef EXEC_ENV_OLS - diskann::get_bin_metadata(reader, file_num_points, file_dim); -#else - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); -#endif - - // since we are loading a new dataset, _empty_slots must be cleared - _empty_slots.clear(); - - if (file_dim != _dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_num_points > _max_points + _num_frozen_pts) - { - // update and tag lock acquired in load() before calling load_data - resize(file_num_points - _num_frozen_pts); - } - -#ifdef EXEC_ENV_OLS - // REFACTOR TODO: Must figure out how to support aligned reader in a clean - // manner. - copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _data_store->get_aligned_dim()); -#else - _data_store->load(filename); // offset == 0. -#endif - return file_num_points; -} - -#ifdef EXEC_ENV_OLS -template -size_t Index::load_delete_set(AlignedFileReader &reader) -{ -#else -template -size_t Index::load_delete_set(const std::string &filename) -{ -#endif - std::unique_ptr delete_list; - size_t npts, ndim; - -#ifdef EXEC_ENV_OLS - diskann::load_bin(reader, delete_list, npts, ndim); -#else - diskann::load_bin(filename, delete_list, npts, ndim); -#endif - assert(ndim == 1); - for (uint32_t i = 0; i < npts; i++) - { - _delete_set->insert(delete_list[i]); - } - return npts; -} - -// load the index from file and update the max_degree, cur (navigating -// node loc), and _final_graph (adjacency list) -template -#ifdef EXEC_ENV_OLS -void Index::load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) -{ -#else -void Index::load(const char *filename, uint32_t num_threads, uint32_t search_l) -{ -#endif - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - _has_built = true; - - size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, label_num_pts = 0; - - std::string mem_index_file(filename); - std::string labels_file = mem_index_file + "_labels.txt"; - std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt"; - std::string labels_map_file = mem_index_file + "_labels_map.txt"; - - if (!_save_as_one_file) - { - // For DLVS Store, we will not support saving the index in multiple - // files. -#ifndef EXEC_ENV_OLS - std::string data_file = std::string(filename) + ".data"; - std::string tags_file = std::string(filename) + ".tags"; - std::string delete_set_file = std::string(filename) + ".del"; - std::string graph_file = std::string(filename); - data_file_num_pts = load_data(data_file); - if (file_exists(delete_set_file)) - { - load_delete_set(delete_set_file); - } - if (_enable_tags) - { - tags_file_num_pts = load_tags(tags_file); - } - graph_num_pts = load_graph(graph_file, data_file_num_pts); -#endif - } - else - { - diskann::cout << "Single index file saving/loading support not yet " - "enabled. Not loading the index." - << std::endl; - return; - } - - if (data_file_num_pts != graph_num_pts || (data_file_num_pts != tags_file_num_pts && _enable_tags)) - { - std::stringstream stream; - stream << "ERROR: When loading index, loaded " << data_file_num_pts << " points from datafile, " - << graph_num_pts << " from graph, and " << tags_file_num_pts - << " tags, with num_frozen_pts being set to " << _num_frozen_pts << " in constructor." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_exists(labels_file)) - { - _label_map = load_label_map(labels_map_file); - parse_label_file(labels_file, label_num_pts); - assert(label_num_pts == data_file_num_pts - _num_frozen_pts); - if (file_exists(labels_to_medoids)) - { - std::ifstream medoid_stream(labels_to_medoids); - std::string line, token; - uint32_t line_cnt = 0; - - _label_to_start_id.clear(); - - while (std::getline(medoid_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - LabelT label; - while (std::getline(iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = (LabelT)std::stoul(token); - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; - } - _label_to_start_id[label] = medoid; - line_cnt++; - } - } - - std::string universal_label_file(filename); - universal_label_file += "_universal_label.txt"; - if (file_exists(universal_label_file)) - { - std::ifstream universal_label_reader(universal_label_file); - universal_label_reader >> _universal_label; - _use_universal_label = true; - universal_label_reader.close(); - } - } - - _nd = data_file_num_pts - _num_frozen_pts; - _empty_slots.clear(); - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)i); - } - - reposition_frozen_point_to_end(); - diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd << " _start: " << _start - << " size(_location_to_tag): " << _location_to_tag.size() - << " size(_tag_to_location):" << _tag_to_location.size() << " Max points: " << _max_points - << std::endl; - - // For incremental index, _query_scratch is initialized in the constructor. - // For the bulk index, the params required to initialize _query_scratch - // are known only at load time, hence this check and the call to - // initialize_q_s(). - if (_query_scratch.size() == 0) - { - initialize_query_scratch(num_threads, search_l, search_l, (uint32_t)_graph_store->get_max_range_of_graph(), - _indexingMaxC, _dim); - } -} - -#ifndef EXEC_ENV_OLS -template -size_t Index::get_graph_num_frozen_points(const std::string &graph_file) -{ - size_t expected_file_size; - uint32_t max_observed_degree, start; - size_t file_frozen_pts; - - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); - - in.open(graph_file, std::ios::binary); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&max_observed_degree, sizeof(uint32_t)); - in.read((char *)&start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); - - return file_frozen_pts; -} -#endif - -#ifdef EXEC_ENV_OLS -template -size_t Index::load_graph(AlignedFileReader &reader, size_t expected_num_points) -{ -#else - -template -size_t Index::load_graph(std::string filename, size_t expected_num_points) -{ -#endif - auto res = _graph_store->load(filename, expected_num_points); - _start = std::get<1>(res); - _num_frozen_pts = std::get<2>(res); - return std::get<0>(res); -} - -template -int Index::_get_vector_by_tag(TagType &tag, DataType &vec) -{ - try - { - TagT tag_val = std::any_cast(tag); - T *vec_val = std::any_cast(vec); - return this->get_vector_by_tag(tag_val, vec_val); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast while performing _get_vector_by_tags() " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template int Index::get_vector_by_tag(TagT &tag, T *vec) -{ - std::shared_lock lock(_tag_lock); - if (_tag_to_location.find(tag) == _tag_to_location.end()) - { - diskann::cout << "Tag " << get_tag_string(tag) << " does not exist" << std::endl; - return -1; - } - - location_t location = _tag_to_location[tag]; - _data_store->get_vector(location, vec); - - return 0; -} - -template uint32_t Index::calculate_entry_point() -{ - // REFACTOR TODO: This function does not support multi-threaded calculation of medoid. - // Must revisit if perf is a concern. - return _data_store->calculate_medoid(); -} - -template std::vector Index::get_init_ids() -{ - std::vector init_ids; - init_ids.reserve(1 + _num_frozen_pts); - - init_ids.emplace_back(_start); - - for (uint32_t frozen = (uint32_t)_max_points; frozen < _max_points + _num_frozen_pts; frozen++) - { - if (frozen != _start) - { - init_ids.emplace_back(frozen); - } - } - - return init_ids; -} - -// Find common filter between a node's labels and a given set of labels, while -// taking into account universal label -template -bool Index::detect_common_filters(uint32_t point_id, bool search_invocation, - const std::vector &incoming_labels) -{ - auto &curr_node_labels = _location_to_labels[point_id]; - std::vector common_filters; - std::set_intersection(incoming_labels.begin(), incoming_labels.end(), curr_node_labels.begin(), - curr_node_labels.end(), std::back_inserter(common_filters)); - if (common_filters.size() > 0) - { - // This is to reduce the repetitive calls. If common_filters size is > 0 , - // we dont need to check further for universal label - return true; - } - if (_use_universal_label) - { - if (!search_invocation) - { - if (std::find(incoming_labels.begin(), incoming_labels.end(), _universal_label) != incoming_labels.end() || - std::find(curr_node_labels.begin(), curr_node_labels.end(), _universal_label) != curr_node_labels.end()) - common_filters.push_back(_universal_label); - } - else - { - if (std::find(curr_node_labels.begin(), curr_node_labels.end(), _universal_label) != curr_node_labels.end()) - common_filters.push_back(_universal_label); - } - } - return (common_filters.size() > 0); -} - -template -std::pair Index::iterate_to_fixed_point( - InMemQueryScratch *scratch, const uint32_t Lsize, const std::vector &init_ids, bool use_filter, - const std::vector &filter_labels, bool search_invocation) -{ - std::vector &expanded_nodes = scratch->pool(); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - best_L_nodes.reserve(Lsize); - tsl::robin_set &inserted_into_pool_rs = scratch->inserted_into_pool_rs(); - boost::dynamic_bitset<> &inserted_into_pool_bs = scratch->inserted_into_pool_bs(); - std::vector &id_scratch = scratch->id_scratch(); - std::vector &dist_scratch = scratch->dist_scratch(); - assert(id_scratch.size() == 0); - - T *aligned_query = scratch->aligned_query(); - - float *pq_dists = nullptr; - - _pq_data_store->preprocess_query(aligned_query, scratch); - - if (expanded_nodes.size() > 0 || id_scratch.size() > 0) - { - throw ANNException("ERROR: Clear scratch space before passing.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // Decide whether to use bitset or robin set to mark visited nodes - auto total_num_points = _max_points + _num_frozen_pts; - bool fast_iterate = total_num_points <= MAX_POINTS_FOR_USING_BITSET; - - if (fast_iterate) - { - if (inserted_into_pool_bs.size() < total_num_points) - { - // hopefully using 2X will reduce the number of allocations. - auto resize_size = - 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET ? MAX_POINTS_FOR_USING_BITSET : 2 * total_num_points; - inserted_into_pool_bs.resize(resize_size); - } - } - - // Lambda to determine if a node has been visited - auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, &inserted_into_pool_rs](const uint32_t id) { - return fast_iterate ? inserted_into_pool_bs[id] == 0 - : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end(); - }; - - // Lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, scratch, pq_dists](const std::vector &ids, std::vector &dists_out) { - _pq_data_store->get_distance(scratch->aligned_query(), ids, dists_out, scratch); - }; - - // Initialize the candidate pool with starting points - for (auto id : init_ids) - { - if (id >= _max_points + _num_frozen_pts) - { - diskann::cerr << "Out of range loc found as an edge : " << id << std::endl; - throw diskann::ANNException(std::string("Wrong loc") + std::to_string(id), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (use_filter) - { - if (!detect_common_filters(id, search_invocation, filter_labels)) - continue; - } - - if (is_not_visited(id)) - { - if (fast_iterate) - { - inserted_into_pool_bs[id] = 1; - } - else - { - inserted_into_pool_rs.insert(id); - } - - float distance; - uint32_t ids[] = {id}; - float distances[] = {std::numeric_limits::max()}; - _pq_data_store->get_distance(aligned_query, ids, 1, distances, scratch); - distance = distances[0]; - - Neighbor nn = Neighbor(id, distance); - best_L_nodes.insert(nn); - } - } - - uint32_t hops = 0; - uint32_t cmps = 0; - - while (best_L_nodes.has_unexpanded_node()) - { - auto nbr = best_L_nodes.closest_unexpanded(); - auto n = nbr.id; - - // Add node to expanded nodes to create pool for prune later - if (!search_invocation) - { - if (!use_filter) - { - expanded_nodes.emplace_back(nbr); - } - else - { // in filter based indexing, the same point might invoke - // multiple iterate_to_fixed_points, so need to be careful - // not to add the same item to pool multiple times. - if (std::find(expanded_nodes.begin(), expanded_nodes.end(), nbr) == expanded_nodes.end()) - { - expanded_nodes.emplace_back(nbr); - } - } - } - - // Find which of the nodes in des have not been visited before - id_scratch.clear(); - dist_scratch.clear(); - if (_dynamic_index) - { - LockGuard guard(_locks[n]); - for (auto id : _graph_store->get_neighbours(n)) - { - assert(id < _max_points + _num_frozen_pts); - - if (use_filter) - { - // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. - if (!detect_common_filters(id, search_invocation, filter_labels)) - continue; - } - - if (is_not_visited(id)) - { - id_scratch.push_back(id); - } - } - } - else - { - _locks[n].lock(); - auto nbrs = _graph_store->get_neighbours(n); - _locks[n].unlock(); - for (auto id : nbrs) - { - assert(id < _max_points + _num_frozen_pts); - - if (use_filter) - { - // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. - if (!detect_common_filters(id, search_invocation, filter_labels)) - continue; - } - - if (is_not_visited(id)) - { - id_scratch.push_back(id); - } - } - } - - // Mark nodes visited - for (auto id : id_scratch) - { - if (fast_iterate) - { - inserted_into_pool_bs[id] = 1; - } - else - { - inserted_into_pool_rs.insert(id); - } - } - - assert(dist_scratch.capacity() >= id_scratch.size()); - compute_dists(id_scratch, dist_scratch); - cmps += (uint32_t)id_scratch.size(); - - // Insert pairs into the pool of candidates - for (size_t m = 0; m < id_scratch.size(); ++m) - { - best_L_nodes.insert(Neighbor(id_scratch[m], dist_scratch[m])); - } - } - return std::make_pair(hops, cmps); -} - -template -void Index::search_for_point_and_prune(int location, uint32_t Lindex, - std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter, - uint32_t filteredLindex) -{ - const std::vector init_ids = get_init_ids(); - const std::vector unused_filter_label; - - if (!use_filter) - { - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch, Lindex, init_ids, false, unused_filter_label, false); - } - else - { - std::shared_lock tl(_tag_lock, std::defer_lock); - if (_dynamic_index) - tl.lock(); - std::vector filter_specific_start_nodes; - for (auto &x : _location_to_labels[location]) - filter_specific_start_nodes.emplace_back(_label_to_start_id[x]); - - if (_dynamic_index) - tl.unlock(); - - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch, filteredLindex, filter_specific_start_nodes, true, - _location_to_labels[location], false); - - // combine candidate pools obtained with filter and unfiltered criteria. - std::set best_candidate_pool; - for (auto filtered_neighbor : scratch->pool()) - { - best_candidate_pool.insert(filtered_neighbor); - } - - // clear scratch for finding unfiltered candidates - scratch->clear(); - - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch, Lindex, init_ids, false, unused_filter_label, false); - - for (auto unfiltered_neighbour : scratch->pool()) - { - // insert if this neighbour is not already in best_candidate_pool - if (best_candidate_pool.find(unfiltered_neighbour) == best_candidate_pool.end()) - { - best_candidate_pool.insert(unfiltered_neighbour); - } - } - - scratch->pool().clear(); - std::copy(best_candidate_pool.begin(), best_candidate_pool.end(), std::back_inserter(scratch->pool())); - } - - auto &pool = scratch->pool(); - - for (uint32_t i = 0; i < pool.size(); i++) - { - if (pool[i].id == (uint32_t)location) - { - pool.erase(pool.begin() + i); - i--; - } - } - - if (pruned_list.size() > 0) - { - throw diskann::ANNException("ERROR: non-empty pruned_list passed", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - prune_neighbors(location, pool, pruned_list, scratch); - - assert(!pruned_list.empty()); - assert(_graph_store->get_total_points() == _max_points + _num_frozen_pts); -} - -template -void Index::occlude_list(const uint32_t location, std::vector &pool, const float alpha, - const uint32_t degree, const uint32_t maxc, std::vector &result, - InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr) -{ - if (pool.size() == 0) - return; - - // Truncate pool at maxc and initialize scratch spaces - assert(std::is_sorted(pool.begin(), pool.end())); - assert(result.size() == 0); - if (pool.size() > maxc) - pool.resize(maxc); - std::vector &occlude_factor = scratch->occlude_factor(); - // occlude_list can be called with the same scratch more than once by - // search_for_point_and_add_link through inter_insert. - occlude_factor.clear(); - // Initialize occlude_factor to pool.size() many 0.0f values for correctness - occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); - - float cur_alpha = 1; - while (cur_alpha <= alpha && result.size() < degree) - { - // used for MIPS, where we store a value of eps in cur_alpha to - // denote pruned out entries which we can skip in later rounds. - float eps = cur_alpha + 0.01f; - - for (auto iter = pool.begin(); result.size() < degree && iter != pool.end(); ++iter) - { - if (occlude_factor[iter - pool.begin()] > cur_alpha) - { - continue; - } - // Set the entry to float::max so that is not considered again - occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); - // Add the entry to the result if its not been deleted, and doesn't - // add a self loop - if (delete_set_ptr == nullptr || delete_set_ptr->find(iter->id) == delete_set_ptr->end()) - { - if (iter->id != location) - { - result.push_back(iter->id); - } - } - - // Update occlude factor for points from iter+1 to pool.end() - for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) - { - auto t = iter2 - pool.begin(); - if (occlude_factor[t] > alpha) - continue; - - bool prune_allowed = true; - if (_filtered_index) - { - uint32_t a = iter->id; - uint32_t b = iter2->id; - if (_location_to_labels.size() < b || _location_to_labels.size() < a) - continue; - for (auto &x : _location_to_labels[b]) - { - if (std::find(_location_to_labels[a].begin(), _location_to_labels[a].end(), x) == - _location_to_labels[a].end()) - { - prune_allowed = false; - } - if (!prune_allowed) - break; - } - } - if (!prune_allowed) - continue; - - float djk = _data_store->get_distance(iter2->id, iter->id); - if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) - { - occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() - : std::max(occlude_factor[t], iter2->distance / djk); - } - else if (_dist_metric == diskann::Metric::INNER_PRODUCT) - { - // Improvization for flipping max and min dist for MIPS - float x = -iter2->distance; - float y = -djk; - if (y > cur_alpha * x) - { - occlude_factor[t] = std::max(occlude_factor[t], eps); - } - } - } - } - cur_alpha *= 1.2f; - } -} - -template -void Index::prune_neighbors(const uint32_t location, std::vector &pool, - std::vector &pruned_list, InMemQueryScratch *scratch) -{ - prune_neighbors(location, pool, _indexingRange, _indexingMaxC, _indexingAlpha, pruned_list, scratch); -} - -template -void Index::prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, - const uint32_t max_candidate_size, const float alpha, - std::vector &pruned_list, InMemQueryScratch *scratch) -{ - if (pool.size() == 0) - { - // if the pool is empty, behave like a noop - pruned_list.clear(); - return; - } - - // If using _pq_build, over-write the PQ distances with actual distances - // REFACTOR PQ: TODO: How to get rid of this!? - if (_pq_dist) - { - for (auto &ngh : pool) - ngh.distance = _data_store->get_distance(ngh.id, location); - } - - // sort the pool based on distance to query and prune it with occlude_list - std::sort(pool.begin(), pool.end()); - pruned_list.clear(); - pruned_list.reserve(range); - - occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, scratch); - assert(pruned_list.size() <= range); - - if (_saturate_graph && alpha > 1) - { - for (const auto &node : pool) - { - if (pruned_list.size() >= range) - break; - if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == pruned_list.end()) && - node.id != location) - pruned_list.push_back(node.id); - } - } -} - -template -void Index::inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, - InMemQueryScratch *scratch) -{ - const auto &src_pool = pruned_list; - - assert(!src_pool.empty()); - - for (auto des : src_pool) - { - // des.loc is the loc of the neighbors of n - assert(des < _max_points + _num_frozen_pts); - // des_pool contains the neighbors of the neighbors of n - std::vector copy_of_neighbors; - bool prune_needed = false; - { - LockGuard guard(_locks[des]); - auto &des_pool = _graph_store->get_neighbours(des); - if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) - { - if (des_pool.size() < (uint64_t)(defaults::GRAPH_SLACK_FACTOR * range)) - { - // des_pool.emplace_back(n); - _graph_store->add_neighbour(des, n); - prune_needed = false; - } - else - { - copy_of_neighbors.reserve(des_pool.size() + 1); - copy_of_neighbors = des_pool; - copy_of_neighbors.push_back(n); - prune_needed = true; - } - } - } // des lock is released by this point - - if (prune_needed) - { - tsl::robin_set dummy_visited(0); - std::vector dummy_pool(0); - - size_t reserveSize = (size_t)(std::ceil(1.05 * defaults::GRAPH_SLACK_FACTOR * range)); - dummy_visited.reserve(reserveSize); - dummy_pool.reserve(reserveSize); - - for (auto cur_nbr : copy_of_neighbors) - { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) - { - float dist = _data_store->get_distance(des, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - std::vector new_out_neighbors; - prune_neighbors(des, dummy_pool, new_out_neighbors, scratch); - { - LockGuard guard(_locks[des]); - - _graph_store->set_neighbours(des, new_out_neighbors); - } - } - } -} - -template -void Index::inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch) -{ - inter_insert(n, pruned_list, _indexingRange, scratch); -} - -template void Index::link() -{ - uint32_t num_threads = _indexingThreads; - if (num_threads != 0) - omp_set_num_threads(num_threads); - - /* visit_order is a vector that is initialized to the entire graph */ - std::vector visit_order; - std::vector pool, tmp; - tsl::robin_set visited; - visit_order.reserve(_nd + _num_frozen_pts); - for (uint32_t i = 0; i < (uint32_t)_nd; i++) - { - visit_order.emplace_back(i); - } - - // If there are any frozen points, add them all. - for (uint32_t frozen = (uint32_t)_max_points; frozen < _max_points + _num_frozen_pts; frozen++) - { - visit_order.emplace_back(frozen); - } - - // if there are frozen points, the first such one is set to be the _start - if (_num_frozen_pts > 0) - _start = (uint32_t)_max_points; - else - _start = calculate_entry_point(); - - diskann::Timer link_timer; - -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) - { - auto node = visit_order[node_ctr]; - - // Find and add appropriate graph edges - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - std::vector pruned_list; - if (_filtered_index) - { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, true, _filterIndexingQueueSize); - } - else - { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch); - } - assert(pruned_list.size() > 0); - - { - LockGuard guard(_locks[node]); - - _graph_store->set_neighbours(node, pruned_list); - assert(_graph_store->get_neighbours((location_t)node).size() <= _indexingRange); - } - - inter_insert(node, pruned_list, scratch); - - if (node_ctr % 100000 == 0) - { - diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) << "% of index build completed." - << std::flush; - } - } - - if (_nd > 0) - { - diskann::cout << "Starting final cleanup.." << std::flush; - } -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) - { - auto node = visit_order[node_ctr]; - if (_graph_store->get_neighbours((location_t)node).size() > _indexingRange) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - tsl::robin_set dummy_visited(0); - std::vector dummy_pool(0); - std::vector new_out_neighbors; - - for (auto cur_nbr : _graph_store->get_neighbours((location_t)node)) - { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) - { - float dist = _data_store->get_distance(node, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - prune_neighbors(node, dummy_pool, new_out_neighbors, scratch); - - _graph_store->clear_neighbours((location_t)node); - _graph_store->set_neighbours((location_t)node, new_out_neighbors); - } - } - if (_nd > 0) - { - diskann::cout << "done. Link time: " << ((double)link_timer.elapsed() / (double)1000000) << "s" << std::endl; - } -} - -template -void Index::prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion_size, - const float alpha) -{ - const uint32_t range = max_degree; - const uint32_t maxc = max_occlusion_size; - - _filtered_index = true; - - diskann::Timer timer; -#pragma omp parallel for - for (int64_t node = 0; node < (int64_t)(_max_points + _num_frozen_pts); node++) - { - if ((size_t)node < _nd || (size_t)node >= _max_points) - { - if (_graph_store->get_neighbours((location_t)node).size() > range) - { - tsl::robin_set dummy_visited(0); - std::vector dummy_pool(0); - std::vector new_out_neighbors; - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - for (auto cur_nbr : _graph_store->get_neighbours((location_t)node)) - { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) - { - float dist = _data_store->get_distance((location_t)node, (location_t)cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - - prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch); - _graph_store->clear_neighbours((location_t)node); - _graph_store->set_neighbours((location_t)node, new_out_neighbors); - } - } - } - - diskann::cout << "Prune time : " << timer.elapsed() / 1000 << "ms" << std::endl; - size_t max = 0, min = 1 << 30, total = 0, cnt = 0; - for (size_t i = 0; i < _max_points + _num_frozen_pts; i++) - { - if (i < _nd || i >= _max_points) - { - const std::vector &pool = _graph_store->get_neighbours((location_t)i); - max = (std::max)(max, pool.size()); - min = (std::min)(min, pool.size()); - total += pool.size(); - if (pool.size() < 2) - cnt++; - } - } - if (min > max) - min = max; - if (_nd > 0) - { - diskann::cout << "Index built with degree: max:" << max - << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) << " min:" << min - << " count(deg<2):" << cnt << std::endl; - } -} - -// REFACTOR -template -void Index::set_start_points(const T *data, size_t data_count) -{ - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (data_count != _num_frozen_pts * _dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); - - // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * - // sizeof(T) * _num_frozen_pts); - for (location_t i = 0; i < _num_frozen_pts; i++) - { - _data_store->set_vector((location_t)(i + _max_points), data + i * _dim); - } - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -} - -template -void Index::_set_start_points_at_random(DataType radius, uint32_t random_seed) -{ - try - { - T radius_to_use = std::any_cast(radius); - this->set_start_points_at_random(radius_to_use, random_seed); - } - catch (const std::bad_any_cast &e) - { - throw ANNException( - "Error: bad any cast while performing _set_start_points_at_random() " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) -{ - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; - - std::vector points_data; - points_data.reserve(_dim * _num_frozen_pts); - std::vector real_vec(_dim); - - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) - { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) - { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; - } - - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } - - set_start_points(points_data.data(), points_data.size()); -} - -template -void Index::build_with_data_populated(const std::vector &tags) -{ - diskann::cout << "Starting index build with " << _nd << " points... " << std::endl; - - if (_nd < 1) - throw ANNException("Error: Trying to build an index with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (_enable_tags && tags.size() != _nd) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _nd << " points from file," - << "but tags vector is of size " << tags.size() << "." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - if (_enable_tags) - { - for (size_t i = 0; i < tags.size(); ++i) - { - _tag_to_location[tags[i]] = (uint32_t)i; - _location_to_tag.set(static_cast(i), tags[i]); - } - } - - uint32_t index_R = _indexingRange; - uint32_t num_threads_index = _indexingThreads; - uint32_t index_L = _indexingQueueSize; - uint32_t maxc = _indexingMaxC; - - if (_query_scratch.size() == 0) - { - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, - _data_store->get_aligned_dim()); - } - - generate_frozen_point(); - link(); - - size_t max = 0, min = SIZE_MAX, total = 0, cnt = 0; - for (size_t i = 0; i < _nd; i++) - { - auto &pool = _graph_store->get_neighbours((location_t)i); - max = std::max(max, pool.size()); - min = std::min(min, pool.size()); - total += pool.size(); - if (pool.size() < 2) - cnt++; - } - diskann::cout << "Index built with degree: max:" << max << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) - << " min:" << min << " count(deg<2):" << cnt << std::endl; - - _has_built = true; -} -template -void Index::_build(const DataType &data, const size_t num_points_to_load, TagVector &tags) -{ - try - { - this->build(std::any_cast(data), num_points_to_load, tags.get>()); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast in while building index. " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error" + std::string(e.what()), -1); - } -} -template -void Index::build(const T *data, const size_t num_points_to_load, const std::vector &tags) -{ - if (num_points_to_load == 0) - { - throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - } - if (_pq_dist) - { - throw ANNException("ERROR: DO not use this build interface with PQ distance", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - std::unique_lock ul(_update_lock); - - { - std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; - - _data_store->populate_data(data, (location_t)num_points_to_load); - } - - build_with_data_populated(tags); -} - -template -void Index::build(const char *filename, const size_t num_points_to_load, const std::vector &tags) -{ - // idealy this should call build_filtered_index based on params passed - - std::unique_lock ul(_update_lock); - - // error checks - if (num_points_to_load == 0) - throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: Data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - size_t file_num_points, file_dim; - if (filename == nullptr) - { - throw diskann::ANNException("Can not build with an empty file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::get_bin_metadata(filename, file_num_points, file_dim); - if (file_num_points > _max_points) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has " << file_num_points - << " points, but " - << "index can support only " << _max_points << " points as specified in constructor." << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (num_points_to_load > file_num_points) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has only " - << file_num_points << " points." << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_dim != _dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // REFACTOR PQ TODO: We can remove this if and add a check in the InMemDataStore - // to not populate_data if it has been called once. - if (_pq_dist) - { -#ifdef EXEC_ENV_OLS - std::stringstream ss; - ss << "PQ Build is not supported in DLVS environment (i.e. if EXEC_ENV_OLS is defined)" << std::endl; - diskann::cerr << ss.str() << std::endl; - throw ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); -#else - // REFACTOR TODO: Both in the previous code and in the current PQDataStore, - // we are writing the PQ files in the same path as the input file. Now we - // may not have write permissions to that folder, but we will always have - // write permissions to the output folder. So we should write the PQ files - // there. The problem is that the Index class gets the output folder prefix - // only at the time of save(), by which time we are too late. So leaving it - // as-is for now. - _pq_data_store->populate_data(filename, 0U); -#endif - } - - _data_store->populate_data(filename, 0U); - diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; - - { - std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; - } - build_with_data_populated(tags); -} - -template -void Index::build(const char *filename, const size_t num_points_to_load, const char *tag_filename) -{ - std::vector tags; - - if (_enable_tags) - { - std::unique_lock tl(_tag_lock); - if (tag_filename == nullptr) - { - throw ANNException("Tag filename is null, while _enable_tags is set", -1, __FUNCSIG__, __FILE__, __LINE__); - } - else - { - if (file_exists(tag_filename)) - { - diskann::cout << "Loading tags from " << tag_filename << " for vamana index build" << std::endl; - TagT *tag_data = nullptr; - size_t npts, ndim; - diskann::load_bin(tag_filename, tag_data, npts, ndim); - if (npts < num_points_to_load) - { - std::stringstream sstream; - sstream << "Loaded " << npts << " tags, insufficient to populate tags for " << num_points_to_load - << " points to load"; - throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - for (size_t i = 0; i < num_points_to_load; i++) - { - tags.push_back(tag_data[i]); - } - delete[] tag_data; - } - else - { - throw diskann::ANNException(std::string("Tag file") + tag_filename + " does not exist", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - } - } - build(filename, num_points_to_load, tags); -} - -template -void Index::build(const std::string &data_file, const size_t num_points_to_load, - IndexFilterParams &filter_params) -{ - size_t points_to_load = num_points_to_load == 0 ? _max_points : num_points_to_load; - - auto s = std::chrono::high_resolution_clock::now(); - if (filter_params.label_file == "") - { - this->build(data_file.c_str(), points_to_load); - } - else - { - // TODO: this should ideally happen in save() - std::string labels_file_to_use = filter_params.save_path_prefix + "_label_formatted.txt"; - std::string mem_labels_int_map_file = filter_params.save_path_prefix + "_labels_map.txt"; - convert_labels_string_to_int(filter_params.label_file, labels_file_to_use, mem_labels_int_map_file, - filter_params.universal_label); - if (filter_params.universal_label != "") - { - LabelT unv_label_as_num = 0; - this->set_universal_label(unv_label_as_num); - } - this->build_filtered_index(data_file.c_str(), labels_file_to_use, points_to_load); - } - std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; - std::cout << "Indexing time: " << diff.count() << "\n"; -} - -template -std::unordered_map Index::load_label_map(const std::string &labels_map_file) -{ - std::unordered_map string_to_int_mp; - std::ifstream map_reader(labels_map_file); - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) - { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = (LabelT)std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; -} - -template -LabelT Index::get_converted_label(const std::string &raw_label) -{ - if (_label_map.find(raw_label) != _label_map.end()) - { - return _label_map[raw_label]; - } - if (_use_universal_label) - { - return _universal_label; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); -} - -template -void Index::parse_label_file(const std::string &label_file, size_t &num_points) -{ - // Format of Label txt file: filters with comma separators - - std::ifstream infile(label_file); - if (infile.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); - } - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(infile, line)) - { - line_cnt++; - } - _location_to_labels.resize(line_cnt, std::vector()); - - infile.clear(); - infile.seekg(0, std::ios::beg); - line_cnt = 0; - - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = (LabelT)std::stoul(token); - lbls.push_back(token_as_num); - _labels.insert(token_as_num); - } - - std::sort(lbls.begin(), lbls.end()); - _location_to_labels[line_cnt] = lbls; - line_cnt++; - } - num_points = (size_t)line_cnt; - diskann::cout << "Identified " << _labels.size() << " distinct label(s)" << std::endl; -} - -template -void Index::_set_universal_label(const LabelType universal_label) -{ - this->set_universal_label(std::any_cast(universal_label)); -} - -template -void Index::set_universal_label(const LabelT &label) -{ - _use_universal_label = true; - _universal_label = label; -} - -template -void Index::build_filtered_index(const char *filename, const std::string &label_file, - const size_t num_points_to_load, const std::vector &tags) -{ - _filtered_index = true; - _label_to_start_id.clear(); - size_t num_points_labels = 0; - - parse_label_file(label_file, - num_points_labels); // determines medoid for each label and identifies - // the points to label mapping - - std::unordered_map> label_to_points; - - for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) - { - for (auto label : _location_to_labels[point_id]) - { - if (label != _universal_label) - { - label_to_points[label].emplace_back(point_id); - } - else - { - for (typename tsl::robin_set::size_type lbl = 0; lbl < _labels.size(); lbl++) - { - auto itr = _labels.begin(); - std::advance(itr, lbl); - auto &x = *itr; - label_to_points[x].emplace_back(point_id); - } - } - } - } - - uint32_t num_cands = 25; - for (auto itr = _labels.begin(); itr != _labels.end(); itr++) - { - uint32_t best_medoid_count = std::numeric_limits::max(); - auto &curr_label = *itr; - uint32_t best_medoid; - auto labeled_points = label_to_points[curr_label]; - for (uint32_t cnd = 0; cnd < num_cands; cnd++) - { - uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; - uint32_t cur_cnt = std::numeric_limits::max(); - if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) - { - _medoid_counts[cur_cnd] = 0; - cur_cnt = 0; - } - else - { - cur_cnt = _medoid_counts[cur_cnd]; - } - if (cur_cnt < best_medoid_count) - { - best_medoid_count = cur_cnt; - best_medoid = cur_cnd; - } - } - _label_to_start_id[curr_label] = best_medoid; - _medoid_counts[best_medoid]++; - } - - this->build(filename, num_points_to_load, tags); -} - -template -std::pair Index::_search(const DataType &query, const size_t K, const uint32_t L, - std::any &indices, float *distances) -{ - try - { - auto typed_query = std::any_cast(query); - if (typeid(uint32_t *) == indices.type()) - { - auto u32_ptr = std::any_cast(indices); - return this->search(typed_query, K, L, u32_ptr, distances); - } - else if (typeid(uint64_t *) == indices.type()) - { - auto u64_ptr = std::any_cast(indices); - return this->search(typed_query, K, L, u64_ptr, distances); - } - else - { - throw ANNException("Error: indices type can only be uint64_t or uint32_t.", -1); - } - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast while searching. " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template -template -std::pair Index::search(const T *query, const size_t K, const uint32_t L, - IdType *indices, float *distances) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } - - const std::vector unused_filter_label; - const std::vector init_ids = get_init_ids(); - - std::shared_lock lock(_update_lock); - - _data_store->preprocess_query(query, scratch); - - auto retval = iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); - - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - if (best_L_nodes[i].id < _max_points) - { - // safe because Index uses uint32_t ids internally - // and IDType will be uint32_t or uint64_t - indices[pos] = (IdType)best_L_nodes[i].id; - if (distances != nullptr) - { -#ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; -#else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; -#endif - } - pos++; - } - if (pos == K) - break; - } - if (pos < K) - { - diskann::cerr << "Found pos: " << pos << "fewer than K elements " << K << " for query" << std::endl; - } - - return retval; -} - -template -std::pair Index::_search_with_filters(const DataType &query, - const std::string &raw_label, const size_t K, - const uint32_t L, std::any &indices, - float *distances) -{ - auto converted_label = this->get_converted_label(raw_label); - if (typeid(uint64_t *) == indices.type()) - { - auto ptr = std::any_cast(indices); - return this->search_with_filters(std::any_cast(query), converted_label, K, L, ptr, distances); - } - else if (typeid(uint32_t *) == indices.type()) - { - auto ptr = std::any_cast(indices); - return this->search_with_filters(std::any_cast(query), converted_label, K, L, ptr, distances); - } - else - { - throw ANNException("Error: Id type can only be uint64_t or uint32_t.", -1); - } -} - -template -template -std::pair Index::search_with_filters(const T *query, const LabelT &filter_label, - const size_t K, const uint32_t L, - IdType *indices, float *distances) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } - - std::vector filter_vec; - std::vector init_ids = get_init_ids(); - - std::shared_lock lock(_update_lock); - std::shared_lock tl(_tag_lock, std::defer_lock); - if (_dynamic_index) - tl.lock(); - - if (_label_to_start_id.find(filter_label) != _label_to_start_id.end()) - { - init_ids.emplace_back(_label_to_start_id[filter_label]); - } - else - { - diskann::cout << "No filtered medoid found. exitting " - << std::endl; // RKNOTE: If universal label found start there - throw diskann::ANNException("No filtered medoid found. exitting ", -1); - } - if (_dynamic_index) - tl.unlock(); - - filter_vec.emplace_back(filter_label); - - _data_store->preprocess_query(query, scratch); - auto retval = iterate_to_fixed_point(scratch, L, init_ids, true, filter_vec, true); - - auto best_L_nodes = scratch->best_l_nodes(); - - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - if (best_L_nodes[i].id < _max_points) - { - indices[pos] = (IdType)best_L_nodes[i].id; - - if (distances != nullptr) - { -#ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; -#else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; -#endif - } - pos++; - } - if (pos == K) - break; - } - if (pos < K) - { - diskann::cerr << "Found fewer than K elements for query" << std::endl; - } - - return retval; -} - -template -size_t Index::_search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, - const TagType &tags, float *distances, DataVector &res_vectors, - bool use_filters, const std::string filter_label) -{ - try - { - return this->search_with_tags(std::any_cast(query), K, L, std::any_cast(tags), distances, - res_vectors.get>(), use_filters, filter_label); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast while performing _search_with_tags() " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template -size_t Index::search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, - float *distances, std::vector &res_vectors, bool use_filters, - const std::string filter_label) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } - - std::shared_lock ul(_update_lock); - - const std::vector init_ids = get_init_ids(); - - //_distance->preprocess_query(query, _data_store->get_dims(), - // scratch->aligned_query()); - _data_store->preprocess_query(query, scratch); - if (!use_filters) - { - const std::vector unused_filter_label; - iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); - } - else - { - std::vector filter_vec; - auto converted_label = this->get_converted_label(filter_label); - filter_vec.push_back(converted_label); - iterate_to_fixed_point(scratch, L, init_ids, true, filter_vec, true); - } - - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - assert(best_L_nodes.size() <= L); - - std::shared_lock tl(_tag_lock); - - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - auto node = best_L_nodes[i]; - - TagT tag; - if (_location_to_tag.try_get(node.id, tag)) - { - tags[pos] = tag; - - if (res_vectors.size() > 0) - { - _data_store->get_vector(node.id, res_vectors[pos]); - } - - if (distances != nullptr) - { -#ifdef EXEC_ENV_OLS - distances[pos] = node.distance; // DLVS expects negative distances -#else - distances[pos] = _dist_metric == INNER_PRODUCT ? -1 * node.distance : node.distance; -#endif - } - pos++; - // If res_vectors.size() < k, clip at the value. - if (pos == K || pos == res_vectors.size()) - break; - } - } - - return pos; -} - -template size_t Index::get_num_points() -{ - std::shared_lock tl(_tag_lock); - return _nd; -} - -template size_t Index::get_max_points() -{ - std::shared_lock tl(_tag_lock); - return _max_points; -} - -template void Index::generate_frozen_point() -{ - if (_num_frozen_pts == 0) - return; - - if (_num_frozen_pts > 1) - { - throw ANNException("More than one frozen point not supported in generate_frozen_point", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - if (_nd == 0) - { - throw ANNException("ERROR: Can not pick a frozen point since nd=0", -1, __FUNCSIG__, __FILE__, __LINE__); - } - size_t res = calculate_entry_point(); - - // REFACTOR PQ: Not sure if we should do this for both stores. - if (_pq_dist) - { - // copy the PQ data corresponding to the point returned by - // calculate_entry_point - // memcpy(_pq_data + _max_points * _num_pq_chunks, - // _pq_data + res * _num_pq_chunks, - // _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); - _pq_data_store->copy_vectors((location_t)res, (location_t)_max_points, 1); - } - else - { - _data_store->copy_vectors((location_t)res, (location_t)_max_points, 1); - } - _frozen_pts_used++; -} - -template int Index::enable_delete() -{ - assert(_enable_tags); - - if (!_enable_tags) - { - diskann::cerr << "Tags must be instantiated for deletions" << std::endl; - return -2; - } - - if (this->_deletes_enabled) - { - return 0; - } - - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - if (_data_compacted) - { - for (uint32_t slot = (uint32_t)_nd; slot < _max_points; ++slot) - { - _empty_slots.insert(slot); - } - } - this->_deletes_enabled = true; - return 0; -} - -template -inline void Index::process_delete(const tsl::robin_set &old_delete_set, size_t loc, - const uint32_t range, const uint32_t maxc, const float alpha, - InMemQueryScratch *scratch) -{ - tsl::robin_set &expanded_nodes_set = scratch->expanded_nodes_set(); - std::vector &expanded_nghrs_vec = scratch->expanded_nodes_vec(); - - // If this condition were not true, deadlock could result - assert(old_delete_set.find((uint32_t)loc) == old_delete_set.end()); - - std::vector adj_list; - { - // Acquire and release lock[loc] before acquiring locks for neighbors - std::unique_lock adj_list_lock; - if (_conc_consolidate) - adj_list_lock = std::unique_lock(_locks[loc]); - adj_list = _graph_store->get_neighbours((location_t)loc); - } - - bool modify = false; - for (auto ngh : adj_list) - { - if (old_delete_set.find(ngh) == old_delete_set.end()) - { - expanded_nodes_set.insert(ngh); - } - else - { - modify = true; - - std::unique_lock ngh_lock; - if (_conc_consolidate) - ngh_lock = std::unique_lock(_locks[ngh]); - for (auto j : _graph_store->get_neighbours((location_t)ngh)) - if (j != loc && old_delete_set.find(j) == old_delete_set.end()) - expanded_nodes_set.insert(j); - } - } - - if (modify) - { - if (expanded_nodes_set.size() <= range) - { - std::unique_lock adj_list_lock(_locks[loc]); - _graph_store->clear_neighbours((location_t)loc); - for (auto &ngh : expanded_nodes_set) - _graph_store->add_neighbour((location_t)loc, ngh); - } - else - { - // Create a pool of Neighbor candidates from the expanded_nodes_set - expanded_nghrs_vec.reserve(expanded_nodes_set.size()); - for (auto &ngh : expanded_nodes_set) - { - expanded_nghrs_vec.emplace_back(ngh, _data_store->get_distance((location_t)loc, (location_t)ngh)); - } - std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); - std::vector &occlude_list_output = scratch->occlude_list_output(); - occlude_list((uint32_t)loc, expanded_nghrs_vec, alpha, range, maxc, occlude_list_output, scratch, - &old_delete_set); - std::unique_lock adj_list_lock(_locks[loc]); - _graph_store->set_neighbours((location_t)loc, occlude_list_output); - } - } -} - -// Returns number of live points left after consolidation -template -consolidation_report Index::consolidate_deletes(const IndexWriteParameters ¶ms) -{ - if (!_enable_tags) - throw diskann::ANNException("Point tag array not instantiated", -1, __FUNCSIG__, __FILE__, __LINE__); - - { - std::shared_lock ul(_update_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - if (_empty_slots.size() + _nd != _max_points) - { - std::string err = "#empty slots + nd != max points"; - diskann::cerr << err << std::endl; - throw ANNException(err, -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_location_to_tag.size() + _delete_set->size() != _nd) - { - diskann::cerr << "Error: _location_to_tag.size (" << _location_to_tag.size() << ") + _delete_set->size (" - << _delete_set->size() << ") != _nd(" << _nd << ") "; - return consolidation_report(diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR, 0, 0, 0, - 0, 0, 0, 0); - } - - if (_location_to_tag.size() != _tag_to_location.size()) - { - throw diskann::ANNException("_location_to_tag and _tag_to_location not of same size", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - } - - std::unique_lock update_lock(_update_lock, std::defer_lock); - if (!_conc_consolidate) - update_lock.lock(); - - std::unique_lock cl(_consolidate_lock, std::defer_lock); - if (!cl.try_lock()) - { - diskann::cerr << "Consildate delete function failed to acquire consolidate lock" << std::endl; - return consolidation_report(diskann::consolidation_report::status_code::LOCK_FAIL, 0, 0, 0, 0, 0, 0, 0); - } - - diskann::cout << "Starting consolidate_deletes... "; - - std::unique_ptr> old_delete_set(new tsl::robin_set); - { - std::unique_lock dl(_delete_lock); - std::swap(_delete_set, old_delete_set); - } - - if (old_delete_set->find(_start) != old_delete_set->end()) - { - throw diskann::ANNException("ERROR: start node has been deleted", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - const uint32_t range = params.max_degree; - const uint32_t maxc = params.max_occlusion_size; - const float alpha = params.alpha; - const uint32_t num_threads = params.num_threads == 0 ? omp_get_num_procs() : params.num_threads; - - uint32_t num_calls_to_process_delete = 0; - diskann::Timer timer; -#pragma omp parallel for num_threads(num_threads) schedule(dynamic, 8192) reduction(+ : num_calls_to_process_delete) - for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) - { - if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && !_empty_slots.is_in_set((uint32_t)loc)) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } - } - for (int64_t loc = _max_points; loc < (int64_t)(_max_points + _num_frozen_pts); loc++) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } - - std::unique_lock tl(_tag_lock); - size_t ret_nd = release_locations(*old_delete_set); - size_t max_points = _max_points; - size_t empty_slots_size = _empty_slots.size(); - - std::shared_lock dl(_delete_lock); - size_t delete_set_size = _delete_set->size(); - size_t old_delete_set_size = old_delete_set->size(); - - if (!_conc_consolidate) - { - update_lock.unlock(); - } - - double duration = timer.elapsed() / 1000000.0; - diskann::cout << " done in " << duration << " seconds." << std::endl; - return consolidation_report(diskann::consolidation_report::status_code::SUCCESS, ret_nd, max_points, - empty_slots_size, old_delete_set_size, delete_set_size, num_calls_to_process_delete, - duration); -} - -template void Index::compact_frozen_point() -{ - if (_nd < _max_points && _num_frozen_pts > 0) - { - reposition_points((uint32_t)_max_points, (uint32_t)_nd, (uint32_t)_num_frozen_pts); - _start = (uint32_t)_nd; - - if (_filtered_index && _dynamic_index) - { - // update medoid id's as frozen points are treated as medoid - for (auto &[label, medoid_id] : _label_to_start_id) - { - /* if (label == _universal_label) - continue;*/ - _label_to_start_id[label] = (uint32_t)_nd + (medoid_id - (uint32_t)_max_points); - } - } - } -} - -// Should be called after acquiring _update_lock -template void Index::compact_data() -{ - if (!_dynamic_index) - throw ANNException("Can not compact a non-dynamic index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (_data_compacted) - { - diskann::cerr << "Warning! Calling compact_data() when _data_compacted is true!" << std::endl; - return; - } - - if (_delete_set->size() > 0) - { - throw ANNException("Can not compact data when index has non-empty _delete_set of " - "size: " + - std::to_string(_delete_set->size()), - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::Timer timer; - - std::vector new_location = std::vector(_max_points + _num_frozen_pts, UINT32_MAX); - - uint32_t new_counter = 0; - std::set empty_locations; - for (uint32_t old_location = 0; old_location < _max_points; old_location++) - { - if (_location_to_tag.contains(old_location)) - { - new_location[old_location] = new_counter; - new_counter++; - } - else - { - empty_locations.insert(old_location); - } - } - for (uint32_t old_location = (uint32_t)_max_points; old_location < _max_points + _num_frozen_pts; old_location++) - { - new_location[old_location] = old_location; - } - - // If start node is removed, throw an exception - if (_start < _max_points && !_location_to_tag.contains(_start)) - { - throw diskann::ANNException("ERROR: Start node deleted.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - size_t num_dangling = 0; - for (uint32_t old = 0; old < _max_points + _num_frozen_pts; ++old) - { - // compact _final_graph - std::vector new_adj_list; - - if ((new_location[old] < _max_points) // If point continues to exist - || (old >= _max_points && old < _max_points + _num_frozen_pts)) - { - new_adj_list.reserve(_graph_store->get_neighbours((location_t)old).size()); - for (auto ngh_iter : _graph_store->get_neighbours((location_t)old)) - { - if (empty_locations.find(ngh_iter) != empty_locations.end()) - { - ++num_dangling; - diskann::cerr << "Error in compact_data(). _final_graph[" << old << "] has neighbor " << ngh_iter - << " which is a location not associated with any tag." << std::endl; - } - else - { - new_adj_list.push_back(new_location[ngh_iter]); - } - } - //_graph_store->get_neighbours((location_t)old).swap(new_adj_list); - _graph_store->set_neighbours((location_t)old, new_adj_list); - - // Move the data and adj list to the correct position - if (new_location[old] != old) - { - assert(new_location[old] < old); - _graph_store->swap_neighbours(new_location[old], (location_t)old); - - if (_filtered_index) - { - _location_to_labels[new_location[old]].swap(_location_to_labels[old]); - } - - _data_store->copy_vectors(old, new_location[old], 1); - } - } - else - { - _graph_store->clear_neighbours((location_t)old); - } - } - diskann::cerr << "#dangling references after data compaction: " << num_dangling << std::endl; - - _tag_to_location.clear(); - for (auto pos = _location_to_tag.find_first(); pos.is_valid(); pos = _location_to_tag.find_next(pos)) - { - const auto tag = _location_to_tag.get(pos); - _tag_to_location[tag] = new_location[pos._key]; - } - _location_to_tag.clear(); - for (const auto &iter : _tag_to_location) - { - _location_to_tag.set(iter.second, iter.first); - } - // remove all cleared up old - for (size_t old = _nd; old < _max_points; ++old) - { - _graph_store->clear_neighbours((location_t)old); - } - if (_filtered_index) - { - for (size_t old = _nd; old < _max_points; old++) - { - _location_to_labels[old].clear(); - } - } - - _empty_slots.clear(); - // mark all slots after _nd as empty - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)i); - } - _data_compacted = true; - diskann::cout << "Time taken for compact_data: " << timer.elapsed() / 1000000. << "s." << std::endl; -} - -// -// Caller must hold unique _tag_lock and _delete_lock before calling this -// -template int Index::reserve_location() -{ - if (_nd >= _max_points) - { - return -1; - } - uint32_t location; - if (_data_compacted && _empty_slots.is_empty()) - { - // This code path is encountered when enable_delete hasn't been - // called yet, so no points have been deleted and _empty_slots - // hasn't been filled in. In that case, just keep assigning - // consecutive locations. - location = (uint32_t)_nd; - } - else - { - assert(_empty_slots.size() != 0); - assert(_empty_slots.size() + _nd == _max_points); - - location = _empty_slots.pop_any(); - _delete_set->erase(location); - } - ++_nd; - return location; -} - -template size_t Index::release_location(int location) -{ - if (_empty_slots.is_in_set(location)) - throw ANNException("Trying to release location, but location already in empty slots", -1, __FUNCSIG__, __FILE__, - __LINE__); - _empty_slots.insert(location); - - _nd--; - return _nd; -} - -template -size_t Index::release_locations(const tsl::robin_set &locations) -{ - for (auto location : locations) - { - if (_empty_slots.is_in_set(location)) - throw ANNException("Trying to release location, but location " - "already in empty slots", - -1, __FUNCSIG__, __FILE__, __LINE__); - _empty_slots.insert(location); - - _nd--; - } - - if (_empty_slots.size() + _nd != _max_points) - throw ANNException("#empty slots + nd != max points", -1, __FUNCSIG__, __FILE__, __LINE__); - - return _nd; -} - -template -void Index::reposition_points(uint32_t old_location_start, uint32_t new_location_start, - uint32_t num_locations) -{ - if (num_locations == 0 || old_location_start == new_location_start) - { - return; - } - - // Update pointers to the moved nodes. Note: the computation is correct even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - - std::vector updated_neighbours_location; - for (uint32_t i = 0; i < _max_points + _num_frozen_pts; i++) - { - auto &i_neighbours = _graph_store->get_neighbours((location_t)i); - std::vector i_neighbours_copy(i_neighbours.begin(), i_neighbours.end()); - for (auto &loc : i_neighbours_copy) - { - if (loc >= old_location_start && loc < old_location_start + num_locations) - loc += location_delta; - } - _graph_store->set_neighbours(i, i_neighbours_copy); - } - - // The [start, end) interval which will contain obsolete points to be - // cleared. - uint32_t mem_clear_loc_start = old_location_start; - uint32_t mem_clear_loc_end_limit = old_location_start + num_locations; - - // Move the adjacency lists. Make sure that overlapping ranges are handled - // correctly. - if (new_location_start < old_location_start) - { - // New location before the old location: copy the entries in order - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = 0; loc_offset < num_locations; loc_offset++) - { - assert(_graph_store->get_neighbours(new_location_start + loc_offset).empty()); - _graph_store->swap_neighbours(new_location_start + loc_offset, old_location_start + loc_offset); - if (_dynamic_index && _filtered_index) - { - _location_to_labels[new_location_start + loc_offset].swap( - _location_to_labels[old_location_start + loc_offset]); - } - } - // If ranges are overlapping, make sure not to clear the newly copied - // data. - if (mem_clear_loc_start < new_location_start + num_locations) - { - // Clear only after the end of the new range. - mem_clear_loc_start = new_location_start + num_locations; - } - } - else - { - // Old location after the new location: copy from the end of the range - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = num_locations; loc_offset > 0; loc_offset--) - { - assert(_graph_store->get_neighbours(new_location_start + loc_offset - 1u).empty()); - _graph_store->swap_neighbours(new_location_start + loc_offset - 1u, old_location_start + loc_offset - 1u); - if (_dynamic_index && _filtered_index) - { - _location_to_labels[new_location_start + loc_offset - 1u].swap( - _location_to_labels[old_location_start + loc_offset - 1u]); - } - } - - // If ranges are overlapping, make sure not to clear the newly copied - // data. - if (mem_clear_loc_end_limit > new_location_start) - { - // Clear only up to the beginning of the new range. - mem_clear_loc_end_limit = new_location_start; - } - } - _data_store->move_vectors(old_location_start, new_location_start, num_locations); -} - -template void Index::reposition_frozen_point_to_end() -{ - if (_num_frozen_pts == 0) - return; - - if (_nd == _max_points) - { - diskann::cout << "Not repositioning frozen point as it is already at the end." << std::endl; - return; - } - - reposition_points((uint32_t)_nd, (uint32_t)_max_points, (uint32_t)_num_frozen_pts); - _start = (uint32_t)_max_points; - - // update medoid id's as frozen points are treated as medoid - if (_filtered_index && _dynamic_index) - { - for (auto &[label, medoid_id] : _label_to_start_id) - { - /*if (label == _universal_label) - continue;*/ - _label_to_start_id[label] = (uint32_t)_max_points + (medoid_id - (uint32_t)_nd); - } - } -} - -template void Index::resize(size_t new_max_points) -{ - const size_t new_internal_points = new_max_points + _num_frozen_pts; - auto start = std::chrono::high_resolution_clock::now(); - assert(_empty_slots.size() == 0); // should not resize if there are empty slots. - - _data_store->resize((location_t)new_internal_points); - _graph_store->resize_graph(new_internal_points); - _locks = std::vector(new_internal_points); - - if (_num_frozen_pts != 0) - { - reposition_points((uint32_t)_max_points, (uint32_t)new_max_points, (uint32_t)_num_frozen_pts); - _start = (uint32_t)new_max_points; - } - - _max_points = new_max_points; - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)i); - } - - auto stop = std::chrono::high_resolution_clock::now(); - diskann::cout << "Resizing took: " << std::chrono::duration(stop - start).count() << "s" << std::endl; -} - -template -int Index::_insert_point(const DataType &point, const TagType tag) -{ - try - { - return this->insert_point(std::any_cast(point), std::any_cast(tag)); - } - catch (const std::bad_any_cast &anycast_e) - { - throw new ANNException("Error:Trying to insert invalid data type" + std::string(anycast_e.what()), -1); - } - catch (const std::exception &e) - { - throw new ANNException("Error:" + std::string(e.what()), -1); - } -} - -template -int Index::_insert_point(const DataType &point, const TagType tag, Labelvector &labels) -{ - try - { - return this->insert_point(std::any_cast(point), std::any_cast(tag), - labels.get>()); - } - catch (const std::bad_any_cast &anycast_e) - { - throw new ANNException("Error:Trying to insert invalid data type" + std::string(anycast_e.what()), -1); - } - catch (const std::exception &e) - { - throw new ANNException("Error:" + std::string(e.what()), -1); - } -} - -template -int Index::insert_point(const T *point, const TagT tag) -{ - std::vector no_labels{0}; - return insert_point(point, tag, no_labels); -} - -template -int Index::insert_point(const T *point, const TagT tag, const std::vector &labels) -{ - - assert(_has_built); - if (tag == 0) - { - throw diskann::ANNException("Do not insert point with tag 0. That is " - "reserved for points hidden " - "from the user.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::shared_lock shared_ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - auto location = reserve_location(); - if (_filtered_index) - { - if (labels.empty()) - { - release_location(location); - std::cerr << "Error: Can't insert point with tag " + get_tag_string(tag) + - " . there are no labels for the point." - << std::endl; - return -1; - } - - _location_to_labels[location] = labels; - - for (LabelT label : labels) - { - if (_labels.find(label) == _labels.end()) - { - if (_frozen_pts_used >= _num_frozen_pts) - { - throw ANNException( - "Error: For dynamic filtered index, the number of frozen points should be atleast equal " - "to number of unique labels.", - -1); - } - - auto fz_location = (int)(_max_points) + _frozen_pts_used; // as first _fz_point - _labels.insert(label); - _label_to_start_id[label] = (uint32_t)fz_location; - _location_to_labels[fz_location] = {label}; - _data_store->set_vector((location_t)fz_location, point); - _frozen_pts_used++; - } - } - } - - if (location == -1) - { -#if EXPAND_IF_FULL - dl.unlock(); - tl.unlock(); - shared_ul.unlock(); - - { - std::unique_lock ul(_update_lock); - tl.lock(); - dl.lock(); - - if (_nd >= _max_points) - { - auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); - resize(new_max_points); - } - - dl.unlock(); - tl.unlock(); - ul.unlock(); - } - - shared_ul.lock(); - tl.lock(); - dl.lock(); - - location = reserve_location(); - if (location == -1) - { - throw diskann::ANNException("Cannot reserve location even after " - "expanding graph. Terminating.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } -#else - return -1; -#endif - } // cant insert as active pts >= max_pts - dl.unlock(); - - // Insert tag and mapping to location - if (_enable_tags) - { - // if tags are enabled and tag is already inserted. so we can't reuse that tag. - if (_tag_to_location.find(tag) != _tag_to_location.end()) - { - release_location(location); - return -1; - } - - _tag_to_location[tag] = location; - _location_to_tag.set(location, tag); - } - tl.unlock(); - - _data_store->set_vector(location, point); // update datastore - - // Find and add appropriate graph edges - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - std::vector pruned_list; // it is the set best candidates to connect to this point - if (_filtered_index) - { - // when filtered the best_candidates will share the same label ( label_present > distance) - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch, true, _filterIndexingQueueSize); - } - else - { - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch); - } - assert(pruned_list.size() > 0); // should find atleast one neighbour (i.e frozen point acting as medoid) - - { - std::shared_lock tlock(_tag_lock, std::defer_lock); - if (_conc_consolidate) - tlock.lock(); - - LockGuard guard(_locks[location]); - _graph_store->clear_neighbours(location); - - std::vector neighbor_links; - for (auto link : pruned_list) - { - if (_conc_consolidate) - if (!_location_to_tag.contains(link)) - continue; - neighbor_links.emplace_back(link); - } - _graph_store->set_neighbours(location, neighbor_links); - assert(_graph_store->get_neighbours(location).size() <= _indexingRange); - - if (_conc_consolidate) - tlock.unlock(); - } - - inter_insert(location, pruned_list, scratch); - - return 0; -} - -template int Index::_lazy_delete(const TagType &tag) -{ - try - { - return lazy_delete(std::any_cast(tag)); - } - catch (const std::bad_any_cast &e) - { - throw ANNException(std::string("Error: ") + e.what(), -1); - } -} - -template -void Index::_lazy_delete(TagVector &tags, TagVector &failed_tags) -{ - try - { - this->lazy_delete(tags.get>(), failed_tags.get>()); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast while performing _lazy_delete() " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template int Index::lazy_delete(const TagT &tag) -{ - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; - - if (_tag_to_location.find(tag) == _tag_to_location.end()) - { - diskann::cerr << "Delete tag not found " << get_tag_string(tag) << std::endl; - return -1; - } - assert(_tag_to_location[tag] < _max_points); - - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); - return 0; -} - -template -void Index::lazy_delete(const std::vector &tags, std::vector &failed_tags) -{ - if (failed_tags.size() > 0) - { - throw ANNException("failed_tags should be passed as an empty list", -1, __FUNCSIG__, __FILE__, __LINE__); - } - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; - - for (auto tag : tags) - { - if (_tag_to_location.find(tag) == _tag_to_location.end()) - { - failed_tags.push_back(tag); - } - else - { - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); - } - } -} - -template bool Index::is_index_saved() -{ - return _is_saved; -} - -template -void Index::_get_active_tags(TagRobinSet &active_tags) -{ - try - { - this->get_active_tags(active_tags.get>()); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad_any cast while performing _get_active_tags() " + std::string(e.what()), -1); - } - catch (const std::exception &e) - { - throw ANNException("Error :" + std::string(e.what()), -1); - } -} - -template -void Index::get_active_tags(tsl::robin_set &active_tags) -{ - active_tags.clear(); - std::shared_lock tl(_tag_lock); - for (auto iter : _tag_to_location) - { - active_tags.insert(iter.first); - } -} - -template -void Index::get_degree_stats(size_t &max_deg, size_t &min_deg, size_t &avg_deg, size_t &cnt_deg) -{ - max_deg = 0; - min_deg = SIZE_MAX; - avg_deg = 0; - cnt_deg = 0; - size_t total = 0; - for (size_t i = 0; i < _nd; i++) - { - auto &pool = _graph_store->get_neighbours((location_t)i); - cnt_deg += (pool.size() < 2); - max_deg = std::max(max_deg, pool.size()); - min_deg = std::min(min_deg, pool.size()); - total += pool.size(); - } - avg_deg = total / _nd; -} - -template -void Index::dump_degree_stats(std::string filename) -{ - std::ofstream file(filename); - if (!file.is_open()) - { - std::cerr << "Error: Could not open file " << filename << " for writing" << std::endl; - return; - } - - // Write each node's degree to the file, one per line - for (size_t i = 0; i < _nd; i++) - { - auto &pool = _graph_store->get_neighbours((location_t)i); - file << pool.size() << std::endl; - } - - file.close(); -} - -template void Index::print_status() -{ - std::shared_lock ul(_update_lock); - std::shared_lock cl(_consolidate_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - - diskann::cout << "------------------- Index object: " << (uint64_t)this << " -------------------" << std::endl; - diskann::cout << "Number of points: " << _nd << std::endl; - diskann::cout << "Graph size: " << _graph_store->get_total_points() << std::endl; - diskann::cout << "Location to tag size: " << _location_to_tag.size() << std::endl; - diskann::cout << "Tag to location size: " << _tag_to_location.size() << std::endl; - diskann::cout << "Number of empty slots: " << _empty_slots.size() << std::endl; - diskann::cout << std::boolalpha << "Data compacted: " << this->_data_compacted << std::endl; - diskann::cout << "---------------------------------------------------------" - "------------" - << std::endl; -} - -template void Index::count_nodes_at_bfs_levels() -{ - std::unique_lock ul(_update_lock); - - boost::dynamic_bitset<> visited(_max_points + _num_frozen_pts); - - size_t MAX_BFS_LEVELS = 32; - auto bfs_sets = new tsl::robin_set[MAX_BFS_LEVELS]; - - bfs_sets[0].insert(_start); - visited.set(_start); - - for (uint32_t i = (uint32_t)_max_points; i < _max_points + _num_frozen_pts; ++i) - { - if (i != _start) - { - bfs_sets[0].insert(i); - visited.set(i); - } - } - - for (size_t l = 0; l < MAX_BFS_LEVELS - 1; ++l) - { - diskann::cout << "Number of nodes at BFS level " << l << " is " << bfs_sets[l].size() << std::endl; - if (bfs_sets[l].size() == 0) - break; - for (auto node : bfs_sets[l]) - { - for (auto nghbr : _graph_store->get_neighbours((location_t)node)) - { - if (!visited.test(nghbr)) - { - visited.set(nghbr); - bfs_sets[l + 1].insert(nghbr); - } - } - } - } - - delete[] bfs_sets; -} - -// REFACTOR: This should be an OptimizedDataStore class -template void Index::optimize_index_layout() -{ // use after build or load - if (_dynamic_index) - { - throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - float *cur_vec = new float[_data_store->get_aligned_dim()]; - std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); - _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); - _neighbor_len = (_graph_store->get_max_observed_degree() + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - auto dist_fast = (DistanceFastL2 *)(_data_store->get_dist_fn()); - for (uint32_t i = 0; i < _nd; i++) - { - char *cur_node_offset = _opt_graph + i * _node_size; - _data_store->get_vector(i, (T *)cur_vec); - float cur_norm = dist_fast->norm((T *)cur_vec, (uint32_t)_data_store->get_aligned_dim()); - std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); - std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); - - cur_node_offset += _data_len; - uint32_t k = (uint32_t)_graph_store->get_neighbours(i).size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _graph_store->get_neighbours(i).data(), k * sizeof(uint32_t)); - // std::vector().swap(_graph_store->get_neighbours(i)); - _graph_store->clear_neighbours(i); - } - _graph_store->clear_graph(); - _graph_store->resize_graph(0); - delete[] cur_vec; -} - -template -void Index::_search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) -{ - try - { - return this->search_with_optimized_layout(std::any_cast(query), K, L, indices); - } - catch (const std::bad_any_cast &e) - { - throw ANNException("Error: bad any cast while performing " - "_search_with_optimized_layout() " + - std::string(e.what()), - -1); - } - catch (const std::exception &e) - { - throw ANNException("Error: " + std::string(e.what()), -1); - } -} - -template -void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ - DistanceFastL2 *dist_fast = (DistanceFastL2 *)(_data_store->get_dist_fn()); - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t MaxM_ep = *neighbors; - neighbors++; - - for (; tmp_l < L && tmp_l < MaxM_ep; tmp_l++) - { - init_ids[tmp_l] = neighbors[tmp_l]; - flags[init_ids[tmp_l]] = true; - } - - while (tmp_l < L) - { - uint32_t id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - T *x = (T *)(_opt_graph + _node_size * id); - float norm_x = *x; - x++; - float dist = dist_fast->compare(x, query, norm_x, (uint32_t)_data_store->get_aligned_dim()); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) - { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; - neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) - { - uint32_t id = neighbors[m]; - if (flags[id]) - continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) - { - indices[i] = retset[i].id; - } -} - -/* Internals of the library */ -template const float Index::INDEX_GROWTH_FACTOR = 1.5f; - -// EXPORTS -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -// Label with short int 2 byte -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT class Index; - -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -// TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -// TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -// TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -// TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/index_factory.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/index_factory.cpp deleted file mode 100644 index 35790f8..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/index_factory.cpp +++ /dev/null @@ -1,213 +0,0 @@ -#include "index_factory.h" -#include "pq_l2_distance.h" - -namespace diskann -{ - -IndexFactory::IndexFactory(const IndexConfig &config) : _config(std::make_unique(config)) -{ - check_config(); -} - -std::unique_ptr IndexFactory::create_instance() -{ - return create_instance(_config->data_type, _config->tag_type, _config->label_type); -} - -void IndexFactory::check_config() -{ - if (_config->dynamic_index && !_config->enable_tags) - { - throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_config->pq_dist_build) - { - if (_config->dynamic_index) - throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (_config->metric == diskann::Metric::INNER_PRODUCT) - throw ANNException("ERROR: Inner product metrics not yet supported " - "with PQ distance " - "base index", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_config->data_type != "float" && _config->data_type != "uint8" && _config->data_type != "int8") - { - throw ANNException("ERROR: invalid data type : + " + _config->data_type + - " is not supported. please select from [float, int8, uint8]", - -1); - } - - if (_config->tag_type != "int32" && _config->tag_type != "uint32" && _config->tag_type != "int64" && - _config->tag_type != "uint64") - { - throw ANNException("ERROR: invalid data type : + " + _config->tag_type + - " is not supported. please select from [int32, uint32, int64, uint64]", - -1); - } -} - -template Distance *IndexFactory::construct_inmem_distance_fn(Metric metric) -{ - if (metric == diskann::Metric::COSINE && std::is_same::value) - { - return (Distance *)new AVXNormalizedCosineDistanceFloat(); - } - else - { - return (Distance *)get_distance_function(metric); - } -} - -template -std::shared_ptr> IndexFactory::construct_datastore(DataStoreStrategy strategy, - size_t total_internal_points, size_t dimension, - Metric metric) -{ - std::unique_ptr> distance; - switch (strategy) - { - case DataStoreStrategy::MEMORY: - distance.reset(construct_inmem_distance_fn(metric)); - return std::make_shared>((location_t)total_internal_points, dimension, - std::move(distance)); - default: - break; - } - return nullptr; -} - -std::unique_ptr IndexFactory::construct_graphstore(const GraphStoreStrategy strategy, - const size_t size, - const size_t reserve_graph_degree) -{ - switch (strategy) - { - case GraphStoreStrategy::MEMORY: - return std::make_unique(size, reserve_graph_degree); - default: - throw ANNException("Error : Current GraphStoreStratagy is not supported.", -1); - } -} - -template -std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreStrategy strategy, size_t num_points, - size_t dimension, Metric m, size_t num_pq_chunks, - bool use_opq) -{ - std::unique_ptr> distance_fn; - std::unique_ptr> quantized_distance_fn; - - quantized_distance_fn = std::move(std::make_unique>((uint32_t)num_pq_chunks, use_opq)); - switch (strategy) - { - case DataStoreStrategy::MEMORY: - distance_fn.reset(construct_inmem_distance_fn(m)); - return std::make_shared>(dimension, (location_t)(num_points), num_pq_chunks, - std::move(distance_fn), std::move(quantized_distance_fn)); - default: - // REFACTOR TODO: We do support diskPQ - so we may need to add a new class for SSDPQDataStore! - break; - } - return nullptr; -} - -template -std::unique_ptr IndexFactory::create_instance() -{ - size_t num_points = _config->max_points + _config->num_frozen_pts; - size_t dim = _config->dimension; - // auto graph_store = construct_graphstore(_config->graph_strategy, num_points); - auto data_store = construct_datastore(_config->data_strategy, num_points, dim, _config->metric); - std::shared_ptr> pq_data_store = nullptr; - - if (_config->data_strategy == DataStoreStrategy::MEMORY && _config->pq_dist_build) - { - pq_data_store = - construct_pq_datastore(_config->data_strategy, num_points + _config->num_frozen_pts, dim, - _config->metric, _config->num_pq_chunks, _config->use_opq); - } - else - { - pq_data_store = data_store; - } - size_t max_reserve_degree = - (size_t)(defaults::GRAPH_SLACK_FACTOR * 1.05 * - (_config->index_write_params == nullptr ? 0 : _config->index_write_params->max_degree)); - std::unique_ptr graph_store = - construct_graphstore(_config->graph_strategy, num_points + _config->num_frozen_pts, max_reserve_degree); - - // REFACTOR TODO: Must construct in-memory PQDatastore if strategy == ONDISK and must construct - // in-mem and on-disk PQDataStore if strategy == ONDISK and diskPQ is required. - return std::make_unique>(*_config, data_store, - std::move(graph_store), pq_data_store); -} - -std::unique_ptr IndexFactory::create_instance(const std::string &data_type, const std::string &tag_type, - const std::string &label_type) -{ - if (data_type == std::string("float")) - { - return create_instance(tag_type, label_type); - } - else if (data_type == std::string("uint8")) - { - return create_instance(tag_type, label_type); - } - else if (data_type == std::string("int8")) - { - return create_instance(tag_type, label_type); - } - else - throw ANNException("Error: unsupported data_type please choose from [float/int8/uint8]", -1); -} - -template -std::unique_ptr IndexFactory::create_instance(const std::string &tag_type, const std::string &label_type) -{ - if (tag_type == std::string("int32")) - { - return create_instance(label_type); - } - else if (tag_type == std::string("uint32")) - { - return create_instance(label_type); - } - else if (tag_type == std::string("int64")) - { - return create_instance(label_type); - } - else if (tag_type == std::string("uint64")) - { - return create_instance(label_type); - } - else - throw ANNException("Error: unsupported tag_type please choose from [int32/uint32/int64/uint64]", -1); -} - -template -std::unique_ptr IndexFactory::create_instance(const std::string &label_type) -{ - if (label_type == std::string("uint16") || label_type == std::string("ushort")) - { - return create_instance(); - } - else if (label_type == std::string("uint32") || label_type == std::string("uint")) - { - return create_instance(); - } - else - throw ANNException("Error: unsupported label_type please choose from [uint/ushort]", -1); -} - -// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); -// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); -// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/linux_aligned_file_reader.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/linux_aligned_file_reader.cpp deleted file mode 100644 index 64e7eee..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/linux_aligned_file_reader.cpp +++ /dev/null @@ -1,230 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "linux_aligned_file_reader.h" -#ifndef __APPLE__ - -#include -#include -#include -#include "tsl/robin_map.h" -#include "utils.h" -#define MAX_EVENTS 1024 - -namespace -{ -typedef struct io_event io_event_t; -typedef struct iocb iocb_t; - -void execute_io(io_context_t ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) -{ -#ifdef DEBUG - for (auto &req : read_reqs) - { - assert(IS_ALIGNED(req.len, 512)); - // std::cout << "request:"<= req.len); - } -#endif - - // break-up requests into chunks of size MAX_EVENTS each - uint64_t n_iters = ROUND_UP(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS; - for (uint64_t iter = 0; iter < n_iters; iter++) - { - uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), (uint64_t)MAX_EVENTS); - std::vector cbs(n_ops, nullptr); - std::vector evts(n_ops); - std::vector cb(n_ops); - for (uint64_t j = 0; j < n_ops; j++) - { - io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, read_reqs[j + iter * MAX_EVENTS].len, - read_reqs[j + iter * MAX_EVENTS].offset); - } - - // initialize `cbs` using `cb` array - // - - for (uint64_t i = 0; i < n_ops; i++) - { - cbs[i] = cb.data() + i; - } - - uint64_t n_tries = 0; - while (n_tries <= n_retries) - { - // issue reads - int64_t ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); - // if requests didn't get accepted - if (ret != (int64_t)n_ops) - { - std::cerr << "io_submit() failed; returned " << ret << ", expected=" << n_ops << ", ernno=" << errno - << "=" << ::strerror(-ret) << ", try #" << n_tries + 1; - std::cout << "ctx: " << ctx << "\n"; - exit(-1); - } - else - { - // wait on io_getevents - ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), nullptr); - // if requests didn't complete - if (ret != (int64_t)n_ops) - { - std::cerr << "io_getevents() failed; returned " << ret << ", expected=" << n_ops - << ", ernno=" << errno << "=" << ::strerror(-ret) << ", try #" << n_tries + 1; - exit(-1); - } - else - { - break; - } - } - } - // disabled since req.buf could be an offset into another buf - /* - for (auto &req : read_reqs) { - // corruption check - assert(malloc_usable_size(req.buf) >= req.len); - } - */ - } -} -} // namespace - -LinuxAlignedFileReader::LinuxAlignedFileReader() -{ - this->file_desc = -1; -} - -LinuxAlignedFileReader::~LinuxAlignedFileReader() -{ - int64_t ret; - // check to make sure file_desc is closed - ret = ::fcntl(this->file_desc, F_GETFD); - if (ret == -1) - { - if (errno != EBADF) - { - std::cerr << "close() not called" << std::endl; - // close file desc - ret = ::close(this->file_desc); - // error checks - if (ret == -1) - { - std::cerr << "close() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) - << std::endl; - } - } - } -} - -io_context_t &LinuxAlignedFileReader::get_ctx() -{ - std::unique_lock lk(ctx_mut); - // perform checks only in DEBUG mode - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) - { - std::cerr << "bad thread access; returning -1 as io_context_t" << std::endl; - return this->bad_ctx; - } - else - { - return ctx_map[std::this_thread::get_id()]; - } -} - -void LinuxAlignedFileReader::register_thread() -{ - auto my_id = std::this_thread::get_id(); - std::unique_lock lk(ctx_mut); - if (ctx_map.find(my_id) != ctx_map.end()) - { - std::cerr << "multiple calls to register_thread from the same thread" << std::endl; - return; - } - io_context_t ctx = 0; - int ret = io_setup(MAX_EVENTS, &ctx); - if (ret != 0) - { - lk.unlock(); - if (ret == -EAGAIN) - { - std::cerr << "io_setup() failed with EAGAIN: Consider increasing /proc/sys/fs/aio-max-nr" << std::endl; - } - else - { - std::cerr << "io_setup() failed; returned " << ret << ": " << ::strerror(-ret) << std::endl; - } - } - else - { - diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id << std::endl; - ctx_map[my_id] = ctx; - } - lk.unlock(); -} - -void LinuxAlignedFileReader::deregister_thread() -{ - auto my_id = std::this_thread::get_id(); - std::unique_lock lk(ctx_mut); - assert(ctx_map.find(my_id) != ctx_map.end()); - - lk.unlock(); - io_context_t ctx = this->get_ctx(); - io_destroy(ctx); - // assert(ret == 0); - lk.lock(); - ctx_map.erase(my_id); - std::cerr << "returned ctx from thread-id:" << my_id << std::endl; - lk.unlock(); -} - -void LinuxAlignedFileReader::deregister_all_threads() -{ - std::unique_lock lk(ctx_mut); - for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) - { - io_context_t ctx = x.value(); - io_destroy(ctx); - // assert(ret == 0); - // lk.lock(); - // ctx_map.erase(my_id); - // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; - } - ctx_map.clear(); - // lk.unlock(); -} - -void LinuxAlignedFileReader::open(const std::string &fname) -{ - int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; - this->file_desc = ::open(fname.c_str(), flags); - // error checks - assert(this->file_desc != -1); - std::cerr << "Opened file : " << fname << std::endl; -} - -void LinuxAlignedFileReader::close() -{ - // int64_t ret; - - // check to make sure file_desc is closed - ::fcntl(this->file_desc, F_GETFD); - // assert(ret != -1); - - ::close(this->file_desc); - // assert(ret != -1); -} - -void LinuxAlignedFileReader::read(std::vector &read_reqs, io_context_t &ctx, bool async) -{ - if (async == true) - { - diskann::cout << "Async currently not supported in linux." << std::endl; - } - assert(this->file_desc != -1); - execute_io(ctx, this->file_desc, read_reqs); -} -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/logger.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/logger.cpp deleted file mode 100644 index 052f548..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/logger.cpp +++ /dev/null @@ -1,97 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include - -#include "logger_impl.h" -#include "windows_customizations.h" - -namespace diskann -{ - -#ifdef ENABLE_CUSTOM_LOGGER -DISKANN_DLLEXPORT ANNStreamBuf coutBuff(stdout); -DISKANN_DLLEXPORT ANNStreamBuf cerrBuff(stderr); - -DISKANN_DLLEXPORT std::basic_ostream cout(&coutBuff); -DISKANN_DLLEXPORT std::basic_ostream cerr(&cerrBuff); -std::function g_logger; - -void SetCustomLogger(std::function logger) -{ - g_logger = logger; - diskann::cout << "Set Custom Logger" << std::endl; -} - -ANNStreamBuf::ANNStreamBuf(FILE *fp) -{ - if (fp == nullptr) - { - throw diskann::ANNException("File pointer passed to ANNStreamBuf() cannot be null", -1); - } - if (fp != stdout && fp != stderr) - { - throw diskann::ANNException("The custom logger only supports stdout and stderr.", -1); - } - _fp = fp; - _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error; - _buf = new char[BUFFER_SIZE + 1]; // See comment in the header - - std::memset(_buf, 0, (BUFFER_SIZE) * sizeof(char)); - setp(_buf, _buf + BUFFER_SIZE - 1); -} - -ANNStreamBuf::~ANNStreamBuf() -{ - sync(); - _fp = nullptr; // we'll not close because we can't. - delete[] _buf; -} - -int ANNStreamBuf::overflow(int c) -{ - std::lock_guard lock(_mutex); - if (c != EOF) - { - *pptr() = (char)c; - pbump(1); - } - flush(); - return c; -} - -int ANNStreamBuf::sync() -{ - std::lock_guard lock(_mutex); - flush(); - return 0; -} - -int ANNStreamBuf::underflow() -{ - throw diskann::ANNException("Attempt to read on streambuf meant only for writing.", -1); -} - -int ANNStreamBuf::flush() -{ - const int num = (int)(pptr() - pbase()); - logImpl(pbase(), num); - pbump(-num); - return num; -} -void ANNStreamBuf::logImpl(char *str, int num) -{ - str[num] = '\0'; // Safe. See the c'tor. - // Invoke the OLS custom logging function. - if (g_logger) - { - g_logger(_logLevel, str); - } -} -#else -using std::cerr; -using std::cout; -#endif - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/math_utils.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/math_utils.cpp deleted file mode 100644 index d8fcda3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/math_utils.cpp +++ /dev/null @@ -1,465 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#ifdef __APPLE__ -#include -#else -#include -#endif -#include "logger.h" -#include "utils.h" - -namespace math_utils -{ - -#ifdef __APPLE__ -typedef int MKL_INT; -#endif - -float calc_distance(float *vec_1, float *vec_2, size_t dim) -{ - float dist = 0; - for (size_t j = 0; j < dim; j++) - { - dist += (vec_1[j] - vec_2[j]) * (vec_1[j] - vec_2[j]); - } - return dist; -} - -// compute l2-squared norms of data stored in row major num_points * dim, -// needs -// to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim) -{ -#pragma omp parallel for schedule(static, 8192) - for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) - { - vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); - vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; - } -} - -void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, - bool transpose_rot) -{ - CBLAS_TRANSPOSE transpose = CblasNoTrans; - if (transpose_rot) - { - diskann::cout << "Transposing rotation matrix.." << std::flush; - transpose = CblasTrans; - } - diskann::cout << "done Rotating data with random matrix.." << std::flush; - - cblas_sgemm(CblasRowMajor, CblasNoTrans, transpose, (MKL_INT)num_points, (MKL_INT)dim, (MKL_INT)dim, 1.0, data, - (MKL_INT)dim, rot_mat, (MKL_INT)dim, 0, new_mat, (MKL_INT)dim); - - diskann::cout << "done." << std::endl; -} - -// calculate k closest centers to data of num_points * dim (row major) -// centers is num_centers * dim (row major) -// data_l2sq has pre-computed squared norms of data -// centers_l2sq has pre-computed squared norms of centers -// pre-allocated center_index will contain id of nearest center -// pre-allocated dist_matrix shound be num_points * num_centers and contain -// squared distances -// Default value of k is 1 - -// Ideally used only by compute_closest_centers -void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k) -{ - if (k > num_centers) - { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; - return; - } - - float *ones_a = new float[num_centers]; - float *ones_b = new float[num_points]; - - for (size_t i = 0; i < num_centers; i++) - { - ones_a[i] = 1.0; - } - for (size_t i = 0; i < num_points; i++) - { - ones_b[i] = 1.0; - } - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, - docs_l2sq, (MKL_INT)1, ones_a, (MKL_INT)1, 0.0f, dist_matrix, (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, - ones_b, (MKL_INT)1, centers_l2sq, (MKL_INT)1, 1.0f, dist_matrix, (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)dim, -2.0f, - data, (MKL_INT)dim, centers, (MKL_INT)dim, 1.0f, dist_matrix, (MKL_INT)num_centers); - - if (k == 1) - { -#pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - float min = std::numeric_limits::max(); - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) - { - if (current[j] < min) - { - center_index[i] = (uint32_t)j; - min = current[j]; - } - } - } - } - else - { -#pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - std::priority_queue top_k_queue; - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) - { - PivotContainer this_piv(j, current[j]); - top_k_queue.push(this_piv); - } - for (size_t j = 0; j < k; j++) - { - PivotContainer this_piv = top_k_queue.top(); - center_index[i * k + j] = (uint32_t)this_piv.piv_id; - top_k_queue.pop(); - } - } - } - delete[] ones_a; - delete[] ones_b; -} - -// Given data in num_points * new_dim row major -// Pivots stored in full_pivot_data as num_centers * new_dim row major -// Calculate the k closest pivot for each point and store it in vector -// closest_centers_ivf (row major, num_points*k) (which needs to be allocated -// outside) Additionally, if inverted index is not null (and pre-allocated), -// it -// will return inverted index for each center, assuming each of the inverted -// indices is an empty vector. Additionally, if pts_norms_squared is not null, -// then it will assume that point norms are pre-computed and use those values - -void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, - size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index, - float *pts_norms_squared) -{ - if (k > num_centers) - { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; - return; - } - - bool is_norm_given_for_pts = (pts_norms_squared != NULL); - - float *pivs_norms_squared = new float[num_centers]; - if (!is_norm_given_for_pts) - pts_norms_squared = new float[num_points]; - - size_t PAR_BLOCK_SIZE = num_points; - size_t N_BLOCKS = - (num_points % PAR_BLOCK_SIZE) == 0 ? (num_points / PAR_BLOCK_SIZE) : (num_points / PAR_BLOCK_SIZE) + 1; - - if (!is_norm_given_for_pts) - math_utils::compute_vecs_l2sq(pts_norms_squared, data, num_points, dim); - math_utils::compute_vecs_l2sq(pivs_norms_squared, pivot_data, num_centers, dim); - uint32_t *closest_centers = new uint32_t[PAR_BLOCK_SIZE * k]; - float *distance_matrix = new float[num_centers * PAR_BLOCK_SIZE]; - - for (size_t cur_blk = 0; cur_blk < N_BLOCKS; cur_blk++) - { - float *data_cur_blk = data + cur_blk * PAR_BLOCK_SIZE * dim; - size_t num_pts_blk = std::min(PAR_BLOCK_SIZE, num_points - cur_blk * PAR_BLOCK_SIZE); - float *pts_norms_blk = pts_norms_squared + cur_blk * PAR_BLOCK_SIZE; - - math_utils::compute_closest_centers_in_block(data_cur_blk, num_pts_blk, dim, pivot_data, num_centers, - pts_norms_blk, pivs_norms_squared, closest_centers, - distance_matrix, k); - -#pragma omp parallel for schedule(static, 1) - for (int64_t j = cur_blk * PAR_BLOCK_SIZE; - j < std::min((int64_t)num_points, (int64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); j++) - { - for (size_t l = 0; l < k; l++) - { - size_t this_center_id = closest_centers[(j - cur_blk * PAR_BLOCK_SIZE) * k + l]; - closest_centers_ivf[j * k + l] = (uint32_t)this_center_id; - if (inverted_index != NULL) - { -#pragma omp critical - inverted_index[this_center_id].push_back(j); - } - } - } - } - delete[] closest_centers; - delete[] distance_matrix; - delete[] pivs_norms_squared; - if (!is_norm_given_for_pts) - delete[] pts_norms_squared; -} - -// if to_subtract is 1, will subtract nearest center from each row. Else will -// add. Output will be in data_load iself. -// Nearest centers need to be provided in closst_centers. -void process_residuals(float *data_load, size_t num_points, size_t dim, float *cur_pivot_data, size_t num_centers, - uint32_t *closest_centers, bool to_subtract) -{ - diskann::cout << "Processing residuals of " << num_points << " points in " << dim << " dimensions using " - << num_centers << " centers " << std::endl; -#pragma omp parallel for schedule(static, 8192) - for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) - { - for (size_t d_iter = 0; d_iter < dim; d_iter++) - { - if (to_subtract == 1) - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] - cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; - else - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] + cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; - } - } -} - -} // namespace math_utils - -namespace kmeans -{ - -// run Lloyds one iteration -// Given data in row major num_points * dim, and centers in row major -// num_centers * dim And squared lengths of data points, output the closest -// center to each data point, update centers, and also return inverted index. -// If -// closest_centers == NULL, will allocate memory and return. Similarly, if -// closest_docs == NULL, will allocate memory and return. - -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, - std::vector *closest_docs, uint32_t *&closest_center) -{ - bool compute_residual = true; - // Timer timer; - - if (closest_center == NULL) - closest_center = new uint32_t[num_points]; - if (closest_docs == NULL) - closest_docs = new std::vector[num_centers]; - else - for (size_t c = 0; c < num_centers; ++c) - closest_docs[c].clear(); - - math_utils::compute_closest_centers(data, num_points, dim, centers, num_centers, 1, closest_center, closest_docs, - docs_l2sq); - - memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); - -#pragma omp parallel for schedule(static, 1) - for (int64_t c = 0; c < (int64_t)num_centers; ++c) - { - float *center = centers + (size_t)c * (size_t)dim; - double *cluster_sum = new double[dim]; - for (size_t i = 0; i < dim; i++) - cluster_sum[i] = 0.0; - for (size_t i = 0; i < closest_docs[c].size(); i++) - { - float *current = data + ((closest_docs[c][i]) * dim); - for (size_t j = 0; j < dim; j++) - { - cluster_sum[j] += (double)current[j]; - } - } - if (closest_docs[c].size() > 0) - { - for (size_t i = 0; i < dim; i++) - center[i] = (float)(cluster_sum[i] / ((double)closest_docs[c].size())); - } - delete[] cluster_sum; - } - - float residual = 0.0; - if (compute_residual) - { - size_t BUF_PAD = 32; - size_t CHUNK_SIZE = 2 * 8192; - size_t nchunks = num_points / CHUNK_SIZE + (num_points % CHUNK_SIZE == 0 ? 0 : 1); - std::vector residuals(nchunks * BUF_PAD, 0.0); - -#pragma omp parallel for schedule(static, 32) - for (int64_t chunk = 0; chunk < (int64_t)nchunks; ++chunk) - for (size_t d = chunk * CHUNK_SIZE; d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) - residuals[chunk * BUF_PAD] += - math_utils::calc_distance(data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, dim); - - for (size_t chunk = 0; chunk < nchunks; ++chunk) - residual += residuals[chunk * BUF_PAD]; - } - - return residual; -} - -// Run Lloyds until max_reps or stopping criterion -// If you pass NULL for closest_docs and closest_center, it will NOT return -// the -// results, else it will assume appriate allocation as closest_docs = new -// vector [num_centers], and closest_center = new size_t[num_points] -// Final centers are output in centers as row major num_centers * dim -// -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, - const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center) -{ - float residual = std::numeric_limits::max(); - bool ret_closest_docs = true; - bool ret_closest_center = true; - if (closest_docs == NULL) - { - closest_docs = new std::vector[num_centers]; - ret_closest_docs = false; - } - if (closest_center == NULL) - { - closest_center = new uint32_t[num_points]; - ret_closest_center = false; - } - - float *docs_l2sq = new float[num_points]; - math_utils::compute_vecs_l2sq(docs_l2sq, data, num_points, dim); - - float old_residual; - // Timer timer; - for (size_t i = 0; i < max_reps; ++i) - { - old_residual = residual; - - residual = lloyds_iter(data, num_points, dim, centers, num_centers, docs_l2sq, closest_docs, closest_center); - - if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || - (residual < std::numeric_limits::epsilon())) - { - diskann::cout << "Residuals unchanged: " << old_residual << " becomes " << residual - << ". Early termination." << std::endl; - break; - } - } - delete[] docs_l2sq; - if (!ret_closest_docs) - delete[] closest_docs; - if (!ret_closest_center) - delete[] closest_center; - return residual; -} - -// assumes memory allocated for pivot_data as new -// float[num_centers*dim] -// and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) -{ - // pivot_data = new float[num_centers * dim]; - - std::vector picked; - std::random_device rd; - auto x = rd(); - std::mt19937 generator(x); - std::uniform_int_distribution distribution(0, num_points - 1); - - size_t tmp_pivot; - for (size_t j = 0; j < num_centers; j++) - { - tmp_pivot = distribution(generator); - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end()) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + j * dim, data + tmp_pivot * dim, dim * sizeof(float)); - } -} - -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) -{ - if (num_points > 1 << 23) - { - diskann::cout << "ERROR: n_pts " << num_points - << " currently not supported for k-means++, maximum is " - "8388608. Falling back to random pivot " - "selection." - << std::endl; - selecting_pivots(data, num_points, dim, pivot_data, num_centers); - return; - } - - std::vector picked; - std::random_device rd; - auto x = rd(); - std::mt19937 generator(x); - std::uniform_real_distribution<> distribution(0, 1); - std::uniform_int_distribution int_dist(0, num_points - 1); - size_t init_id = int_dist(generator); - size_t num_picked = 1; - - picked.push_back(init_id); - std::memcpy(pivot_data, data + init_id * dim, dim * sizeof(float)); - - float *dist = new float[num_points]; - -#pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - dist[i] = math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); - } - - double dart_val; - size_t tmp_pivot; - bool sum_flag = false; - - while (num_picked < num_centers) - { - dart_val = distribution(generator); - - double sum = 0; - for (size_t i = 0; i < num_points; i++) - { - sum = sum + dist[i]; - } - if (sum == 0) - sum_flag = true; - - dart_val *= sum; - - double prefix_sum = 0; - for (size_t i = 0; i < (num_points); i++) - { - tmp_pivot = i; - if (dart_val >= prefix_sum && dart_val < prefix_sum + dist[i]) - { - break; - } - - prefix_sum += dist[i]; - } - - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && (sum_flag == false)) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, dim * sizeof(float)); - -#pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - dist[i] = (std::min)(dist[i], math_utils::calc_distance(data + i * dim, data + tmp_pivot * dim, dim)); - } - num_picked++; - } - delete[] dist; -} - -} // namespace kmeans diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/memory_mapper.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/memory_mapper.cpp deleted file mode 100644 index d1c5ef9..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/memory_mapper.cpp +++ /dev/null @@ -1,107 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "logger.h" -#include "memory_mapper.h" -#include -#include - -using namespace diskann; - -MemoryMapper::MemoryMapper(const std::string &filename) : MemoryMapper(filename.c_str()) -{ -} - -MemoryMapper::MemoryMapper(const char *filename) -{ -#ifndef _WINDOWS - _fd = open(filename, O_RDONLY); - if (_fd <= 0) - { - std::cerr << "Inner vertices file not found" << std::endl; - return; - } - struct stat sb; - if (fstat(_fd, &sb) != 0) - { - std::cerr << "Inner vertices file not dound. " << std::endl; - return; - } - _fileSize = sb.st_size; - diskann::cout << "File Size: " << _fileSize << std::endl; - _buf = (char *)mmap(NULL, _fileSize, PROT_READ, MAP_PRIVATE, _fd, 0); -#else - _bareFile = - CreateFileA(filename, GENERIC_READ | GENERIC_EXECUTE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - if (_bareFile == nullptr) - { - std::ostringstream message; - message << "CreateFileA(" << filename << ") failed with error " << GetLastError() << std::endl; - std::cerr << message.str(); - throw std::exception(message.str().c_str()); - } - - _fd = CreateFileMapping(_bareFile, NULL, PAGE_EXECUTE_READ, 0, 0, NULL); - if (_fd == nullptr) - { - std::ostringstream message; - message << "CreateFileMapping(" << filename << ") failed with error " << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } - - _buf = (char *)MapViewOfFile(_fd, FILE_MAP_READ, 0, 0, 0); - if (_buf == nullptr) - { - std::ostringstream message; - message << "MapViewOfFile(" << filename << ") failed with error: " << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } - - LARGE_INTEGER fSize; - if (TRUE == GetFileSizeEx(_bareFile, &fSize)) - { - _fileSize = fSize.QuadPart; // take the 64-bit value - diskann::cout << "File Size: " << _fileSize << std::endl; - } - else - { - std::cerr << "Failed to get size of file " << filename << std::endl; - } -#endif -} -char *MemoryMapper::getBuf() -{ - return _buf; -} - -size_t MemoryMapper::getFileSize() -{ - return _fileSize; -} - -MemoryMapper::~MemoryMapper() -{ -#ifndef _WINDOWS - if (munmap(_buf, _fileSize) != 0) - std::cerr << "ERROR unmapping. CHECK!" << std::endl; - close(_fd); -#else - if (FALSE == UnmapViewOfFile(_buf)) - { - std::cerr << "Unmap view of file failed. Error: " << GetLastError() << std::endl; - } - - if (FALSE == CloseHandle(_fd)) - { - std::cerr << "Failed to close memory mapped file. Error: " << GetLastError() << std::endl; - } - - if (FALSE == CloseHandle(_bareFile)) - { - std::cerr << "Failed to close file: " << _fileName << " Error: " << GetLastError() << std::endl; - } - -#endif -} diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_map.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_map.cpp deleted file mode 100644 index a996dcf..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_map.cpp +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include - -#include "natural_number_map.h" -#include "tag_uint128.h" - -namespace diskann -{ -static constexpr auto invalid_position = boost::dynamic_bitset<>::npos; - -template -natural_number_map::natural_number_map() - : _size(0), _values_bitset(std::make_unique>()) -{ -} - -template void natural_number_map::reserve(size_t count) -{ - _values_vector.reserve(count); - _values_bitset->reserve(count); -} - -template size_t natural_number_map::size() const -{ - return _size; -} - -template void natural_number_map::set(Key key, Value value) -{ - if (key >= _values_bitset->size()) - { - _values_bitset->resize(static_cast(key) + 1); - _values_vector.resize(_values_bitset->size()); - } - - _values_vector[key] = value; - const bool was_present = _values_bitset->test_set(key, true); - - if (!was_present) - { - ++_size; - } -} - -template void natural_number_map::erase(Key key) -{ - if (key < _values_bitset->size()) - { - const bool was_present = _values_bitset->test_set(key, false); - - if (was_present) - { - --_size; - } - } -} - -template bool natural_number_map::contains(Key key) const -{ - return key < _values_bitset->size() && _values_bitset->test(key); -} - -template bool natural_number_map::try_get(Key key, Value &value) const -{ - if (!contains(key)) - { - return false; - } - - value = _values_vector[key]; - return true; -} - -template -typename natural_number_map::position natural_number_map::find_first() const -{ - return position{_size > 0 ? _values_bitset->find_first() : invalid_position, 0}; -} - -template -typename natural_number_map::position natural_number_map::find_next( - const position &after_position) const -{ - return position{after_position._keys_already_enumerated < _size ? _values_bitset->find_next(after_position._key) - : invalid_position, - after_position._keys_already_enumerated + 1}; -} - -template bool natural_number_map::position::is_valid() const -{ - return _key != invalid_position; -} - -template Value natural_number_map::get(const position &pos) const -{ - assert(pos.is_valid()); - return _values_vector[pos._key]; -} - -template void natural_number_map::clear() -{ - _size = 0; - _values_vector.clear(); - _values_bitset->clear(); -} - -// Instantiate used templates. -template class natural_number_map; -template class natural_number_map; -template class natural_number_map; -template class natural_number_map; -template class natural_number_map; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_set.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_set.cpp deleted file mode 100644 index b36cb52..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/natural_number_set.cpp +++ /dev/null @@ -1,70 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include - -#include "ann_exception.h" -#include "natural_number_set.h" - -namespace diskann -{ -template -natural_number_set::natural_number_set() : _values_bitset(std::make_unique>()) -{ -} - -template bool natural_number_set::is_empty() const -{ - return _values_vector.empty(); -} - -template void natural_number_set::reserve(size_t count) -{ - _values_vector.reserve(count); - _values_bitset->reserve(count); -} - -template void natural_number_set::insert(T id) -{ - _values_vector.emplace_back(id); - - if (id >= _values_bitset->size()) - _values_bitset->resize(static_cast(id) + 1); - - _values_bitset->set(id, true); -} - -template T natural_number_set::pop_any() -{ - if (_values_vector.empty()) - { - throw diskann::ANNException("No values available", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - const T id = _values_vector.back(); - _values_vector.pop_back(); - - _values_bitset->set(id, false); - - return id; -} - -template void natural_number_set::clear() -{ - _values_vector.clear(); - _values_bitset->clear(); -} - -template size_t natural_number_set::size() const -{ - return _values_vector.size(); -} - -template bool natural_number_set::is_in_set(T id) const -{ - return _values_bitset->test(id); -} - -// Instantiate used templates. -template class natural_number_set; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/partition.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/partition.cpp deleted file mode 100644 index 7e100ad..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/partition.cpp +++ /dev/null @@ -1,657 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include - -#include -#include "tsl/robin_map.h" -#include "tsl/robin_set.h" - -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) -#include "gperftools/malloc_extension.h" -#endif - -#include "utils.h" -#include "math_utils.h" -#include "index.h" -#include "parameters.h" -#include "memory_mapper.h" -#include "partition.h" -#ifdef _WINDOWS -#include -#endif - -// block size for reading/ processing large files and matrices in blocks -#define BLOCK_SIZE 5000000 - -// #define SAVE_INFLATED_PQ true - -template -void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(base_file.c_str(), read_blk_size); - std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), std::ios::binary); - std::ofstream sample_id_writer(std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); - - std::random_device rd; // Will be used to obtain a seed for the random number engine - auto x = rd(); - std::mt19937 generator(x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - size_t npts, nd; - uint32_t npts_u32, nd_u32; - uint32_t num_sampled_pts_u32 = 0; - uint32_t one_const = 1; - - base_reader.read((char *)&npts_u32, sizeof(uint32_t)); - base_reader.read((char *)&nd_u32, sizeof(uint32_t)); - diskann::cout << "Loading base " << base_file << ". #points: " << npts_u32 << ". #dim: " << nd_u32 << "." - << std::endl; - sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.write((char *)&nd_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&one_const, sizeof(uint32_t)); - - npts = npts_u32; - nd = nd_u32; - std::unique_ptr cur_row = std::make_unique(nd); - - for (size_t i = 0; i < npts; i++) - { - base_reader.read((char *)cur_row.get(), sizeof(T) * nd); - float sample = distribution(generator); - if (sample < sampling_rate) - { - sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); - uint32_t cur_i_u32 = (uint32_t)i; - sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); - num_sampled_pts_u32++; - } - } - sample_writer.seekp(0, std::ios::beg); - sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_id_writer.seekp(0, std::ios::beg); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.close(); - sample_id_writer.close(); - diskann::cout << "Wrote " << num_sampled_pts_u32 << " points to sample file: " << output_prefix + "_data.bin" - << std::endl; -} - -// streams data from the file, and samples each vector with probability p_val -// and returns a matrix of size slice_size* ndims as floating point type. -// the slice_size and ndims are set inside the function. - -/*********************************** - * Reimplement using gen_random_slice(const T* inputdata,...) - ************************************/ - -template -void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, - size_t &ndims) -{ - size_t npts; - uint32_t npts32, ndims32; - std::vector> sampled_vectors; - - // amount to read in one shot - size_t read_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file.c_str(), read_blk_size); - - // metadata: npts, ndims - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&ndims32, sizeof(uint32_t)); - npts = npts32; - ndims = ndims32; - - std::unique_ptr cur_vector_T = std::make_unique(ndims); - p_val = p_val < 1 ? p_val : 1; - - std::random_device rd; // Will be used to obtain a seed for the random number - size_t x = rd(); - std::mt19937 generator((uint32_t)x); - std::uniform_real_distribution distribution(0, 1); - - for (size_t i = 0; i < npts; i++) - { - base_reader.read((char *)cur_vector_T.get(), ndims * sizeof(T)); - float rnd_val = distribution(generator); - if (rnd_val < p_val) - { - std::vector cur_vector_float; - for (size_t d = 0; d < ndims; d++) - cur_vector_float.push_back(cur_vector_T[d]); - sampled_vectors.push_back(cur_vector_float); - } - } - slice_size = sampled_vectors.size(); - sampled_data = new float[slice_size * ndims]; - for (size_t i = 0; i < slice_size; i++) - { - for (size_t j = 0; j < ndims; j++) - { - sampled_data[i * ndims + j] = sampled_vectors[i][j]; - } - } -} - -// same as above, but samples from the matrix inputdata instead of a file of -// npts*ndims to return sampled_data of size slice_size*ndims. -template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, - size_t &slice_size) -{ - std::vector> sampled_vectors; - const T *cur_vector_T; - - p_val = p_val < 1 ? p_val : 1; - - std::random_device rd; // Will be used to obtain a seed for the random number engine - size_t x = rd(); - std::mt19937 generator((uint32_t)x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - for (size_t i = 0; i < npts; i++) - { - cur_vector_T = inputdata + ndims * i; - float rnd_val = distribution(generator); - if (rnd_val < p_val) - { - std::vector cur_vector_float; - for (size_t d = 0; d < ndims; d++) - cur_vector_float.push_back(cur_vector_T[d]); - sampled_vectors.push_back(cur_vector_float); - } - } - slice_size = sampled_vectors.size(); - sampled_data = new float[slice_size * ndims]; - for (size_t i = 0; i < slice_size; i++) - { - for (size_t j = 0; j < ndims; j++) - { - sampled_data[i * ndims + j] = sampled_vectors[i][j]; - } - } -} - -int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, - const size_t test_dim, const size_t k_base, std::vector &cluster_sizes) -{ - cluster_sizes.clear(); - - size_t *shard_counts = new size_t[num_centers]; - - for (size_t i = 0; i < num_centers; i++) - { - shard_counts[i] = 0; - } - - size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; - uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; - float *block_data_float; - - size_t num_blocks = DIV_ROUND_UP(num_test, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_test); - size_t cur_blk_size = end_id - start_id; - - block_data_float = test_data_float + start_id * test_dim; - - math_utils::compute_closest_centers(block_data_float, cur_blk_size, test_dim, pivots, num_centers, k_base, - block_closest_centers); - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - shard_counts[shard_id]++; - } - } - } - - diskann::cout << "Estimated cluster sizes: "; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - cluster_sizes.push_back((size_t)cur_shard_count); - diskann::cout << cur_shard_count << " "; - } - diskann::cout << std::endl; - delete[] shard_counts; - delete[] block_closest_centers; - return 0; -} - -template -int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, - const size_t k_base, std::string prefix_path) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) - { - diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; - return -1; - } - - std::unique_ptr shard_counts = std::make_unique(num_centers); - std::vector shard_data_writer(num_centers); - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) - { - std::string data_filename = prefix_path + "_subshard-" + std::to_string(i) + ".bin"; - std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_data_writer[i] = std::ofstream(data_filename.c_str(), std::ios::binary); - shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_data_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer[i].write((char *)&basedim32, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_data_writer[shard_id].write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); - shard_counts[shard_id]++; - } - } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_data_writer[i].seekp(0); - shard_data_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_data_writer[i].close(); - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); - } - - diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " << std::endl; - return 0; -} - -// useful for partitioning large dataset. we first generate only the IDS for -// each shard, and retrieve the actual vectors on demand. -template -int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::string prefix_path) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) - { - diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; - return -1; - } - - std::unique_ptr shard_counts = std::make_unique(num_centers); - - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) - { - std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); - shard_counts[shard_id]++; - } - } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); - } - - diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " << std::endl; - return 0; -} - -template -int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - uint32_t dummy_size = 0; - - std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); - shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); - - uint32_t *shard_ids; - size_t shard_size, tmp; - diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); - - uint32_t cur_pos = 0; - uint32_t num_written = 0; - std::cout << "Shard has " << shard_size << " points" << std::endl; - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - - for (size_t p = 0; p < cur_blk_size; p++) - { - uint32_t original_point_map_id = (uint32_t)(start_id + p); - if (cur_pos == shard_size) - break; - if (original_point_map_id == shard_ids[cur_pos]) - { - cur_pos++; - shard_data_writer.write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); - num_written++; - } - } - if (cur_pos == shard_size) - break; - } - - diskann::cout << "Written file with " << num_written << " points" << std::endl; - - shard_data_writer.seekp(0); - shard_data_writer.write((char *)&num_written, sizeof(uint32_t)); - shard_data_writer.close(); - delete[] shard_ids; - return 0; -} - -// partitions a large base file into many shards using k-means hueristic -// on a random sample generated using sampling_rate probability. After this, it -// assignes each base point to the closest k_base nearest centers and creates -// the shards. -// The total number of points across all shards will be k_base * num_points. - -template -int partition(const std::string data_file, const float sampling_rate, size_t num_parts, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base) -{ - size_t train_dim; - size_t num_train; - float *train_data_float; - - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); - - float *pivot_data; - - std::string cur_file = std::string(prefix_path); - std::string output_file; - - // kmeans_partitioning on training data - - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; - - pivot_data = new float[num_parts * train_dim]; - - // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); - - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); - - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); - - // now pivots are ready. need to stream base points and assign them to - // closest clusters. - - shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - return 0; -} - -template -int partition_with_ram_budget(const std::string data_file, const double sampling_rate, double ram_budget, - size_t graph_degree, const std::string prefix_path, size_t k_base) -{ - size_t train_dim; - size_t num_train; - float *train_data_float; - size_t max_k_means_reps = 10; - - int num_parts = 3; - bool fit_in_ram = false; - - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); - - size_t test_dim; - size_t num_test; - float *test_data_float; - gen_random_slice(data_file, sampling_rate, test_data_float, num_test, test_dim); - - float *pivot_data = nullptr; - - std::string cur_file = std::string(prefix_path); - std::string output_file; - - // kmeans_partitioning on training data - - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; - - while (!fit_in_ram) - { - fit_in_ram = true; - - double max_ram_usage = 0; - if (pivot_data != nullptr) - delete[] pivot_data; - - pivot_data = new float[num_parts * train_dim]; - // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); - - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); - - // now pivots are ready. need to stream base points and assign them to - // closest clusters. - - std::vector cluster_sizes; - estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, k_base, cluster_sizes); - - for (auto &p : cluster_sizes) - { - // to account for the fact that p is the size of the shard over the - // testing sample. - p = (uint64_t)(p / sampling_rate); - double cur_shard_ram_estimate = - diskann::estimate_ram_usage(p, (uint32_t)train_dim, sizeof(T), (uint32_t)graph_degree); - - if (cur_shard_ram_estimate > max_ram_usage) - max_ram_usage = cur_shard_ram_estimate; - } - diskann::cout << "With " << num_parts - << " parts, max estimated RAM usage: " << max_ram_usage / (1024 * 1024 * 1024) - << "GB, budget given is " << ram_budget << std::endl; - if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) - { - fit_in_ram = false; - num_parts += 2; - } - } - - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); - - shard_data_into_clusters_only_ids(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - delete[] test_data_float; - return num_parts; -} - -// Instantations of supported templates - -template void DISKANN_DLLEXPORT gen_random_slice(const std::string base_file, const std::string output_prefix, - double sampling_rate); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string base_file, const std::string output_prefix, - double sampling_rate); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string base_file, const std::string output_prefix, - double sampling_rate); - -template void DISKANN_DLLEXPORT gen_random_slice(const float *inputdata, size_t npts, size_t ndims, double p_val, - float *&sampled_data, size_t &slice_size); -template void DISKANN_DLLEXPORT gen_random_slice(const uint8_t *inputdata, size_t npts, size_t ndims, - double p_val, float *&sampled_data, size_t &slice_size); -template void DISKANN_DLLEXPORT gen_random_slice(const int8_t *inputdata, size_t npts, size_t ndims, - double p_val, float *&sampled_data, size_t &slice_size); - -template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, - size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims); - -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); - -template DISKANN_DLLEXPORT int partition_with_ram_budget(const std::string data_file, - const double sampling_rate, double ram_budget, - size_t graph_degree, const std::string prefix_path, - size_t k_base); -template DISKANN_DLLEXPORT int partition_with_ram_budget(const std::string data_file, - const double sampling_rate, double ram_budget, - size_t graph_degree, const std::string prefix_path, - size_t k_base); -template DISKANN_DLLEXPORT int partition_with_ram_budget(const std::string data_file, const double sampling_rate, - double ram_budget, size_t graph_degree, - const std::string prefix_path, size_t k_base); - -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/pq.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/pq.cpp deleted file mode 100644 index d8fbc7f..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/pq.cpp +++ /dev/null @@ -1,1214 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#ifdef __APPLE__ -#include -#else -#include "mkl.h" -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) -#include "gperftools/malloc_extension.h" -#endif -#endif -#include "pq.h" -#include "partition.h" -#include "math_utils.h" -#include "tsl/robin_map.h" - -// block size for reading/processing large files and matrices in blocks -#define BLOCK_SIZE 5000000 - -namespace diskann -{ - -#ifdef __APPLE__ -typedef long long int MKL_INT; -#endif - -FixedChunkPQTable::FixedChunkPQTable() -{ -} - -FixedChunkPQTable::~FixedChunkPQTable() -{ -#ifndef EXEC_ENV_OLS - if (tables != nullptr) - delete[] tables; - if (tables_tr != nullptr) - delete[] tables_tr; - if (chunk_offsets != nullptr) - delete[] chunk_offsets; - if (centroid != nullptr) - delete[] centroid; - if (rotmat_tr != nullptr) - delete[] rotmat_tr; -#endif -} - -#ifdef EXEC_ENV_OLS -void FixedChunkPQTable::load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks) -{ -#else -void FixedChunkPQTable::load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks) -{ -#endif - - size_t nr, nc; - std::string rotmat_file = std::string(pq_table_file) + "_rotation_matrix.bin"; - -#ifdef EXEC_ENV_OLS - size_t *file_offset_data; // since load_bin only sets the pointer, no need - // to delete. - diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); -#else - std::unique_ptr file_offset_data; - diskann::load_bin(pq_table_file, file_offset_data, nr, nc); -#endif - - bool use_old_filetype = false; - - if (nr != 4 && nr != 5) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 - << " or " << 5; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (nr == 4) - { - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << std::endl; - } - else if (nr == 5) - { - use_old_filetype = true; - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << file_offset_data[4] << std::endl; - } - else - { - throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); - } - -#ifdef EXEC_ENV_OLS - - diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); -#else - diskann::load_bin(pq_table_file, tables, nr, nc, file_offset_data[0]); -#endif - - if ((nr != NUM_PQ_CENTROIDS)) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". file_num_centers = " << nr - << " but expecting " << NUM_PQ_CENTROIDS << " centers"; - throw diskann::ANNException("Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - this->ndims = nc; - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); -#else - diskann::load_bin(pq_table_file, centroid, nr, nc, file_offset_data[1]); -#endif - - if ((nr != this->ndims) || (nc != 1)) - { - diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". file_dim = " << nr - << ", file_cols = " << nc << " but expecting " << this->ndims << " entries in 1 dimension."; - throw diskann::ANNException("Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - int chunk_offsets_index = 2; - if (use_old_filetype) - { - chunk_offsets_index = 3; - } -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#else - diskann::load_bin(pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#endif - - if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) - { - diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr - << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; - throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - this->n_chunks = nr - 1; - diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->ndims - << ", #chunks: " << this->n_chunks << std::endl; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(rotmat_file)) - { - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); -#else - if (file_exists(rotmat_file)) - { - diskann::load_bin(rotmat_file, rotmat_tr, nr, nc); -#endif - if (nr != this->ndims || nc != this->ndims) - { - diskann::cerr << "Error loading rotation matrix file" << std::endl; - throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - use_rotation = true; - } - - // alloc and compute transpose - tables_tr = new float[256 * this->ndims]; - for (size_t i = 0; i < 256; i++) - { - for (size_t j = 0; j < this->ndims; j++) - { - tables_tr[j * 256 + i] = tables[i * this->ndims + j]; - } - } -} - -uint32_t FixedChunkPQTable::get_num_chunks() -{ - return static_cast(n_chunks); -} - -void FixedChunkPQTable::preprocess_query(float *query_vec) -{ - for (uint32_t d = 0; d < ndims; d++) - { - query_vec[d] -= centroid[d]; - } - std::vector tmp(ndims, 0); - if (use_rotation) - { - for (uint32_t d = 0; d < ndims; d++) - { - for (uint32_t d1 = 0; d1 < ndims; d1++) - { - tmp[d] += query_vec[d1] * rotmat_tr[d1 * ndims + d]; - } - } - std::memcpy(query_vec, tmp.data(), ndims * sizeof(float)); - } -} - -// assumes pre-processed query -void FixedChunkPQTable::populate_chunk_distances(const float *query_vec, float *dist_vec) -{ - memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); - // chunk wise distance computation - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - // sum (q-c)^2 for the dimensions associated with this chunk - float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - for (size_t idx = 0; idx < 256; idx++) - { - double diff = centers_dim_vec[idx] - (query_vec[j]); - chunk_dists[idx] += (float)(diff * diff); - } - } - } -} - -float FixedChunkPQTable::l2_distance(const float *query_vec, uint8_t *base_vec) -{ - float res = 0; - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); - res += diff * diff; - } - } - return res; -} - -float FixedChunkPQTable::inner_product(const float *query_vec, uint8_t *base_vec) -{ - float res = 0; - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] * query_vec[j]; // assumes centroid is 0 to - // prevent translation errors - res += diff; - } - } - return -res; // returns negative value to simulate distances (max -> min - // conversion) -} - -// assumes no rotation is involved -void FixedChunkPQTable::inflate_vector(uint8_t *base_vec, float *out_vec) -{ - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - out_vec[j] = centers_dim_vec[base_vec[chunk]] + centroid[j]; - } - } -} - -void FixedChunkPQTable::populate_chunk_inner_products(const float *query_vec, float *dist_vec) -{ - memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); - // chunk wise distance computation - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - // sum (q-c)^2 for the dimensions associated with this chunk - float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - for (size_t idx = 0; idx < 256; idx++) - { - double prod = centers_dim_vec[idx] * query_vec[j]; // assumes that we are not - // shifting the vectors to - // mean zero, i.e., centroid - // array should be all zeros - chunk_dists[idx] -= (float)prod; // returning negative to keep the search code - // clean (max inner product vs min distance) - } - } - } -} - -void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out) -{ - for (size_t i = 0; i < ids.size(); i++) - { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); - } -} - -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - std::vector &dists_out) -{ - //_mm_prefetch((char*) dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - dists_out.clear(); - dists_out.resize(n_pts, 0); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) - { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) - { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) - { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; - } - } -} - -// Need to replace calls to these functions with calls to vector& based -// functions above -void aggregate_coords(const uint32_t *ids, const uint64_t n_ids, const uint8_t *all_coords, const uint64_t ndims, - uint8_t *out) -{ - for (size_t i = 0; i < n_ids; i++) - { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); - } -} - -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - float *dists_out) -{ - _mm_prefetch((char *)dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - memset(dists_out, 0, n_pts * sizeof(float)); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) - { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) - { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) - { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; - } - } -} - -// generate_pq_pivots_simplified is a simplified version of generate_pq_pivots. -// Input is provided in the in-memory buffer train_data. -// Output is stored in the in-memory buffer pivot_data_vector. -// Simplification is based on the following assumptions: -// dim % num_pq_chunks == 0 -// num_centers == 256 by default -// KMEANS_ITERS_FOR_PQ == 15 by default -// make_zero_mean is false by default. -// These assumptions allow to make the function much simpler and avoid storing -// array of chunk_offsets and centroids. -// The compiler pragma for multi-threading support is removed from this implementation -// for the purpose of integration into systems that strictly control resource allocation. -int generate_pq_pivots_simplified(const float *train_data, size_t num_train, size_t dim, size_t num_pq_chunks, - std::vector &pivot_data_vector) -{ - if (num_pq_chunks > dim || dim % num_pq_chunks != 0) - { - return -1; - } - - const size_t num_centers = 256; - const size_t cur_chunk_size = dim / num_pq_chunks; - const uint32_t KMEANS_ITERS_FOR_PQ = 15; - - pivot_data_vector.resize(num_centers * dim); - std::vector cur_pivot_data_vector(num_centers * cur_chunk_size); - std::vector cur_data_vector(num_train * cur_chunk_size); - std::vector closest_center_vector(num_train); - - float *pivot_data = &pivot_data_vector[0]; - float *cur_pivot_data = &cur_pivot_data_vector[0]; - float *cur_data = &cur_data_vector[0]; - uint32_t *closest_center = &closest_center_vector[0]; - - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t chunk_offset = cur_chunk_size * i; - - for (int32_t j = 0; j < num_train; j++) - { - std::memcpy(cur_data + j * cur_chunk_size, train_data + j * dim + chunk_offset, - cur_chunk_size * sizeof(float)); - } - - kmeans::kmeanspp_selecting_pivots(cur_data, num_train, cur_chunk_size, cur_pivot_data, num_centers); - - kmeans::run_lloyds(cur_data, num_train, cur_chunk_size, cur_pivot_data, num_centers, KMEANS_ITERS_FOR_PQ, NULL, - closest_center); - - for (uint64_t j = 0; j < num_centers; j++) - { - std::memcpy(pivot_data + j * dim + chunk_offset, cur_pivot_data + j * cur_chunk_size, - cur_chunk_size * sizeof(float)); - } - } - - return 0; -} - -// given training data in train_data of dimensions num_train * dim, generate -// PQ pivots using k-means algorithm to partition the co-ordinates into -// num_pq_chunks (if it divides dimension, else rounded) chunks, and runs -// k-means in each chunk to compute the PQ pivots and stores in bin format in -// file pq_pivots_path as a s num_centers*dim floating point binary file -int generate_pq_pivots(const float *const passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, uint32_t max_k_means_reps, std::string pq_pivots_path, - bool make_zero_mean) -{ - if (num_pq_chunks > dim) - { - diskann::cout << " Error: number of chunks more than dimension" << std::endl; - return -1; - } - - std::unique_ptr train_data = std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); - - std::unique_ptr full_pivot_data; - - if (file_exists(pq_pivots_path)) - { - size_t file_dim, file_num_centers; - diskann::load_bin(pq_pivots_path, full_pivot_data, file_num_centers, file_dim, METADATA_SIZE); - if (file_dim == dim && file_num_centers == num_centers) - { - diskann::cout << "PQ pivot file exists. Not generating again" << std::endl; - return -1; - } - } - - // Calculate centroid and center the training data - std::unique_ptr centroid = std::make_unique(dim); - for (uint64_t d = 0; d < dim; d++) - { - centroid[d] = 0; - } - if (make_zero_mean) - { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; - } - - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - train_data[p * dim + d] -= centroid[d]; - } - } - } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) - { - if (dim_to_bin.find(d) != dim_to_bin.end()) - continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) - { - cur_best = b; - cur_best_load = bin_loads[b]; - } - } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) - { - cur_num_high++; - if (cur_num_high == max_num_high) - cur_bin_threshold = low_val; - } - } - - chunk_offsets.clear(); - chunk_offsets.push_back(0); - - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); - } - chunk_offsets.push_back(dim); - - full_pivot_data.reset(new float[num_centers * dim]); - - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - - if (cur_chunk_size == 0) - continue; - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); - std::unique_ptr closest_center = std::make_unique(num_train); - - diskann::cout << "Processing chunk " << i << " with dimensions [" << chunk_offsets[i] << ", " - << chunk_offsets[i + 1] << ")" << std::endl; - -#pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)num_train; j++) - { - std::memcpy(cur_data.get() + j * cur_chunk_size, train_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); - } - - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers); - - kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers, - max_k_means_reps, NULL, closest_center.get()); - - for (uint64_t j = 0; j < num_centers; j++) - { - std::memcpy(full_pivot_data.get() + j * dim + chunk_offsets[i], cur_pivot_data.get() + j * cur_chunk_size, - cur_chunk_size * sizeof(float)); - } - } - - std::vector cumul_bytes(4, 0); - cumul_bytes[0] = METADATA_SIZE; - cumul_bytes[1] = cumul_bytes[0] + diskann::save_bin(pq_pivots_path.c_str(), full_pivot_data.get(), - (size_t)num_centers, dim, cumul_bytes[0]); - cumul_bytes[2] = cumul_bytes[1] + - diskann::save_bin(pq_pivots_path.c_str(), centroid.get(), (size_t)dim, 1, cumul_bytes[1]); - cumul_bytes[3] = cumul_bytes[2] + diskann::save_bin(pq_pivots_path.c_str(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(pq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); - - diskann::cout << "Saved pq pivot data to " << pq_pivots_path << " of size " << cumul_bytes[cumul_bytes.size() - 1] - << "B." << std::endl; - - return 0; -} - -int generate_opq_pivots(const float *passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, std::string opq_pivots_path, bool make_zero_mean) -{ - if (num_pq_chunks > dim) - { - diskann::cout << " Error: number of chunks more than dimension" << std::endl; - return -1; - } - - std::unique_ptr train_data = std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); - - std::unique_ptr rotated_train_data = std::make_unique(num_train * dim); - std::unique_ptr rotated_and_quantized_train_data = std::make_unique(num_train * dim); - - std::unique_ptr full_pivot_data; - - // rotation matrix for OPQ - std::unique_ptr rotmat_tr; - - // matrices for SVD - std::unique_ptr Umat = std::make_unique(dim * dim); - std::unique_ptr Vmat_T = std::make_unique(dim * dim); - std::unique_ptr singular_values = std::make_unique(dim); - std::unique_ptr correlation_matrix = std::make_unique(dim * dim); - - // Calculate centroid and center the training data - std::unique_ptr centroid = std::make_unique(dim); - for (uint64_t d = 0; d < dim; d++) - { - centroid[d] = 0; - } - if (make_zero_mean) - { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; - } - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - train_data[p * dim + d] -= centroid[d]; - } - } - } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) - { - if (dim_to_bin.find(d) != dim_to_bin.end()) - continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) - { - cur_best = b; - cur_best_load = bin_loads[b]; - } - } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) - { - cur_num_high++; - if (cur_num_high == max_num_high) - cur_bin_threshold = low_val; - } - } - - chunk_offsets.clear(); - chunk_offsets.push_back(0); - - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); - } - chunk_offsets.push_back(dim); - - full_pivot_data.reset(new float[num_centers * dim]); - rotmat_tr.reset(new float[dim * dim]); - - std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); - for (uint32_t d1 = 0; d1 < dim; d1++) - *(rotmat_tr.get() + d1 * dim + d1) = 1; - - for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) - { - // rotate the training data using the current rotation matrix - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, (MKL_INT)dim, (MKL_INT)dim, 1.0f, - train_data.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, rotated_train_data.get(), - (MKL_INT)dim); - - // compute the PQ pivots on the rotated space - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - - if (cur_chunk_size == 0) - continue; - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); - std::unique_ptr closest_center = std::make_unique(num_train); - - diskann::cout << "Processing chunk " << i << " with dimensions [" << chunk_offsets[i] << ", " - << chunk_offsets[i + 1] << ")" << std::endl; - -#pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)num_train; j++) - { - std::memcpy(cur_data.get() + j * cur_chunk_size, rotated_train_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); - } - - if (rnd == 0) - { - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), - num_centers); - } - else - { - for (uint64_t j = 0; j < num_centers; j++) - { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); - } - } - - uint32_t num_lloyds_iters = 8; - kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers, - num_lloyds_iters, NULL, closest_center.get()); - - for (uint64_t j = 0; j < num_centers; j++) - { - std::memcpy(full_pivot_data.get() + j * dim + chunk_offsets[i], - cur_pivot_data.get() + j * cur_chunk_size, cur_chunk_size * sizeof(float)); - } - - for (size_t j = 0; j < num_train; j++) - { - std::memcpy(rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], - cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, - cur_chunk_size * sizeof(float)); - } - } - - // compute the correlation matrix between the original data and the - // quantized data to compute the new rotation - cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)num_train, 1.0f, - train_data.get(), (MKL_INT)dim, rotated_and_quantized_train_data.get(), (MKL_INT)dim, 0.0f, - correlation_matrix.get(), (MKL_INT)dim); - - // compute the SVD of the correlation matrix to help determine the new - // rotation matrix - -#ifdef __APPLE__ - uint32_t errcode = (uint32_t)LAPACKE_sgesdd(LAPACK_ROW_MAJOR, 'A', (clp_int)dim, (clp_int)dim, - correlation_matrix.get(), (clp_int)dim, singular_values.get(), - Umat.get(), (clp_int)dim, Vmat_T.get(), (clp_int)dim); - -#else - uint32_t errcode = (uint32_t)LAPACKE_sgesdd(LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, - correlation_matrix.get(), (MKL_INT)dim, singular_values.get(), - Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); -#endif - - if (errcode > 0) - { - std::cout << "SVD failed to converge." << std::endl; - exit(-1); - } - - // compute the new rotation matrix from the singular vectors as R^T = U - // V^T - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)dim, 1.0f, - Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim, 0.0f, rotmat_tr.get(), (MKL_INT)dim); - } - - std::vector cumul_bytes(4, 0); - cumul_bytes[0] = METADATA_SIZE; - cumul_bytes[1] = cumul_bytes[0] + diskann::save_bin(opq_pivots_path.c_str(), full_pivot_data.get(), - (size_t)num_centers, dim, cumul_bytes[0]); - cumul_bytes[2] = cumul_bytes[1] + - diskann::save_bin(opq_pivots_path.c_str(), centroid.get(), (size_t)dim, 1, cumul_bytes[1]); - cumul_bytes[3] = cumul_bytes[2] + diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); - - diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " << cumul_bytes[cumul_bytes.size() - 1] - << "B." << std::endl; - - std::string rotmat_path = opq_pivots_path + "_rotation_matrix.bin"; - diskann::save_bin(rotmat_path.c_str(), rotmat_tr.get(), dim, dim); - - return 0; -} - -// generate_pq_data_from_pivots_simplified is a simplified version of generate_pq_data_from_pivots. -// Input is provided in the in-memory buffers data and pivot_data. -// Output is stored in the in-memory buffer pq. -// Simplification is based on the following assumptions: -// supporting only float data type -// dim % num_pq_chunks == 0, which results in a fixed chunk_size -// num_centers == 256 by default -// make_zero_mean is false by default. -// These assumptions allow to make the function much simpler and avoid using -// array of chunk_offsets and centroids. -// The compiler pragma for multi-threading support is removed from this implementation -// for the purpose of integration into systems that strictly control resource allocation. -int generate_pq_data_from_pivots_simplified(const float *data, const size_t num, const float *pivot_data, - const size_t pivots_num, const size_t dim, const size_t num_pq_chunks, - std::vector &pq) -{ - if (num_pq_chunks == 0 || num_pq_chunks > dim || dim % num_pq_chunks != 0) - { - return -1; - } - - const size_t num_centers = 256; - const size_t chunk_size = dim / num_pq_chunks; - - if (pivots_num != num_centers * dim) - { - return -1; - } - - pq.resize(num * num_pq_chunks); - - std::vector cur_pivot_vector(num_centers * chunk_size); - std::vector cur_data_vector(num * chunk_size); - std::vector closest_center_vector(num); - - float *cur_pivot_data = &cur_pivot_vector[0]; - float *cur_data = &cur_data_vector[0]; - uint32_t *closest_center = &closest_center_vector[0]; - - for (size_t i = 0; i < num_pq_chunks; i++) - { - const size_t chunk_offset = chunk_size * i; - - for (int j = 0; j < num_centers; j++) - { - std::memcpy(cur_pivot_data + j * chunk_size, pivot_data + j * dim + chunk_offset, - chunk_size * sizeof(float)); - } - - for (int j = 0; j < num; j++) - { - for (size_t k = 0; k < chunk_size; k++) - { - cur_data[j * chunk_size + k] = data[j * dim + chunk_offset + k]; - } - } - - math_utils::compute_closest_centers(cur_data, num, chunk_size, cur_pivot_data, num_centers, 1, closest_center); - - for (int j = 0; j < num; j++) - { - assert(closest_center[j] < num_centers); - pq[j * num_pq_chunks + i] = closest_center[j]; - } - } - - return 0; -} - -// streams the base file (data_file), and computes the closest centers in each -// chunk to generate the compressed data_file and stores it in -// pq_compressed_vectors_path. -// If the numbber of centers is < 256, it stores as byte vector, else as -// 4-byte vector in binary format. -template -int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, - bool use_opq) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - std::unique_ptr full_pivot_data; - std::unique_ptr rotmat_tr; - std::unique_ptr centroid; - std::unique_ptr chunk_offsets; - - std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; - - if (!file_exists(pq_pivots_path)) - { - std::cout << "ERROR: PQ k-means pivot file not found" << std::endl; - throw diskann::ANNException("PQ k-means pivot file not found", -1); - } - else - { - size_t nr, nc; - std::unique_ptr file_offset_data; - - diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, 0); - - if (nr != 4) - { - diskann::cout << "Error reading pq_pivots file " << pq_pivots_path - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting 4."; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - diskann::load_bin(pq_pivots_path.c_str(), full_pivot_data, nr, nc, file_offset_data[0]); - - if ((nr != num_centers) || (nc != dim)) - { - diskann::cout << "Error reading pq_pivots file " << pq_pivots_path << ". file_num_centers = " << nr - << ", file_dim = " << nc << " but expecting " << num_centers << " centers in " << dim - << " dimensions."; - throw diskann::ANNException("Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, file_offset_data[1]); - - if ((nr != dim) || (nc != 1)) - { - diskann::cout << "Error reading pq_pivots file " << pq_pivots_path << ". file_dim = " << nr - << ", file_cols = " << nc << " but expecting " << dim << " entries in 1 dimension."; - throw diskann::ANNException("Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, file_offset_data[2]); - - if (nr != (uint64_t)num_pq_chunks + 1 || nc != 1) - { - diskann::cout << "Error reading pq_pivots file at chunk offsets; file has nr=" << nr << ",nc=" << nc - << ", expecting nr=" << num_pq_chunks + 1 << ", nc=1." << std::endl; - throw diskann::ANNException("Error reading pq_pivots file at chunk offsets.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (use_opq) - { - std::string rotmat_path = pq_pivots_path + "_rotation_matrix.bin"; - diskann::load_bin(rotmat_path.c_str(), rotmat_tr, nr, nc); - if (nr != (uint64_t)dim || nc != dim) - { - diskann::cout << "Error reading rotation matrix file." << std::endl; - throw diskann::ANNException("Error reading rotation matrix file.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - - diskann::cout << "Loaded PQ pivot information" << std::endl; - } - - std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); - uint32_t num_pq_chunks_u32 = num_pq_chunks; - - compressed_file_writer.write((char *)&num_points, sizeof(uint32_t)); - compressed_file_writer.write((char *)&num_pq_chunks_u32, sizeof(uint32_t)); - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - -#ifdef SAVE_INFLATED_PQ - std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); - inflated_file_writer.write((char *)&num_points, sizeof(uint32_t)); - inflated_file_writer.write((char *)&basedim32, sizeof(uint32_t)); - - std::unique_ptr block_inflated_base = std::make_unique(block_size * dim); - std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); -#endif - - std::unique_ptr block_compressed_base = - std::make_unique(block_size * (size_t)num_pq_chunks); - std::memset(block_compressed_base.get(), 0, block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); - - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - std::unique_ptr block_data_tmp = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)(block_data_T.get()), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_tmp.get(), cur_blk_size, dim); - - diskann::cout << "Processing points [" << start_id << ", " << end_id << ").." << std::flush; - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (uint64_t d = 0; d < dim; d++) - { - block_data_tmp[p * dim + d] -= centroid[d]; - } - } - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (uint64_t d = 0; d < dim; d++) - { - block_data_float[p * dim + d] = block_data_tmp[p * dim + d]; - } - } - - if (use_opq) - { - // rotate the current block with the trained rotation matrix before - // PQ - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, - 1.0f, block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, - block_data_tmp.get(), (MKL_INT)dim); - std::memcpy(block_data_float.get(), block_data_tmp.get(), cur_blk_size * dim * sizeof(float)); - } - - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - if (cur_chunk_size == 0) - continue; - - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(cur_blk_size * cur_chunk_size); - std::unique_ptr closest_center = std::make_unique(cur_blk_size); - -#pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) - { - for (size_t k = 0; k < cur_chunk_size; k++) - cur_data[j * cur_chunk_size + k] = block_data_float[j * dim + chunk_offsets[i] + k]; - } - -#pragma omp parallel for schedule(static, 1) - for (int64_t j = 0; j < (int64_t)num_centers; j++) - { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); - } - - math_utils::compute_closest_centers(cur_data.get(), cur_blk_size, cur_chunk_size, cur_pivot_data.get(), - num_centers, 1, closest_center.get()); - -#pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) - { - block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; -#ifdef SAVE_INFLATED_PQ - for (size_t k = 0; k < cur_chunk_size; k++) - block_inflated_base[j * dim + chunk_offsets[i] + k] = - cur_pivot_data[closest_center[j] * cur_chunk_size + k] + centroid[chunk_offsets[i] + k]; -#endif - } - } - - if (num_centers > 256) - { - compressed_file_writer.write((char *)(block_compressed_base.get()), - cur_blk_size * num_pq_chunks * sizeof(uint32_t)); - } - else - { - std::unique_ptr pVec = std::make_unique(cur_blk_size * num_pq_chunks); - diskann::convert_types(block_compressed_base.get(), pVec.get(), cur_blk_size, - num_pq_chunks); - compressed_file_writer.write((char *)(pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); - } -#ifdef SAVE_INFLATED_PQ - inflated_file_writer.write((char *)(block_inflated_base.get()), cur_blk_size * dim * sizeof(float)); -#endif - diskann::cout << ".done." << std::endl; - } -// Gopal. Splitting diskann_dll into separate DLLs for search and build. -// This code should only be available in the "build" DLL. -#if defined(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); -#endif - compressed_file_writer.close(); -#ifdef SAVE_INFLATED_PQ - inflated_file_writer.close(); -#endif - return 0; -} - -template -void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, - const double p_val, size_t &disk_pq_dims) -{ - size_t train_size, train_dim; - float *train_data; - - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; - - if (disk_pq_dims > train_dim) - disk_pq_dims = train_dim; - - std::cout << "Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, 256, (uint32_t)disk_pq_dims, NUM_KMEANS_REPS_PQ, - disk_pq_pivots_path, false); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - else - generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - - delete[] train_data; -} - -template -void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, diskann::Metric compareMetric, - const double p_val, const uint64_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix) -{ - size_t train_size, train_dim; - float *train_data; - if (!file_exists(codebook_prefix)) - { - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; - - bool make_zero_mean = true; - if (compareMetric == diskann::Metric::INNER_PRODUCT) - make_zero_mean = false; - if (use_opq) // we also do not center the data for OPQ - make_zero_mean = false; - - if (!use_opq) - { - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - NUM_KMEANS_REPS_PQ, pq_pivots_path, make_zero_mean); - } - else - { - generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - pq_pivots_path, make_zero_mean); - } - delete[] train_data; - } - else - { - diskann::cout << "Skip Training with predefined pivots in: " << pq_pivots_path << std::endl; - if (!file_exists(pq_compressed_vectors_path)) - { - diskann::cout << "! Pivot exists, but compressed vectors do not exist, please check the file path" - << std::endl; - diskann::cout << "It's " << pq_compressed_vectors_path << " and " << pq_pivots_path << std::endl; - assert(false); - } - return; - } - generate_pq_data_from_pivots(data_file_to_use, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path, use_opq); -} - -// Instantations of supported templates - -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, - const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data( - const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, - const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const uint64_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const uint64_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const uint64_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_data_store.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/pq_data_store.cpp deleted file mode 100644 index 491975e..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_data_store.cpp +++ /dev/null @@ -1,260 +0,0 @@ -#include - -#include "pq_data_store.h" -#include "pq.h" -#include "pq_scratch.h" -#include "utils.h" -#include "distance.h" - -namespace diskann -{ - -// REFACTOR TODO: Assuming that num_pq_chunks is known already. Must verify if -// this is true. -template -PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, - std::unique_ptr> distance_fn, - std::unique_ptr> pq_distance_fn) - : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), - _distance_metric(distance_fn->get_metric()) -{ - if (num_pq_chunks > dim) - { - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); - } - _distance_fn = std::move(distance_fn); - _pq_distance_fn = std::move(pq_distance_fn); -} - -template PQDataStore::~PQDataStore() -{ - if (_quantized_data != nullptr) - { - aligned_free(_quantized_data); - _quantized_data = nullptr; - } -} - -template location_t PQDataStore::load(const std::string &filename) -{ - return load_impl(filename); -} -template size_t PQDataStore::save(const std::string &filename, const location_t num_points) -{ - return diskann::save_bin(filename, _quantized_data, this->capacity(), _num_chunks, 0); -} - -template size_t PQDataStore::get_aligned_dim() const -{ - return this->get_dims(); -} - -// Populate quantized data from regular data. -template void PQDataStore::populate_data(const data_t *vectors, const location_t num_pts) -{ - throw std::logic_error("Not implemented yet"); -} - -template void PQDataStore::populate_data(const std::string &filename, const size_t offset) -{ - if (_quantized_data != nullptr) - { - aligned_free(_quantized_data); - } - - size_t file_num_points = 0, file_dim = 0; - get_bin_metadata(filename, file_num_points, file_dim, offset); - this->_capacity = (location_t)file_num_points; - this->_dim = file_dim; - - double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); - - auto pivots_file = _pq_distance_fn->get_pivot_data_filename(filename); - auto compressed_file = _pq_distance_fn->get_quantized_vectors_filename(filename); - - generate_quantized_data(filename, pivots_file, compressed_file, _distance_metric, p_val, _num_chunks, - _pq_distance_fn->is_opq()); - - // REFACTOR TODO: Not sure of the alignment. Just copying from index.cpp - alloc_aligned(((void **)&_quantized_data), file_num_points * _num_chunks * sizeof(uint8_t), 1); - copy_aligned_data_from_file(compressed_file.c_str(), _quantized_data, file_num_points, _num_chunks, - _num_chunks); -#ifdef EXEC_ENV_OLS - throw ANNException("load_pq_centroid_bin should not be called when " - "EXEC_ENV_OLS is defined.", - -1, __FUNCSIG__, __FILE__, __LINE__); -#else - _pq_distance_fn->load_pivot_data(pivots_file.c_str(), _num_chunks); -#endif -} - -template -void PQDataStore::extract_data_to_bin(const std::string &filename, const location_t num_pts) -{ - throw std::logic_error("Not implemented yet"); -} - -template void PQDataStore::get_vector(const location_t i, data_t *target) const -{ - // REFACTOR TODO: Should we inflate the compressed vector here? - if (i < this->capacity()) - { - throw std::logic_error("Not implemented yet."); - } - else - { - std::stringstream ss; - ss << "Requested vector " << i << " but only " << this->capacity() << " vectors are present"; - throw diskann::ANNException(ss.str(), -1); - } -} -template void PQDataStore::set_vector(const location_t i, const data_t *const vector) -{ - // REFACTOR TODO: Should we accept a normal vector and compress here? - // memcpy (_data + i * _num_chunks, vector, _num_chunks * sizeof(data_t)); - throw std::logic_error("Not implemented yet"); -} - -template void PQDataStore::prefetch_vector(const location_t loc) -{ - const uint8_t *ptr = _quantized_data + ((size_t)loc) * _num_chunks * sizeof(data_t); - diskann::prefetch_vector((const char *)ptr, _num_chunks * sizeof(data_t)); -} - -template -void PQDataStore::move_vectors(const location_t old_location_start, const location_t new_location_start, - const location_t num_points) -{ - // REFACTOR TODO: Moving vectors is only for in-mem fresh. - throw std::logic_error("Not implemented yet"); -} - -template -void PQDataStore::copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) -{ - // REFACTOR TODO: Is the number of bytes correct? - memcpy(_quantized_data + to_loc * _num_chunks, _quantized_data + from_loc * _num_chunks, _num_chunks * num_points); -} - -// REFACTOR TODO: Currently, we take aligned_query as parameter, but this -// function should also do the alignment. -template -void PQDataStore::preprocess_query(const data_t *aligned_query, AbstractScratch *scratch) const -{ - if (scratch == nullptr) - { - throw diskann::ANNException("Scratch space is null", -1); - } - - PQScratch *pq_scratch = scratch->pq_scratch(); - - if (pq_scratch == nullptr) - { - throw diskann::ANNException("PQScratch space has not been set in the scratch object.", -1); - } - - _pq_distance_fn->preprocess_query(aligned_query, (location_t)this->get_dims(), *pq_scratch); -} - -template float PQDataStore::get_distance(const data_t *query, const location_t loc) const -{ - throw std::logic_error("Not implemented yet"); -} - -template float PQDataStore::get_distance(const location_t loc1, const location_t loc2) const -{ - throw std::logic_error("Not implemented yet"); -} - -template -void PQDataStore::get_distance(const data_t *preprocessed_query, const location_t *locations, - const uint32_t location_count, float *distances, - AbstractScratch *scratch_space) const -{ - if (scratch_space == nullptr) - { - throw diskann::ANNException("Scratch space is null", -1); - } - PQScratch *pq_scratch = scratch_space->pq_scratch(); - if (pq_scratch == nullptr) - { - throw diskann::ANNException("PQScratch not set in scratch space.", -1); - } - diskann::aggregate_coords(locations, location_count, _quantized_data, this->_num_chunks, - pq_scratch->aligned_pq_coord_scratch); - _pq_distance_fn->preprocessed_distance(*pq_scratch, location_count, distances); -} - -template -void PQDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const -{ - if (scratch_space == nullptr) - { - throw diskann::ANNException("Scratch space is null", -1); - } - PQScratch *pq_scratch = scratch_space->pq_scratch(); - if (pq_scratch == nullptr) - { - throw diskann::ANNException("PQScratch not set in scratch space.", -1); - } - diskann::aggregate_coords(ids, _quantized_data, this->_num_chunks, pq_scratch->aligned_pq_coord_scratch); - _pq_distance_fn->preprocessed_distance(*pq_scratch, (location_t)ids.size(), distances); -} - -template location_t PQDataStore::calculate_medoid() const -{ - // REFACTOR TODO: Must calculate this just like we do with data store. - size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); - return (uint32_t)(r % (size_t)this->capacity()); -} - -template size_t PQDataStore::get_alignment_factor() const -{ - return 1; -} - -template Distance *PQDataStore::get_dist_fn() const -{ - return _distance_fn.get(); -} - -template location_t PQDataStore::load_impl(const std::string &file_prefix) -{ - if (_quantized_data != nullptr) - { - aligned_free(_quantized_data); - } - auto quantized_vectors_file = _pq_distance_fn->get_quantized_vectors_filename(file_prefix); - - size_t num_points; - load_aligned_bin(quantized_vectors_file, _quantized_data, num_points, _num_chunks, _num_chunks); - this->_capacity = (location_t)num_points; - - auto pivots_file = _pq_distance_fn->get_pivot_data_filename(file_prefix); - _pq_distance_fn->load_pivot_data(pivots_file, _num_chunks); - - return this->_capacity; -} - -template location_t PQDataStore::expand(const location_t new_size) -{ - throw std::logic_error("Not implemented yet"); -} - -template location_t PQDataStore::shrink(const location_t new_size) -{ - throw std::logic_error("Not implemented yet"); -} - -#ifdef EXEC_ENV_OLS -template location_t PQDataStore::load_impl(AlignedFileReader &reader) -{ -} -#endif - -template DISKANN_DLLEXPORT class PQDataStore; -template DISKANN_DLLEXPORT class PQDataStore; -template DISKANN_DLLEXPORT class PQDataStore; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_flash_index.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/pq_flash_index.cpp deleted file mode 100644 index bfb0abb..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_flash_index.cpp +++ /dev/null @@ -1,2964 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "common_includes.h" - -#include -#include -#include - -#include "timer.h" -#include "pq.h" -#include "pq_scratch.h" -#include "pq_flash_index.h" -#include "cosine_similarity.h" -#include "embedding.pb.h" // from embedding.proto -> embedding.pb.h -#include -#include -#include -#include -#include - -#ifdef _WINDOWS -#include "windows_aligned_file_reader.h" -#elif __APPLE__ -#include "apple_aligned_file_reader.h" -#else -#include "linux_aligned_file_reader.h" -#endif - -#define READ_U64(stream, val) stream.read((char *)&val, sizeof(uint64_t)) -#define READ_U32(stream, val) stream.read((char *)&val, sizeof(uint32_t)) -#define READ_UNSIGNED(stream, val) stream.read((char *)&val, sizeof(unsigned)) - -// sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_NO(id) (((uint64_t)(id)) / _nvecs_per_sector + _reorder_data_start_sector) - -// sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_OFFSET(id) ((((uint64_t)(id)) % _nvecs_per_sector) * _data_dim * sizeof(float)) - -namespace diskann -{ -static std::mutex log_file_mutex; -static std::atomic search_counter(0); - -template -PQFlashIndex::PQFlashIndex(std::shared_ptr &fileReader, - std::shared_ptr &graphReader, diskann::Metric m) - : reader(fileReader), graph_reader(graphReader), metric(m), _thread_data(nullptr) -{ - diskann::Metric metric_to_invoke = m; - if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) - { - if (std::is_floating_point::value) - { - diskann::cout << "Since data is floating point, we assume that it has been appropriately pre-processed " - "(normalization for cosine, and convert-to-l2 by adding extra dimension for MIPS). So we " - "shall invoke an l2 distance function." - << std::endl; - metric_to_invoke = diskann::Metric::L2; - } - else - { - diskann::cerr << "WARNING: Cannot normalize integral data types." - << " This may result in erroneous results or poor recall." - << " Consider using L2 distance with integral data types." << std::endl; - } - } - - this->_dist_cmp.reset(diskann::get_distance_function(metric_to_invoke)); - this->_dist_cmp_float.reset(diskann::get_distance_function(metric_to_invoke)); -} - -template PQFlashIndex::~PQFlashIndex() -{ -#ifndef EXEC_ENV_OLS - if (data != nullptr) - { - delete[] data; - } -#endif - - if (_centroid_data != nullptr) - aligned_free(_centroid_data); - // delete backing bufs for nhood and coord cache - if (_nhood_cache_buf != nullptr) - { - delete[] _nhood_cache_buf; - diskann::aligned_free(_coord_cache_buf); - } - - if (_load_flag) - { - diskann::cout << "Clearing scratch" << std::endl; - ScratchStoreManager> manager(this->_thread_data); - manager.destroy(); - this->reader->deregister_all_threads(); - reader->close(); - } - if (_pts_to_label_offsets != nullptr) - { - delete[] _pts_to_label_offsets; - } - if (_pts_to_label_counts != nullptr) - { - delete[] _pts_to_label_counts; - } - if (_pts_to_labels != nullptr) - { - delete[] _pts_to_labels; - } - if (_medoids != nullptr) - { - delete[] _medoids; - } -} - -template inline uint64_t PQFlashIndex::get_node_sector(uint64_t node_id) -{ - return 1 + (_nnodes_per_sector > 0 ? node_id / _nnodes_per_sector - : node_id * DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN)); -} - -template -inline char *PQFlashIndex::offset_to_node(char *sector_buf, uint64_t node_id) -{ - return sector_buf + (_nnodes_per_sector == 0 ? 0 : (node_id % _nnodes_per_sector) * _max_node_len); -} - -template inline uint32_t *PQFlashIndex::offset_to_node_nhood(char *node_buf) -{ - return (unsigned *)(node_buf + _disk_bytes_per_point); -} - -template inline T *PQFlashIndex::offset_to_node_coords(char *node_buf) -{ - return (T *)(node_buf); -} - -template -void PQFlashIndex::setup_thread_data(uint64_t nthreads, uint64_t visited_reserve) -{ - diskann::cout << "Setting up thread-specific contexts for nthreads: " << nthreads << std::endl; -// omp parallel for to generate unique thread IDs -#pragma omp parallel for num_threads((int)nthreads) - for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) - { -#pragma omp critical - { - SSDThreadData *data = new SSDThreadData(this->_aligned_dim, visited_reserve); - this->reader->register_thread(); - data->ctx = this->reader->get_ctx(); - this->_thread_data.push(data); - } - } - _load_flag = true; -} - -template -std::vector PQFlashIndex::read_nodes(const std::vector &node_ids, - std::vector &coord_buffers, - std::vector> &nbr_buffers) -{ - std::vector read_reqs; - std::vector retval(node_ids.size(), true); - - char *buf = nullptr; - auto num_sectors = _nnodes_per_sector > 0 ? 1 : DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN); - - // borrow thread data and issue reads - ScratchStoreManager> manager(this->_thread_data); - auto this_thread_data = manager.scratch_space(); - IOContext &ctx = this_thread_data->ctx; - -#if 1 - // -- If not partition_read, this is the normal DiskANN approach: - if (!_use_partition) - { -#endif - // (1) read each node's 4 KB from offset = get_node_sector(node_id)*4096 - alloc_aligned((void **)&buf, node_ids.size() * num_sectors * defaults::SECTOR_LEN, defaults::SECTOR_LEN); - - // create read requests - for (size_t i = 0; i < node_ids.size(); ++i) - { - auto node_id = node_ids[i]; - - AlignedRead read; - read.len = num_sectors * defaults::SECTOR_LEN; - read.buf = buf + i * num_sectors * defaults::SECTOR_LEN; - read.offset = get_node_sector(node_id) * defaults::SECTOR_LEN; - read_reqs.push_back(read); - } - - reader->read(read_reqs, ctx); - - // copy reads into buffers - for (uint32_t i = 0; i < read_reqs.size(); i++) - { -#if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle failed reads in - // production settings - if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) - { - retval[i] = false; - continue; - } -#endif - - char *node_buf = offset_to_node((char *)read_reqs[i].buf, node_ids[i]); - - if (coord_buffers[i] != nullptr) - { - T *node_coords = offset_to_node_coords(node_buf); - memcpy(coord_buffers[i], node_coords, _disk_bytes_per_point); - } - - if (nbr_buffers[i].second != nullptr) - { - uint32_t *node_nhood = offset_to_node_nhood(node_buf); - auto num_nbrs = *node_nhood; - nbr_buffers[i].first = num_nbrs; - memcpy(nbr_buffers[i].second, node_nhood + 1, num_nbrs * sizeof(uint32_t)); - } - } - aligned_free(buf); - - if (!_use_partition) - { - // done with the normal path - return retval; - } -#if 1 - } -#endif - - { - // 计算每个节点的分区偏移 - std::vector> offsets(node_ids.size()); - std::vector valid_nodes(node_ids.size(), true); - - // 按分区分组，减少重复读取相同分区 - std::map> partition_to_indices; - - // 遍历所有节点，获取其分区信息 - for (size_t i = 0; i < node_ids.size(); i++) - { - uint32_t node_id = node_ids[i]; - if (nbr_buffers[i].second != nullptr) - { - // 使用read_neighbors中的逻辑获取分区ID - uint32_t partition_id = _id2partition[node_id]; - if (partition_id >= _num_partitions) - { - valid_nodes[i] = false; - retval[i] = false; - continue; - } - - // 将节点按分区ID分组 - partition_to_indices[partition_id].push_back(i); - } - } - - // 对每个分区执行一次读取 - for (const auto &pair : partition_to_indices) - { - uint32_t partition_id = pair.first; - const auto &indices = pair.second; - - // 计算扇区偏移 (与read_neighbors中相同) - uint64_t sector_offset = (partition_id + 1) * defaults::SECTOR_LEN; - - // 读取分区扇区 - char *sector_buf = nullptr; - alloc_aligned((void **)§or_buf, defaults::SECTOR_LEN, defaults::SECTOR_LEN); - - AlignedRead read; - read.len = defaults::SECTOR_LEN; - read.buf = sector_buf; - read.offset = sector_offset; - - std::vector single_read = {read}; - graph_reader->read(single_read, ctx); - - // 处理该分区中的所有节点 - for (size_t idx : indices) - { - uint32_t node_id = node_ids[idx]; - - // 查找节点在分区内的位置 (与read_neighbors中相同) - const auto &part_list = _graph_partitions[partition_id]; - auto it = std::find(part_list.begin(), part_list.end(), node_id); - if (it == part_list.end()) - { - retval[idx] = false; - continue; - } - size_t j = std::distance(part_list.begin(), it); - - // 计算节点在扇区内的偏移 (与read_neighbors中相同) - uint64_t node_offset = j * _graph_node_len; - if (node_offset + 4 > defaults::SECTOR_LEN) - { - retval[idx] = false; - continue; - } - - // 读取邻居数量 - char *adjacency_ptr = sector_buf + node_offset; - uint32_t neighbor_count = *reinterpret_cast(adjacency_ptr); - - // 检查邻居数据是否超出扇区范围 - size_t needed = neighbor_count * sizeof(uint32_t); - if (node_offset + 4 + needed > defaults::SECTOR_LEN) - { - retval[idx] = false; - continue; - } - - // 拷贝邻居数据 - nbr_buffers[idx].first = neighbor_count; - memcpy(nbr_buffers[idx].second, adjacency_ptr + 4, needed); - } - - aligned_free(sector_buf); - } - } - - return retval; -} - -template void PQFlashIndex::load_cache_list(std::vector &node_list) -{ - diskann::cout << "Loading the cache list into memory.." << std::flush; - size_t num_cached_nodes = node_list.size(); - - // Allocate space for neighborhood cache - _nhood_cache_buf = new uint32_t[num_cached_nodes * (_max_degree + 1)]; - memset(_nhood_cache_buf, 0, num_cached_nodes * (_max_degree + 1)); - - // Allocate space for coordinate cache - size_t coord_cache_buf_len = num_cached_nodes * _aligned_dim; - diskann::alloc_aligned((void **)&_coord_cache_buf, coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); - memset(_coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); - - size_t BLOCK_SIZE = 8; - size_t num_blocks = 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); - - // Copy offset into buffers to read into - std::vector nodes_to_read; - std::vector coord_buffers; - std::vector> nbr_buffers; - for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) - { - nodes_to_read.push_back(node_list[node_idx]); - coord_buffers.push_back(_coord_cache_buf + node_idx * _aligned_dim); - nbr_buffers.emplace_back(0, _nhood_cache_buf + node_idx * (_max_degree + 1)); - } - - // issue the reads - auto read_status = read_nodes(nodes_to_read, coord_buffers, nbr_buffers); - - // check for success and insert into the cache. - 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])); - _nhood_cache.insert(std::make_pair(nodes_to_read[i], nbr_buffers[i])); - } - } - } - diskann::cout << "..done." << std::endl; -} - -#ifdef EXEC_ENV_OLS -template -void PQFlashIndex::generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, - uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list) -{ -#else -template -void PQFlashIndex::generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, - uint32_t nthreads, - std::vector &node_list) -{ -#endif - if (num_nodes_to_cache >= this->_num_points) - { - // for small num_points and big num_nodes_to_cache, use below way to get the node_list quickly - node_list.resize(this->_num_points); - for (uint32_t i = 0; i < this->_num_points; ++i) - { - node_list[i] = i; - } - return; - } - - this->_count_visited_nodes = true; - this->_node_visit_counter.clear(); - this->_node_visit_counter.resize(this->_num_points); - for (uint32_t i = 0; i < _node_visit_counter.size(); i++) - { - this->_node_visit_counter[i].first = i; - this->_node_visit_counter[i].second = 0; - } - - size_t sample_num, sample_dim, sample_aligned_dim; - T *samples; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(sample_bin)) - { - diskann::load_aligned_bin(files, sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); - } -#else - if (file_exists(sample_bin)) - { - diskann::load_aligned_bin(sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); - } -#endif - else - { - diskann::cerr << "Sample bin file not found. Not generating cache." << std::endl; - return; - } - - std::vector tmp_result_ids_64(sample_num, 0); - std::vector tmp_result_dists(sample_num, 0); - - bool filtered_search = false; - std::vector random_query_filters(sample_num); - if (_filter_to_medoid_ids.size() != 0) - { - filtered_search = true; - generate_random_labels(random_query_filters, (uint32_t)sample_num, nthreads); - } - -#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (int64_t i = 0; i < (int64_t)sample_num; i++) - { - auto &label_for_search = random_query_filters[i]; - // run a search on the sample query with a random label (sampled from base label distribution), and it will - // concurrently update the node_visit_counter to track most visited nodes. The last false is to not use the - // "use_reorder_data" option which enables a final reranking if the disk index itself contains only PQ data. - cached_beam_search(samples + (i * sample_aligned_dim), 1, l_search, tmp_result_ids_64.data() + i, - tmp_result_dists.data() + i, beamwidth, filtered_search, label_for_search, false); - } - - std::sort(this->_node_visit_counter.begin(), _node_visit_counter.end(), - [](std::pair &left, std::pair &right) { - return left.second > right.second; - }); - node_list.clear(); - node_list.shrink_to_fit(); - num_nodes_to_cache = std::min((size_t)num_nodes_to_cache, this->_node_visit_counter.size()); - node_list.reserve(num_nodes_to_cache); - for (uint64_t i = 0; i < num_nodes_to_cache; i++) - { - node_list.push_back(this->_node_visit_counter[i].first); - } - this->_count_visited_nodes = false; - - diskann::aligned_free(samples); -} - -template -void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, - const bool shuffle) -{ - std::random_device rng; - std::mt19937 urng(rng()); - - tsl::robin_set node_set; - - // Do not cache more than 10% of the nodes in the index - uint64_t tenp_nodes = (uint64_t)(std::round(this->_num_points * 0.1)); - if (num_nodes_to_cache > tenp_nodes) - { - diskann::cout << "Reducing nodes to cache from: " << num_nodes_to_cache << " to: " << tenp_nodes - << "(10 percent of total nodes:" << this->_num_points << ")" << std::endl; - num_nodes_to_cache = tenp_nodes == 0 ? 1 : tenp_nodes; - } - diskann::cout << "Caching " << num_nodes_to_cache << "..." << std::endl; - - std::unique_ptr> cur_level, prev_level; - cur_level = std::make_unique>(); - prev_level = std::make_unique>(); - - for (uint64_t miter = 0; miter < _num_medoids && cur_level->size() < num_nodes_to_cache; miter++) - { - cur_level->insert(_medoids[miter]); - } - - if ((_filter_to_medoid_ids.size() > 0) && (cur_level->size() < num_nodes_to_cache)) - { - for (auto &x : _filter_to_medoid_ids) - { - for (auto &y : x.second) - { - cur_level->insert(y); - if (cur_level->size() == num_nodes_to_cache) - break; - } - if (cur_level->size() == num_nodes_to_cache) - break; - } - } - - uint64_t lvl = 1; - uint64_t prev_node_set_size = 0; - while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && cur_level->size() != 0) - { - // swap prev_level and cur_level - std::swap(prev_level, cur_level); - // clear cur_level - cur_level->clear(); - - std::vector nodes_to_expand; - - for (const uint32_t &id : *prev_level) - { - if (node_set.find(id) != node_set.end()) - { - continue; - } - node_set.insert(id); - nodes_to_expand.push_back(id); - } - - if (shuffle) - std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); - else - std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); - - diskann::cout << "Level: " << lvl << std::flush; - bool finish_flag = false; - - size_t BLOCK_SIZE = 1024; - size_t nblocks = DIV_ROUND_UP(nodes_to_expand.size(), BLOCK_SIZE); - for (size_t block = 0; block < nblocks && !finish_flag; block++) - { - diskann::cout << "." << std::flush; - size_t start = block * BLOCK_SIZE; - size_t end = (std::min)((block + 1) * BLOCK_SIZE, nodes_to_expand.size()); - - std::vector nodes_to_read; - std::vector coord_buffers(end - start, nullptr); - std::vector> nbr_buffers; - - for (size_t cur_pt = start; cur_pt < end; cur_pt++) - { - nodes_to_read.push_back(nodes_to_expand[cur_pt]); - nbr_buffers.emplace_back(0, new uint32_t[_max_degree + 1]); - } - - // issue read requests - auto read_status = read_nodes(nodes_to_read, coord_buffers, nbr_buffers); - - // process each nhood buf - for (uint32_t i = 0; i < read_status.size(); i++) - { - if (read_status[i] == false) - { - continue; - } - else - { - uint32_t nnbrs = nbr_buffers[i].first; - uint32_t *nbrs = nbr_buffers[i].second; - - // explore next level - for (uint32_t j = 0; j < nnbrs && !finish_flag; j++) - { - if (node_set.find(nbrs[j]) == node_set.end()) - { - cur_level->insert(nbrs[j]); - } - if (cur_level->size() + node_set.size() >= num_nodes_to_cache) - { - finish_flag = true; - } - } - } - delete[] nbr_buffers[i].second; - } - } - - diskann::cout << ". #nodes: " << node_set.size() - prev_node_set_size - << ", #nodes thus far: " << node_set.size() << std::endl; - prev_node_set_size = node_set.size(); - lvl++; - } - - 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); - for (auto node : *cur_level) - node_list.push_back(node); - - diskann::cout << "Level: " << lvl << std::flush; - diskann::cout << ". #nodes: " << node_list.size() - prev_node_set_size << ", #nodes thus far: " << node_list.size() - << std::endl; - diskann::cout << "done" << std::endl; -} - -template void PQFlashIndex::use_medoids_data_as_centroids() -{ - if (_centroid_data != nullptr) - aligned_free(_centroid_data); - alloc_aligned(((void **)&_centroid_data), _num_medoids * _aligned_dim * sizeof(float), 32); - std::memset(_centroid_data, 0, _num_medoids * _aligned_dim * sizeof(float)); - - diskann::cout << "Loading centroid data from medoids vector data of " << _num_medoids << " medoid(s)" << std::endl; - - std::vector nodes_to_read; - std::vector medoid_bufs; - std::vector> nbr_bufs; - - for (uint64_t cur_m = 0; cur_m < _num_medoids; cur_m++) - { - nodes_to_read.push_back(_medoids[cur_m]); - medoid_bufs.push_back(new T[_data_dim]); - nbr_bufs.emplace_back(0, nullptr); - } - - auto read_status = read_nodes(nodes_to_read, medoid_bufs, nbr_bufs); - - for (uint64_t cur_m = 0; cur_m < _num_medoids; cur_m++) - { - if (read_status[cur_m] == true) - { - if (!_use_disk_index_pq) - { - for (uint32_t i = 0; i < _data_dim; i++) - _centroid_data[cur_m * _aligned_dim + i] = medoid_bufs[cur_m][i]; - } - else - { - _disk_pq_table.inflate_vector((uint8_t *)medoid_bufs[cur_m], (_centroid_data + cur_m * _aligned_dim)); - } - } - else - { - throw ANNException("Unable to read a medoid", -1, __FUNCSIG__, __FILE__, __LINE__); - } - delete[] medoid_bufs[cur_m]; - } -} - -template -void PQFlashIndex::generate_random_labels(std::vector &labels, const uint32_t num_labels, - const uint32_t nthreads) -{ - std::random_device rd; - labels.clear(); - labels.resize(num_labels); - - uint64_t num_total_labels = _pts_to_label_offsets[_num_points - 1] + _pts_to_label_counts[_num_points - 1]; - std::mt19937 gen(rd()); - if (num_total_labels == 0) - { - std::stringstream stream; - stream << "No labels found in data. Not sampling random labels "; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - std::uniform_int_distribution dis(0, num_total_labels - 1); - -#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (int64_t i = 0; i < num_labels; i++) - { - uint64_t rnd_loc = dis(gen); - labels[i] = (LabelT)_pts_to_labels[rnd_loc]; - } -} - -template -std::unordered_map PQFlashIndex::load_label_map(std::basic_istream &map_reader) -{ - std::unordered_map string_to_int_mp; - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) - { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = (LabelT)std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; -} - -template -LabelT PQFlashIndex::get_converted_label(const std::string &filter_label) -{ - if (_label_map.find(filter_label) != _label_map.end()) - { - return _label_map[filter_label]; - } - if (_use_universal_label) - { - return _universal_filter_label; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); -} - -template -void PQFlashIndex::reset_stream_for_reading(std::basic_istream &infile) -{ - infile.clear(); - infile.seekg(0); -} - -template -void PQFlashIndex::get_label_file_metadata(const std::string &fileContent, uint32_t &num_pts, - uint32_t &num_total_labels) -{ - num_pts = 0; - num_total_labels = 0; - - size_t file_size = fileContent.length(); - - std::string label_str; - size_t cur_pos = 0; - size_t next_pos = 0; - while (cur_pos < file_size && cur_pos != std::string::npos) - { - next_pos = fileContent.find('\n', cur_pos); - if (next_pos == std::string::npos) - { - break; - } - - size_t lbl_pos = cur_pos; - size_t next_lbl_pos = 0; - while (lbl_pos < next_pos && lbl_pos != std::string::npos) - { - next_lbl_pos = fileContent.find(',', lbl_pos); - if (next_lbl_pos == std::string::npos) // the last label - { - next_lbl_pos = next_pos; - } - - num_total_labels++; - - lbl_pos = next_lbl_pos + 1; - } - - cur_pos = next_pos + 1; - - num_pts++; - } - - diskann::cout << "Labels file metadata: num_points: " << num_pts << ", #total_labels: " << num_total_labels - << std::endl; -} - -template -inline bool PQFlashIndex::point_has_label(uint32_t point_id, LabelT label_id) -{ - uint32_t start_vec = _pts_to_label_offsets[point_id]; - uint32_t num_lbls = _pts_to_label_counts[point_id]; - bool ret_val = false; - for (uint32_t i = 0; i < num_lbls; i++) - { - if (_pts_to_labels[start_vec + i] == label_id) - { - ret_val = true; - break; - } - } - return ret_val; -} - -template -void PQFlashIndex::parse_label_file(std::basic_istream &infile, size_t &num_points_labels) -{ - infile.seekg(0, std::ios::end); - size_t file_size = infile.tellg(); - - std::string buffer(file_size, ' '); - - infile.seekg(0, std::ios::beg); - infile.read(&buffer[0], file_size); - - std::string line; - uint32_t line_cnt = 0; - - uint32_t num_pts_in_label_file; - uint32_t num_total_labels; - get_label_file_metadata(buffer, num_pts_in_label_file, num_total_labels); - - _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; - _pts_to_label_counts = new uint32_t[num_pts_in_label_file]; - _pts_to_labels = new LabelT[num_total_labels]; - uint32_t labels_seen_so_far = 0; - - std::string label_str; - size_t cur_pos = 0; - size_t next_pos = 0; - while (cur_pos < file_size && cur_pos != std::string::npos) - { - next_pos = buffer.find('\n', cur_pos); - if (next_pos == std::string::npos) - { - break; - } - - _pts_to_label_offsets[line_cnt] = labels_seen_so_far; - uint32_t &num_lbls_in_cur_pt = _pts_to_label_counts[line_cnt]; - num_lbls_in_cur_pt = 0; - - size_t lbl_pos = cur_pos; - size_t next_lbl_pos = 0; - while (lbl_pos < next_pos && lbl_pos != std::string::npos) - { - next_lbl_pos = buffer.find(',', lbl_pos); - if (next_lbl_pos == std::string::npos) // the last label in the whole file - { - next_lbl_pos = next_pos; - } - - if (next_lbl_pos > next_pos) // the last label in one line, just read to the end - { - next_lbl_pos = next_pos; - } - - label_str.assign(buffer.c_str() + lbl_pos, next_lbl_pos - lbl_pos); - if (label_str[label_str.length() - 1] == '\t') // '\t' won't exist in label file? - { - label_str.erase(label_str.length() - 1); - } - - LabelT token_as_num = (LabelT)std::stoul(label_str); - _pts_to_labels[labels_seen_so_far++] = (LabelT)token_as_num; - num_lbls_in_cur_pt++; - - // move to next label - lbl_pos = next_lbl_pos + 1; - } - - // move to next line - cur_pos = next_pos + 1; - - if (num_lbls_in_cur_pt == 0) - { - diskann::cout << "No label found for point " << line_cnt << std::endl; - exit(-1); - } - - line_cnt++; - } - - num_points_labels = line_cnt; - reset_stream_for_reading(infile); -} - -template void PQFlashIndex::set_universal_label(const LabelT &label) -{ - _use_universal_label = true; - _universal_filter_label = label; -} - -#ifdef EXEC_ENV_OLS -template -int PQFlashIndex::load(MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix, - const char *pq_prefix) -{ -#else -template -int PQFlashIndex::load(uint32_t num_threads, const char *index_prefix, const char *pq_prefix, - const char *partition_prefix) -{ -#endif - if (pq_prefix == nullptr || strcmp(pq_prefix, "") == 0) - { - pq_prefix = index_prefix; - } - if (partition_prefix != nullptr && strcmp(partition_prefix, "") != 0) - { - _use_partition = true; - } - std::string pq_table_bin = std::string(pq_prefix) + "_pq_pivots.bin"; - std::string pq_compressed_vectors = std::string(pq_prefix) + "_pq_compressed.bin"; - std::string _disk_index_file = std::string(index_prefix) + "_disk.index"; - std::string graph_file = std::string(partition_prefix) + "_disk_graph.index"; - std::string partition_file = std::string(partition_prefix) + "_partition.bin"; -#ifdef EXEC_ENV_OLS - return load_from_separate_paths(files, num_threads, _disk_index_file.c_str(), pq_table_bin.c_str(), - pq_compressed_vectors.c_str(), graph_file.c_str(), partition_file.c_str()); -#else - return load_from_separate_paths(num_threads, _disk_index_file.c_str(), pq_table_bin.c_str(), - pq_compressed_vectors.c_str(), graph_file.c_str(), partition_file.c_str()); -#endif -} - -template -int PQFlashIndex::read_partition_info(const std::string &partition_bin) -{ - std::ifstream pf(partition_bin, std::ios::binary); - if (!pf.is_open()) - { - diskann::cout << "Cannot open partition.bin: " << partition_bin << std::endl; - return 1; - } - diskann::cout << "Loading partition info from " << partition_bin << std::endl; - uint64_t C, nd; - READ_U64(pf, C); - READ_U64(pf, _num_partitions); - READ_U64(pf, nd); - std::cout << "[partition.bin header] C=" << C << ", partition_nums=" << _num_partitions << ", nd=" << nd - << std::endl; - - // 读取分区节点列表 - _graph_partitions.resize(_num_partitions); - for (uint64_t i = 0; i < _num_partitions; i++) - { - uint32_t psize; - READ_U32(pf, psize); - _graph_partitions[i].resize(psize); - pf.read(reinterpret_cast(_graph_partitions[i].data()), psize * sizeof(uint32_t)); - } - // 读取 _id2partition[node], 大小= nd - _id2partition.resize(nd); - pf.read(reinterpret_cast(_id2partition.data()), nd * sizeof(uint32_t)); - pf.close(); - std::cout << "Done loading partition info.\n"; - - return 0; -} - -template -int PQFlashIndex::load_graph_index(const std::string &graph_index_file) -{ - std::ifstream gf(graph_index_file, std::ios::binary); - if (!gf.is_open()) - { - diskann::cout << "Cannot open disk_graph.index: " << graph_index_file << std::endl; - return 1; - } - diskann::cout << "Loading graph index from " << graph_index_file << std::endl; - - // (a) sector0 => read 2 ints for meta_n and meta_dim - int meta_n, meta_dim; - gf.read((char *)&meta_n, sizeof(int)); - gf.read((char *)&meta_dim, sizeof(int)); - diskann::cout << "[debug] meta_n=" << meta_n << ", meta_dim=" << meta_dim << "\n"; - - // (b) Read uint64_t meta_n times - std::vector meta_info(meta_n); - gf.read(reinterpret_cast(meta_info.data()), meta_n * sizeof(uint64_t)); - for (int i = 0; i < meta_n; i++) - { - diskann::cout << " meta_info[" << i << "]= " << meta_info[i] << "\n"; - } - - size_t file_size = get_file_size(graph_index_file); - diskann::cout << "[disk_graph.index size] " << file_size << " bytes\n"; - - uint64_t nd_in_meta = meta_info[0]; - uint64_t dim_in_meta = meta_info[1]; - uint64_t max_node_len = meta_info[3]; - uint64_t c_in_meta = meta_info[4]; - uint64_t entire_file_sz = meta_info[8]; - - diskann::cout << "Based on meta_info:\n" - << " nd_in_meta= " << nd_in_meta << ", dim_in_meta= " << dim_in_meta - << ", max_node_len= " << max_node_len << ", c_in_meta= " << c_in_meta - << ", entire_file_size= " << entire_file_sz << "\n"; - - uint64_t dim_size = dim_in_meta * sizeof(float); - - _graph_node_len = max_node_len - dim_size; - -#if 0 - assert(max_node_len == _max_node_len); - assert(dim_size == _disk_bytes_per_point); - assert(_graph_node_len / sizeof(float) == _max_degree + 1); -#endif - - // Compensate the losting info from old meta_info - _max_degree = _graph_node_len / sizeof(float) - 1; - _disk_bytes_per_point = dim_size; - _max_node_len = max_node_len; - - diskann::cout << " => graph_node_len= " << _graph_node_len << "\n\n"; - - return 0; -} - -#ifdef EXEC_ENV_OLS -template -int PQFlashIndex::load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath, const char *graph_filepath) -{ -#else -template -int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath, - const char *graph_file, const char *partition_file) -{ -#endif - std::string pq_table_bin = pivots_filepath; - std::string pq_compressed_vectors = compressed_filepath; - std::string _disk_index_file = index_filepath; - // medoids, etc. - std::string medoids_file = std::string(_disk_index_file) + "_medoids.bin"; - std::string centroids_file = std::string(_disk_index_file) + "_centroids.bin"; - - std::string labels_file = std::string(_disk_index_file) + "_labels.txt"; - std::string labels_to_medoids = std::string(_disk_index_file) + "_labels_to_medoids.txt"; - std::string dummy_map_file = std::string(_disk_index_file) + "_dummy_map.txt"; - std::string labels_map_file = std::string(_disk_index_file) + "_labels_map.txt"; - - size_t num_pts_in_label_file = 0; - - size_t pq_file_dim = 0, pq_file_num_centroids = 0; -#ifdef EXEC_ENV_OLS - get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); -#else - get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); -#endif - - this->_disk_index_file = _disk_index_file; - - if (pq_file_num_centroids != 256) - { - diskann::cout << "Got " << pq_file_num_centroids << " PQ centroids, loading from " << pq_table_bin << std::endl; - diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." << std::endl; - return -1; - } - - this->_data_dim = pq_file_dim; - // will change later if we use PQ on disk or if we are using - // inner product without PQ - this->_disk_bytes_per_point = this->_data_dim * sizeof(T); - this->_aligned_dim = ROUND_UP(pq_file_dim, 8); - - size_t npts_u64, nchunks_u64; -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, nchunks_u64); -#else - diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, nchunks_u64); -#endif - - this->_num_points = npts_u64; - this->_n_chunks = nchunks_u64; -#ifdef EXEC_ENV_OLS - if (files.fileExists(labels_file)) - { - FileContent &content_labels = files.getContent(labels_file); - std::stringstream infile(std::string((const char *)content_labels._content, content_labels._size)); -#else - if (file_exists(labels_file)) - { - std::ifstream infile(labels_file, std::ios::binary); - if (infile.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + labels_file, -1); - } -#endif - parse_label_file(infile, num_pts_in_label_file); - assert(num_pts_in_label_file == this->_num_points); - -#ifndef EXEC_ENV_OLS - infile.close(); -#endif - -#ifdef EXEC_ENV_OLS - FileContent &content_labels_map = files.getContent(labels_map_file); - std::stringstream map_reader(std::string((const char *)content_labels_map._content, content_labels_map._size)); -#else - std::ifstream map_reader(labels_map_file); -#endif - _label_map = load_label_map(map_reader); - -#ifndef EXEC_ENV_OLS - map_reader.close(); -#endif - -#ifdef EXEC_ENV_OLS - if (files.fileExists(labels_to_medoids)) - { - FileContent &content_labels_to_meoids = files.getContent(labels_to_medoids); - std::stringstream medoid_stream( - std::string((const char *)content_labels_to_meoids._content, content_labels_to_meoids._size)); -#else - if (file_exists(labels_to_medoids)) - { - std::ifstream medoid_stream(labels_to_medoids); - assert(medoid_stream.is_open()); -#endif - std::string line, token; - - _filter_to_medoid_ids.clear(); - try - { - while (std::getline(medoid_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - std::vector medoids; - LabelT label; - while (std::getline(iss, token, ',')) - { - if (cnt == 0) - label = (LabelT)std::stoul(token); - else - medoids.push_back((uint32_t)stoul(token)); - cnt++; - } - _filter_to_medoid_ids[label].swap(medoids); - } - } - catch (std::system_error &e) - { - throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - std::string univ_label_file = std ::string(_disk_index_file) + "_universal_label.txt"; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(univ_label_file)) - { - FileContent &content_univ_label = files.getContent(univ_label_file); - std::stringstream universal_label_reader( - std::string((const char *)content_univ_label._content, content_univ_label._size)); -#else - if (file_exists(univ_label_file)) - { - std::ifstream universal_label_reader(univ_label_file); - assert(universal_label_reader.is_open()); -#endif - std::string univ_label; - universal_label_reader >> univ_label; -#ifndef EXEC_ENV_OLS - universal_label_reader.close(); -#endif - LabelT label_as_num = (LabelT)std::stoul(univ_label); - set_universal_label(label_as_num); - } - -#ifdef EXEC_ENV_OLS - if (files.fileExists(dummy_map_file)) - { - FileContent &content_dummy_map = files.getContent(dummy_map_file); - std::stringstream dummy_map_stream( - std::string((const char *)content_dummy_map._content, content_dummy_map._size)); -#else - if (file_exists(dummy_map_file)) - { - std::ifstream dummy_map_stream(dummy_map_file); - assert(dummy_map_stream.is_open()); -#endif - std::string line, token; - - while (std::getline(dummy_map_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t dummy_id; - uint32_t real_id; - while (std::getline(iss, token, ',')) - { - if (cnt == 0) - dummy_id = (uint32_t)stoul(token); - else - real_id = (uint32_t)stoul(token); - cnt++; - } - _dummy_pts.insert(dummy_id); - _has_dummy_pts.insert(real_id); - _dummy_to_real_map[dummy_id] = real_id; - - if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) - _real_to_dummy_map[real_id] = std::vector(); - - _real_to_dummy_map[real_id].emplace_back(dummy_id); - } -#ifndef EXEC_ENV_OLS - dummy_map_stream.close(); -#endif - diskann::cout << "Loaded dummy map" << std::endl; - } - } - -#ifdef EXEC_ENV_OLS - _pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64); -#else - _pq_table.load_pq_centroid_bin(pq_table_bin.c_str(), nchunks_u64); -#endif - - diskann::cout << "Loaded PQ centroids and in-memory compressed vectors. #points: " << _num_points - << " #dim: " << _data_dim << " #aligned_dim: " << _aligned_dim << " #chunks: " << _n_chunks - << std::endl; - - if (_n_chunks > MAX_PQ_CHUNKS) - { - std::stringstream stream; - stream << "Error loading index. Ensure that max PQ bytes for in-memory " - "PQ data does not exceed " - << MAX_PQ_CHUNKS << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::string disk_pq_pivots_path = this->_disk_index_file + "_pq_pivots.bin"; -#ifdef EXEC_ENV_OLS - if (files.fileExists(disk_pq_pivots_path)) - { - _use_disk_index_pq = true; - // giving 0 chunks to make the _pq_table infer from the - // chunk_offsets file the correct value - _disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); -#else - if (file_exists(disk_pq_pivots_path)) - { - _use_disk_index_pq = true; - // giving 0 chunks to make the _pq_table infer from the - // chunk_offsets file the correct value - _disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); -#endif - _disk_pq_n_chunks = _disk_pq_table.get_num_chunks(); - _disk_bytes_per_point = - _disk_pq_n_chunks * sizeof(uint8_t); // revising disk_bytes_per_point since DISK PQ is used. - diskann::cout << "Disk index uses PQ data compressed down to " << _disk_pq_n_chunks << " bytes per point." - << std::endl; - } - -// read index metadata -#ifdef EXEC_ENV_OLS - // This is a bit tricky. We have to read the header from the - // disk_index_file. But this is now exclusively a preserve of the - // DiskPriorityIO class. So, we need to estimate how many - // bytes are needed to store the header and read in that many using our - // 'standard' aligned file reader approach. - reader->open(_disk_index_file); - this->setup_thread_data(num_threads); - this->_max_nthreads = num_threads; - - char *bytes = getHeaderBytes(); - ContentBuf buf(bytes, HEADER_SIZE); - std::basic_istream index_metadata(&buf); -#else - diskann::cout << "Loading index metadata from " << _disk_index_file << std::endl; - std::ifstream index_metadata(_disk_index_file, std::ios::binary); -#endif - - size_t medoid_id_on_file; -#if 1 - if (!_use_partition) - { -#endif - if (!index_metadata.is_open()) - { - diskann::cout << "Error: Could not open index metadata file: " << _disk_index_file << std::endl; - return -1; - } - - uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of - // metadata, nc should be 1) - READ_U32(index_metadata, nr); - READ_U32(index_metadata, nc); - - uint64_t disk_nnodes; - uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true - READ_U64(index_metadata, disk_nnodes); - READ_U64(index_metadata, disk_ndims); - - if (disk_nnodes != _num_points) - { - diskann::cout << "Mismatch in #points for compressed data file and disk " - "index file: " - << disk_nnodes << " vs " << _num_points << std::endl; - return -1; - } - - READ_U64(index_metadata, medoid_id_on_file); - READ_U64(index_metadata, _max_node_len); - READ_U64(index_metadata, _nnodes_per_sector); - _max_degree = ((_max_node_len - _disk_bytes_per_point) / sizeof(uint32_t)) - 1; - - if (_max_degree > defaults::MAX_GRAPH_DEGREE) - { - std::stringstream stream; - stream << "Error loading index. Ensure that max graph degree (R) does " - "not exceed " - << defaults::MAX_GRAPH_DEGREE << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // setting up concept of frozen points in disk index for streaming-DiskANN - READ_U64(index_metadata, this->_num_frozen_points); - uint64_t file_frozen_id; - READ_U64(index_metadata, file_frozen_id); - if (this->_num_frozen_points == 1) - this->_frozen_location = file_frozen_id; - if (this->_num_frozen_points == 1) - { - diskann::cout << " Detected frozen point in index at location " << this->_frozen_location - << ". Will not output it at search time." << std::endl; - } - - READ_U64(index_metadata, this->_reorder_data_exists); - if (this->_reorder_data_exists) - { - if (this->_use_disk_index_pq == false) - { - throw ANNException("Reordering is designed for used with disk PQ " - "compression option", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - READ_U64(index_metadata, this->_reorder_data_start_sector); - READ_U64(index_metadata, this->_ndims_reorder_vecs); - READ_U64(index_metadata, this->_nvecs_per_sector); - } - - diskann::cout << "Disk-Index File Meta-data: "; - diskann::cout << "# nodes per sector: " << _nnodes_per_sector; - diskann::cout << ", max node len (bytes): " << _max_node_len; - diskann::cout << ", max node degree: " << _max_degree << std::endl; - -#ifdef EXEC_ENV_OLS - delete[] bytes; -#else - index_metadata.close(); -#endif - -#ifndef EXEC_ENV_OLS - // open AlignedFileReader handle to index_file - std::string index_fname(_disk_index_file); - reader->open(index_fname); - - diskann::cout << "Disk-Index Meta: nodes per sector: " << _nnodes_per_sector - << ", max node len: " << _max_node_len << ", max node degree: " << _max_degree << std::endl; - -#endif - -#if 1 - } -#endif - - this->setup_thread_data(num_threads); - this->_max_nthreads = num_threads; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(medoids_file)) - { - size_t tmp_dim; - diskann::load_bin(files, norm_file, medoids_file, _medoids, _num_medoids, tmp_dim); -#else - if (file_exists(medoids_file)) - { - size_t tmp_dim; - diskann::load_bin(medoids_file, _medoids, _num_medoids, tmp_dim); -#endif - - if (tmp_dim != 1) - { - std::stringstream stream; - stream << "Error loading medoids file. Expected bin format of m times " - "1 vector of uint32_t." - << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -#ifdef EXEC_ENV_OLS - if (!files.fileExists(centroids_file)) - { -#else - if (!file_exists(centroids_file)) - { -#endif - diskann::cout << "Centroid data file not found. Using corresponding vectors " - "for the medoids " - << std::endl; - use_medoids_data_as_centroids(); - } - else - { - size_t num_centroids, aligned_tmp_dim; -#ifdef EXEC_ENV_OLS - diskann::load_aligned_bin(files, centroids_file, _centroid_data, num_centroids, tmp_dim, - aligned_tmp_dim); -#else - diskann::load_aligned_bin(centroids_file, _centroid_data, num_centroids, tmp_dim, aligned_tmp_dim); -#endif - if (aligned_tmp_dim != _aligned_dim || num_centroids != _num_medoids) - { - std::stringstream stream; - stream << "Error loading centroids data file. Expected bin format " - "of " - "m times data_dim vector of float, where m is number of " - "medoids " - "in medoids file."; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - } - else - { - if (_use_partition) - { - assert(false); // We do not have a valid medoid id in the partition file. - } - _num_medoids = 1; - _medoids = new uint32_t[1]; - _medoids[0] = (uint32_t)(medoid_id_on_file); - use_medoids_data_as_centroids(); - } - - std::string norm_file = std::string(_disk_index_file) + "_max_base_norm.bin"; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) - { - uint64_t dumr, dumc; - float *norm_val; - diskann::load_bin(files, norm_val, dumr, dumc); -#else - if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) - { - size_t dumr, dumc; - float *norm_val; - diskann::load_bin(norm_file, norm_val, dumr, dumc); -#endif - this->_max_base_norm = norm_val[0]; - diskann::cout << "Setting re-scaling factor of base vectors to " << this->_max_base_norm << std::endl; - delete[] norm_val; - } - - if (_use_partition) - { - read_partition_info(partition_file); - - this->_graph_index_file = graph_file; - graph_reader->open(this->_graph_index_file); - load_graph_index(this->_graph_index_file); - } - - diskann::cout << "load_from_separate_paths done." << std::endl; - return 0; -} - -#ifdef USE_BING_INFRA -bool getNextCompletedRequest(std::shared_ptr &reader, IOContext &ctx, size_t size, - int &completedIndex) -{ - if ((*ctx.m_pRequests)[0].m_callback) - { - bool waitsRemaining = false; - long completeCount = ctx.m_completeCount; - do - { - for (int i = 0; i < size; i++) - { - auto ithStatus = (*ctx.m_pRequestsStatus)[i]; - if (ithStatus == IOContext::Status::READ_SUCCESS) - { - completedIndex = i; - return true; - } - else if (ithStatus == IOContext::Status::READ_WAIT) - { - waitsRemaining = true; - } - } - - // if we didn't find one in READ_SUCCESS, wait for one to complete. - if (waitsRemaining) - { - WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100); - // this assumes the knowledge of the reader behavior (implicit - // contract). need better factoring? - } - } while (waitsRemaining); - - completedIndex = -1; - return false; - } - else - { - reader->wait(ctx, completedIndex); - return completedIndex != -1; - } -} -#endif - -template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_reorder_data, QueryStats *stats, - bool USE_DEFERRED_FETCH, bool skip_search_reorder, - bool recompute_beighbor_embeddings, bool dedup_node_dis, - float prune_ratio, const bool batch_recompute, bool global_pruning) -{ - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, std::numeric_limits::max(), - use_reorder_data, stats, USE_DEFERRED_FETCH, skip_search_reorder, recompute_beighbor_embeddings, - dedup_node_dis, prune_ratio, batch_recompute, global_pruning); -} - -template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data, QueryStats *stats, - bool USE_DEFERRED_FETCH, bool skip_search_reorder, - bool recompute_beighbor_embeddings, bool dedup_node_dis, - float prune_ratio, const bool batch_recompute, bool global_pruning) -{ - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, use_filter, filter_label, - std::numeric_limits::max(), use_reorder_data, stats, USE_DEFERRED_FETCH, - skip_search_reorder, recompute_beighbor_embeddings, dedup_node_dis, prune_ratio, batch_recompute, - global_pruning); -} - -template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data, - QueryStats *stats, bool USE_DEFERRED_FETCH, bool skip_search_reorder, - bool recompute_beighbor_embeddings, bool dedup_node_dis, - float prune_ratio, const bool batch_recompute, bool global_pruning) -{ - LabelT dummy_filter = 0; - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, false, dummy_filter, io_limit, - use_reorder_data, stats, USE_DEFERRED_FETCH, skip_search_reorder, recompute_beighbor_embeddings, - dedup_node_dis, prune_ratio, batch_recompute, global_pruning); -} - -// A helper callback for cURL -static size_t WriteCallback(void *contents, size_t size, size_t nmemb, void *userp) -{ - ((std::string *)userp)->append((char *)contents, size * nmemb); - return size * nmemb; -} - -static void *g_zmq_context = zmq_ctx_new(); - -struct ZmqContextManager -{ - ~ZmqContextManager() - { - if (g_zmq_context) - { - zmq_ctx_destroy(g_zmq_context); - g_zmq_context = nullptr; - } - } -}; -static ZmqContextManager g_zmq_manager; - -bool fetch_embeddings_zmq(const std::vector &node_ids, std::vector> &out_embeddings, - int zmq_port) -{ - // 1. Protobuf 序列化：创建请求消息 - protoembedding::NodeEmbeddingRequest req_proto; - for (const auto id : node_ids) - { - req_proto.add_node_ids(id); - } - std::string req_str; - if (!req_proto.SerializeToString(&req_str)) - { - std::cerr << "ZMQ_FETCH_ERROR: Failed to serialize NodeEmbeddingRequest.\n"; - return false; - } - - // 2. 使用线程本地(thread_local)的 Socket，实现连接复用 - // 每个线程将拥有自己独立的、持久化的 Socket - thread_local void *tl_socket = nullptr; - - // 如果当前线程的 Socket 还未创建，则初始化并连接 - if (tl_socket == nullptr) - { - // 从全局 Context 创建 Socket - tl_socket = zmq_socket(g_zmq_context, ZMQ_REQ); - if (!tl_socket) - { - std::cerr << "ZMQ_FETCH_ERROR: zmq_socket() failed: " << zmq_strerror(zmq_errno()) << "\n"; - return false; - } - - int timeout = 30000; // 30 秒超时 - zmq_setsockopt(tl_socket, ZMQ_RCVTIMEO, &timeout, sizeof(timeout)); - zmq_setsockopt(tl_socket, ZMQ_SNDTIMEO, &timeout, sizeof(timeout)); - - std::string endpoint = "tcp://127.0.0.1:" + std::to_string(zmq_port); - if (zmq_connect(tl_socket, endpoint.c_str()) != 0) - { - std::cerr << "ZMQ_FETCH_ERROR: zmq_connect() to " << endpoint << " failed: " << zmq_strerror(zmq_errno()) - << "\n"; - zmq_close(tl_socket); - tl_socket = nullptr; // 重置为空指针，以便下次调用时可以尝试重建 - return false; - } - } - - // 3. 使用已建立的连接发送请求 - if (zmq_send(tl_socket, req_str.data(), req_str.size(), 0) < 0) - { - std::cerr << "ZMQ_FETCH_ERROR: zmq_send() failed: " << zmq_strerror(zmq_errno()) << "\n"; - zmq_close(tl_socket); // 连接可能已失效，关闭它 - tl_socket = nullptr; // 重置，强制下次重建 - return false; - } - - // 4. 接收响应 - zmq_msg_t response_msg; - zmq_msg_init(&response_msg); - bool success = true; - - if (zmq_msg_recv(&response_msg, tl_socket, 0) < 0) - { - std::cerr << "ZMQ_FETCH_ERROR: zmq_msg_recv() failed: " << zmq_strerror(zmq_errno()) << "\n"; - zmq_close(tl_socket); // 同样，接收超时后连接也可能无效 - tl_socket = nullptr; // 重置，强制下次重建 - success = false; - } - else - { - // 5. Protobuf 反序列化并提取数据 - protoembedding::NodeEmbeddingResponse resp_proto; - if (!resp_proto.ParseFromArray(zmq_msg_data(&response_msg), static_cast(zmq_msg_size(&response_msg)))) - { - std::cerr << "ZMQ_FETCH_ERROR: Failed to parse NodeEmbeddingResponse from server.\n"; - success = false; - } - else - { - if (resp_proto.dimensions_size() == 2) - { - int batch_size = resp_proto.dimensions(0); - int embedding_dim = resp_proto.dimensions(1); - const std::string &emb_data = resp_proto.embeddings_data(); - size_t expected_bytes = (size_t)batch_size * embedding_dim * sizeof(float); - - if (batch_size >= 0 && emb_data.size() == expected_bytes) - { - out_embeddings.resize(batch_size); - if (batch_size > 0) - { - const float *float_data = reinterpret_cast(emb_data.data()); - for (int i = 0; i < batch_size; ++i) - { - out_embeddings[i].resize(embedding_dim); - std::memcpy(out_embeddings[i].data(), float_data + (size_t)i * embedding_dim, - embedding_dim * sizeof(float)); - } - } - } - else - { - std::cerr << "ZMQ_FETCH_ERROR: Embedding data size mismatch. Expected " << expected_bytes - << " bytes, got " << emb_data.size() << ".\n"; - success = false; - } - } - else - { - std::cerr << "ZMQ_FETCH_ERROR: Server response has invalid dimensions size.\n"; - success = false; - } - } - } - - // 6. 清理消息对象，但保持 Socket 和 Context 开放以备下次复用 - zmq_msg_close(&response_msg); - - return success; -} - -/** - * fetch_embeddings_http: Function for backward compatibility, now uses ZMQ exclusively - */ -bool fetch_embeddings_http(const std::vector &node_ids, std::vector> &out_embeddings) -{ - // Use ZMQ implementation exclusively - return fetch_embeddings_zmq(node_ids, out_embeddings, 5555); -} - -//! Should be aligned with utils.h::prepare_base_for_inner_products -void preprocess_fetched_embeddings(std::vector> &embeddings, diskann::Metric metric, - float max_base_norm, uint32_t data_dim) -{ - for (auto &emb : embeddings) - { - // Ensure embedding has correct size - if (emb.size() < data_dim - 1) - { - // Pad with zeros if needed - emb.resize(data_dim - 1, 0); - } - - if (metric == diskann::Metric::INNER_PRODUCT) - { - // For inner product, apply same preprocessing as in prepare_base_for_inner_products - - // Calculate original norm - float norm_sq = 0; - for (size_t i = 0; i < data_dim - 1; i++) - { - norm_sq += emb[i] * emb[i]; - } - - // Normalize by max_base_norm (same as in index construction) - for (size_t i = 0; i < data_dim - 1; i++) - { - emb[i] /= max_base_norm; - } - - // Add the extra coordinate for MIPS->L2 conversion - float res = 1 - (norm_sq / (max_base_norm * max_base_norm)); - res = res <= 0 ? 0 : std::sqrt(res); - emb.resize(data_dim, res); - } - else if (metric == diskann::Metric::COSINE) - { - // For cosine similarity, just normalize the vector - float norm = 0; - for (auto val : emb) - { - norm += val * val; - } - norm = std::sqrt(norm); - - if (norm > 0) - { - for (size_t i = 0; i < emb.size(); i++) - { - emb[i] /= norm; - } - } - } - // For L2, no preprocessing needed - } -} - -template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const uint32_t io_limit, const bool use_reorder_data, - QueryStats *stats, bool USE_DEFERRED_FETCH, bool skip_search_reorder, - bool recompute_beighbor_embeddings, const bool dedup_node_dis, - float prune_ratio, const bool batch_recompute, bool global_pruning) -{ - // printf("cached_beam_search\n"); - // diskann::cout << "cached_beam_search" << std::endl; - // diskann out prune_ratio - prune_ratio = 1 - prune_ratio; - diskann::cout << "reserve ratio: " << prune_ratio << std::endl; - // prune_ratio = 0.8; - uint64_t num_sector_per_nodes = DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN); - if (beam_width > num_sector_per_nodes * defaults::MAX_N_SECTOR_READS) - throw ANNException("Beamwidth can not be higher than defaults::MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__, - __LINE__); - - ScratchStoreManager> manager(this->_thread_data); - auto data = manager.scratch_space(); - IOContext &ctx = data->ctx; - auto query_scratch = &(data->scratch); - auto pq_query_scratch = query_scratch->pq_scratch(); - - // reset query scratch - query_scratch->reset(); - - // copy query to thread specific aligned and allocated memory (for distance - // calculations we need aligned data) - float query_norm = 0; - T *aligned_query_T = query_scratch->aligned_query_T(); - float *query_float = pq_query_scratch->aligned_query_float; - float *query_rotated = pq_query_scratch->rotated_query; - - // Add cache hit tracking variables - uint64_t total_nodes_requested = 0; - uint64_t total_nodes_from_cache = 0; - - // normalization step. for cosine, we simply normalize the query - // for mips, we normalize the first d-1 dims, and add a 0 for last dim, since an extra coordinate was used to - // convert MIPS to L2 search - if (metric == diskann::Metric::INNER_PRODUCT || metric == diskann::Metric::COSINE) - { - uint64_t inherent_dim = (metric == diskann::Metric::COSINE) ? this->_data_dim : (uint64_t)(this->_data_dim - 1); - for (size_t i = 0; i < inherent_dim; i++) - { - aligned_query_T[i] = query1[i]; - query_norm += query1[i] * query1[i]; - } - if (metric == diskann::Metric::INNER_PRODUCT) - aligned_query_T[this->_data_dim - 1] = 0; - - query_norm = std::sqrt(query_norm); - - for (size_t i = 0; i < inherent_dim; i++) - { - aligned_query_T[i] = (T)(aligned_query_T[i] / query_norm); - } - pq_query_scratch->initialize(this->_data_dim, aligned_query_T); - } - else - { - for (size_t i = 0; i < this->_data_dim; i++) - { - aligned_query_T[i] = query1[i]; - } - pq_query_scratch->initialize(this->_data_dim, aligned_query_T); - } - - // pointers to buffers for data - T *data_buf = query_scratch->coord_scratch; - _mm_prefetch((char *)data_buf, _MM_HINT_T1); - - // sector scratch - char *sector_scratch = query_scratch->sector_scratch; - size_t §or_scratch_idx = query_scratch->sector_idx; - const uint64_t num_sectors_per_node = - _nnodes_per_sector > 0 ? 1 : DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN); - - // query <-> PQ chunk centers distances - _pq_table.preprocess_query(query_rotated); // center the query and rotate if - // we have a rotation matrix - float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - _pq_table.populate_chunk_distances(query_rotated, pq_dists); - // Preprocess Distance b/w Query Vector and Centroids - // Chunk 1 | Chunk 2 | Chunk 3 - // Centroid 1 d[1][1] d[1][2] d[1][3] - // Centroid 2 - // Centroid 3 - // Centroid 4 - // Centroid 5 - // Centroid 6 - // Centroid 7 - // Centroid 8 - - // query <-> neighbor list - float *dist_scratch = pq_query_scratch->aligned_dist_scratch; - uint8_t *pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; - - std::map node_distances; - - // Lambda to batch compute query<->node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists, aligned_query_T, recompute_beighbor_embeddings, data_buf, - &node_distances, &total_nodes_requested, &total_nodes_from_cache, - dedup_node_dis](const uint32_t *ids, const uint64_t n_ids, float *dists_out) { - // Vector[0], {3, 6, 2} - // Distance = d[3][1] + d[6][2] + d[2][3] - // recompute_beighbor_embeddings = true; - if (!recompute_beighbor_embeddings) - { - diskann::aggregate_coords(ids, n_ids, this->data, this->_n_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->_n_chunks, pq_dists, dists_out); - } - else - { - // Fetch the embeddings from the embedding server using n_ids - std::vector node_ids; - - // Update total nodes requested counter - total_nodes_requested += n_ids; - - // Build a map from node_id to original position for O(1) lookup - // Handle deduplication if enabled - std::vector cached_node_idx(n_ids, false); - if (dedup_node_dis) - { - // First pass: use cached distances where available - for (size_t i = 0; i < n_ids; i++) - { - if (node_distances.find(ids[i]) != node_distances.end()) - { - // Use cached distance - dists_out[i] = node_distances[ids[i]]; - cached_node_idx[i] = true; - total_nodes_from_cache++; // Count cache hits - } - else - { - // Not in cache, need to compute - node_ids.push_back(ids[i]); - } - } - - // If all distances are cached, we can return early - if (node_ids.empty()) - return; - } - else - { - node_ids = std::vector(ids, ids + n_ids); - } - - // Fetch embeddings from the embedding server - std::vector> embeddings; - bool success = fetch_embeddings_http(node_ids, embeddings); - - if (!success || embeddings.size() != node_ids.size()) - { - diskann::cout << "Failed to fetch embeddings from the embedding server" << std::endl; - // Fallback to PQ-based distance computation if fetching fails - diskann::aggregate_coords(ids, n_ids, this->data, this->_n_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->_n_chunks, pq_dists, dists_out); - return; - } - - // Preprocess the fetched embeddings to match the format used in diskann - preprocess_fetched_embeddings(embeddings, this->metric, this->_max_base_norm, this->_data_dim); - - // Compute distances for fetched embeddings - if (dedup_node_dis) - { - // Process each node that needs computation - uint32_t idx = 0; - for (size_t i = 0; i < n_ids; i++) - { - if (cached_node_idx[i]) - { - continue; - } - // Prepare embedding for distance computation - embeddings[idx].resize(this->_aligned_dim, 0); - memcpy(data_buf, embeddings[idx].data(), this->_aligned_dim * sizeof(T)); - - // Compute distance - float distance = - this->_dist_cmp->compare(aligned_query_T, data_buf, static_cast(this->_aligned_dim)); - - // Store results - dists_out[i] = distance; - node_distances[node_ids[i]] = distance; - idx++; - } - } - else - { - // Without deduplication, embeddings match the original order - for (size_t i = 0; i < n_ids; i++) - { - // Prepare embedding for distance computation - embeddings[i].resize(this->_aligned_dim, 0); - memcpy(data_buf, embeddings[i].data(), this->_aligned_dim * sizeof(T)); - - // Compute distance - float distance = - this->_dist_cmp->compare(aligned_query_T, data_buf, static_cast(this->_aligned_dim)); - - // Store results - dists_out[i] = distance; - } - } - } - }; - - // Add logic of global pruning - // Using a priority queue to record the PQ distance - use min heap for nearest neighbors - std::priority_queue, std::vector>, - std::greater>> - aq_priority_queue; - tsl::robin_set &visited = query_scratch->visited; - - // TODO: implement this function - // 1. Based on some heristic to prune the node_nbrs and nnbrs that is not promising - // 1.1 heruistic 1: use higher compression PQ to prune the node_nbrs and nnbrs that is not promising in path - // /powerrag/scaling_out/embeddings/facebook/contriever-msmarco/rpj_wiki/compressed_2/ - // 1.2 heruistic 2: use a lightweight reranker to rerank the node_nbrs and nnbrs that is not promising - auto prune_node_nbrs = [this, pq_coord_scratch, pq_dists, recompute_beighbor_embeddings, dedup_node_dis, - prune_ratio, global_pruning, &aq_priority_queue, - &visited](uint32_t *&node_nbrs, uint64_t &nnbrs) { - if (!recompute_beighbor_embeddings) - { - return; - } - if (nnbrs <= 10) - { - // Don't prune if there are very few neighbors - return; - } - - // Allocate space for distance calculations - float *dists_out = new float[nnbrs]; - - // Compute distances using PQ directly instead of compute_dists - diskann::aggregate_coords(node_nbrs, nnbrs, this->data, this->_n_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, nnbrs, this->_n_chunks, pq_dists, dists_out); - - if (global_pruning) - { - // Add the distance and node_id to the priority queue - for (uint64_t i = 0; i < nnbrs; i++) - { - aq_priority_queue.push(std::make_pair(dists_out[i], node_nbrs[i])); - } - // select all ratio=prune_ratio in aq_priority_queue but need to check if the node_id is already visited, - // dont need to pop - std::vector> promising_nodes; - - std::vector> roll_back_nodes; - // 1. visit top prune_ratio*length of aq_priority_queue nodes in aq_priority_queue and put the node_id not - // visited into a vector - uint64_t original_size = aq_priority_queue.size(); - for (uint64_t i = 0; i < prune_ratio * original_size; i++) - { - auto top_node = aq_priority_queue.top(); - roll_back_nodes.push_back(top_node); - aq_priority_queue.pop(); - if (visited.find(top_node.second) == visited.end()) - { - float distance = top_node.first; - uint32_t node_id = top_node.second; - promising_nodes.push_back(std::make_pair(distance, node_id)); - } - } - // push all roll_back_nodes back to aq_priority_queue - for (uint64_t i = 0; i < roll_back_nodes.size(); i++) - { - aq_priority_queue.push(roll_back_nodes[i]); - } - - // 2. assing the node_id and distance to node_nbrs and nnbrs - for (uint64_t i = 0; i < promising_nodes.size(); i++) - { - node_nbrs[i] = promising_nodes[i].second; - } - nnbrs = promising_nodes.size(); - // then return corresponding node_nbrs and nnbrs - - delete[] dists_out; - return; - } - // Create a vector of pairs (node_id, distance) - std::vector> scored_nbrs; - scored_nbrs.reserve(nnbrs); - - for (uint64_t i = 0; i < nnbrs; i++) - { - scored_nbrs.emplace_back(node_nbrs[i], dists_out[i]); - } - - // Sort by distance (lower is better) - std::sort(scored_nbrs.begin(), scored_nbrs.end(), - [](const std::pair &a, const std::pair &b) { - return a.second < b.second; - }); - - // Keep only the top portion of neighbors based on prune_ratio (or at least 10) - uint64_t new_nnbrs = std::max(10UL, static_cast(nnbrs * prune_ratio)); - if (new_nnbrs < nnbrs) - { - // Update the original node_nbrs array with pruned neighbors - for (uint64_t i = 0; i < new_nnbrs; i++) - { - node_nbrs[i] = scored_nbrs[i].first; - } - - // Update the count of neighbors - nnbrs = new_nnbrs; - } - - // Free the allocated memory - delete[] dists_out; - }; - Timer query_timer, io_timer, cpu_timer; - - NeighborPriorityQueue &retset = query_scratch->retset; - retset.reserve(l_search); - std::vector &full_retset = query_scratch->full_retset; - std::vector points_to_compute; // Store points for later embedding computation - -#if 0 - std::vector exact_dist_retset; - std::vector> exact_embeddings; -#endif - - uint32_t best_medoid = 0; - float best_dist = (std::numeric_limits::max)(); - if (!use_filter) - { - for (uint64_t cur_m = 0; cur_m < _num_medoids; cur_m++) - { - float cur_expanded_dist = - _dist_cmp_float->compare(query_float, _centroid_data + _aligned_dim * cur_m, (uint32_t)_aligned_dim); - if (cur_expanded_dist < best_dist) - { - best_medoid = _medoids[cur_m]; - best_dist = cur_expanded_dist; - } - } - } - else - { - if (_filter_to_medoid_ids.find(filter_label) != _filter_to_medoid_ids.end()) - { - const auto &medoid_ids = _filter_to_medoid_ids[filter_label]; - for (uint64_t cur_m = 0; cur_m < medoid_ids.size(); cur_m++) - { - // for filtered index, we dont store global centroid data as for unfiltered index, so we use PQ distance - // as approximation to decide closest medoid matching the query filter. - compute_dists(&medoid_ids[cur_m], 1, dist_scratch); - float cur_expanded_dist = dist_scratch[0]; - if (cur_expanded_dist < best_dist) - { - best_medoid = medoid_ids[cur_m]; - best_dist = cur_expanded_dist; - } - } - } - else - { - throw ANNException("Cannot find medoid for specified filter.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - - compute_dists(&best_medoid, 1, dist_scratch); - retset.insert(Neighbor(best_medoid, dist_scratch[0])); - visited.insert(best_medoid); - - uint32_t cmps = 0; - uint32_t hops = 0; - uint32_t num_ios = 0; - - // cleared every iteration - std::vector frontier; - frontier.reserve(2 * beam_width); - std::vector> frontier_nhoods; - frontier_nhoods.reserve(2 * beam_width); - std::vector frontier_read_reqs; - frontier_read_reqs.reserve(2 * beam_width); - std::vector>> cached_nhoods; - cached_nhoods.reserve(2 * beam_width); - - float *batched_dists = nullptr; - if (batch_recompute) - { - batched_dists = new float[_max_degree * beam_width + 5]; - } - - while (retset.has_unexpanded_node() && num_ios < io_limit) - { - // clear iteration state - frontier.clear(); - frontier_nhoods.clear(); - frontier_read_reqs.clear(); - cached_nhoods.clear(); - sector_scratch_idx = 0; - // find new beam - uint32_t num_seen = 0; - while (retset.has_unexpanded_node() && frontier.size() < beam_width && num_seen < beam_width) - { - auto nbr = retset.closest_unexpanded(); - num_seen++; - auto iter = _nhood_cache.find(nbr.id); - if (iter != _nhood_cache.end()) - { - cached_nhoods.push_back(std::make_pair(nbr.id, iter->second)); - if (stats != nullptr) - { - stats->n_cache_hits++; - } - } - else - { - frontier.push_back(nbr.id); - } - if (this->_count_visited_nodes) - { - reinterpret_cast &>(this->_node_visit_counter[nbr.id].second).fetch_add(1); - } - } - - std::vector graph_read_reqs; - std::map node_offsets; // id -> offset - std::map> node_nbrs_ori; - std::map> node_cords; - - // read nhoods of frontier ids - if (!frontier.empty()) - { - if (stats != nullptr) - stats->n_hops++; - - for (uint64_t i = 0; i < frontier.size(); i++) - { - auto id = frontier[i]; - std::pair fnhood; - fnhood.first = id; - fnhood.second = sector_scratch + num_sectors_per_node * sector_scratch_idx * defaults::SECTOR_LEN; - sector_scratch_idx++; - frontier_nhoods.push_back(fnhood); -#if 1 - if (!_use_partition) - { -#endif - frontier_read_reqs.emplace_back(get_node_sector((size_t)id) * defaults::SECTOR_LEN, - num_sectors_per_node * defaults::SECTOR_LEN, fnhood.second); -#if 1 - } -#endif - if (stats != nullptr) - { - stats->n_4k++; - stats->n_ios++; - } - num_ios++; - } - - if (_use_partition) - { - sector_scratch_idx = 0; - for (auto &frontier_nhood : frontier_nhoods) - { - uint32_t node_id = frontier_nhood.first; - uint32_t partition_id = _id2partition[node_id]; - if (partition_id >= _num_partitions) - { - diskann::cout << "Warning: partition_id is invalid: " << partition_id << std::endl; - assert(false); - } - - std::vector part_list = _graph_partitions[partition_id]; - auto it = std::find(part_list.begin(), part_list.end(), node_id); - if (it == part_list.end()) - { - diskann::cerr << "Error: node " << node_id << " not found in partition " << partition_id - << std::endl; - assert(false); - } - size_t j = std::distance(part_list.begin(), it); - node_offsets[node_id] = j; - - uint64_t sector_offset = (partition_id + 1) * defaults::SECTOR_LEN; - // ! Keep it same with frontier_nhood.second - char *sector_buffer = sector_scratch + sector_scratch_idx * defaults::SECTOR_LEN; - sector_scratch_idx++; - - AlignedRead partition_read; - partition_read.len = defaults::SECTOR_LEN; - partition_read.buf = sector_buffer; - partition_read.offset = sector_offset; - - graph_read_reqs.emplace_back(partition_read); - } - } - - io_timer.reset(); -#if 1 - if (!_use_partition) - { -#endif -#ifdef USE_BING_INFRA - reader->read(frontier_read_reqs, ctx, - true); // asynhronous reader for Bing. -#else - reader->read(frontier_read_reqs, ctx); // synchronous IO linux -#endif -#if 1 - } -#endif - -#if 0 - for (auto &[node_id, disk_buf] : frontier_nhoods) - { - char *node_disk_buf = offset_to_node(disk_buf, node_id); - uint32_t *nhood_buf = offset_to_node_nhood(node_disk_buf); - uint32_t neighbor_count = *nhood_buf; - node_nbrs_ori[node_id] = std::vector(nhood_buf + 1, nhood_buf + 1 + neighbor_count); - node_cords[node_id] = - std::vector(offset_to_node_coords(node_disk_buf), - offset_to_node_coords(node_disk_buf) + _disk_bytes_per_point / sizeof(float)); - } -#endif - if (_use_partition) - { - graph_reader->read(graph_read_reqs, ctx); - } - - if (stats != nullptr) - { - stats->io_us += (float)io_timer.elapsed(); - } - } - - // process cached nhoods - for (auto &cached_nhood : cached_nhoods) - { - auto global_cache_iter = _coord_cache.find(cached_nhood.first); - uint32_t node_id = cached_nhood.first; - T *node_fp_coords_copy = global_cache_iter->second; - float cur_expanded_dist; - float exact_expanded_dist = 0; - - if (skip_search_reorder) - { - compute_dists(&node_id, 1, dist_scratch); - cur_expanded_dist = dist_scratch[0]; - } - else if (USE_DEFERRED_FETCH) - { - cur_expanded_dist = 0.0f; - } - else if (!_use_disk_index_pq) - { - cur_expanded_dist = _dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)_aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = _disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); - else - cur_expanded_dist = _disk_pq_table.l2_distance( // disk_pq does not support OPQ yet - query_float, (uint8_t *)node_fp_coords_copy); - } - full_retset.push_back(Neighbor(node_id, cur_expanded_dist)); - -#if 0 - if (!_use_disk_index_pq) - { - exact_expanded_dist = _dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)_aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - exact_expanded_dist = _disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); - else - exact_expanded_dist = _disk_pq_table.l2_distance(query_float, (uint8_t *)node_fp_coords_copy); - } - exact_dist_retset.push_back(Neighbor(node_id, exact_expanded_dist)); - exact_embeddings.push_back(std::vector(node_fp_coords_copy, node_fp_coords_copy + _aligned_dim)); -#endif - - uint64_t nnbrs = cached_nhood.second.first; - uint32_t *node_nbrs = cached_nhood.second.second; - - // compute node_nbrs <-> query dists in PQ space - cpu_timer.reset(); - compute_dists(node_nbrs, nnbrs, dist_scratch); - if (stats != nullptr) - { - stats->n_cmps += (uint32_t)nnbrs; - stats->cpu_us += (float)cpu_timer.elapsed(); - } - - // process prefetched nhood - for (uint64_t m = 0; m < nnbrs; ++m) - { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) - { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) - continue; - - if (use_filter && !(point_has_label(id, filter_label)) && - (!_use_universal_label || !point_has_label(id, _universal_filter_label))) - continue; - cmps++; - float dist = dist_scratch[m]; - Neighbor nn(id, dist); - retset.insert(nn); - } - } - } -#ifdef USE_BING_INFRA - // process each frontier nhood - compute distances to unvisited nodes - int completedIndex = -1; - long requestCount = static_cast(frontier_read_reqs.size()); - // If we issued read requests and if a read is complete or there are - // reads in wait state, then enter the while loop. - while (requestCount > 0 && getNextCompletedRequest(reader, ctx, requestCount, completedIndex)) - { - assert(completedIndex >= 0); - auto &frontier_nhood = frontier_nhoods[completedIndex]; - (*ctx.m_pRequestsStatus)[completedIndex] = IOContext::PROCESS_COMPLETE; -#else - std::vector batched_node_ids; - - for (auto &frontier_nhood : frontier_nhoods) - { -#endif - uint32_t node_id = frontier_nhood.first; - char *disk_buf = frontier_nhood.second; - char *node_disk_buf = offset_to_node(disk_buf, node_id); - - float cur_expanded_dist; - - // If skip_reorder is true, compute both PQ distance and exact distance - if (skip_search_reorder) - { - compute_dists(&node_id, 1, dist_scratch); - cur_expanded_dist = dist_scratch[0]; - } - else if (USE_DEFERRED_FETCH) - { - cur_expanded_dist = 0.0f; - } - else if (recompute_beighbor_embeddings && dedup_node_dis && _use_partition) - { - // For _use_partition = True, we must rely on node_distances to get the distance - // Since we are using graph-structure only reading. - // ! Use node_distances to get the distance - cur_expanded_dist = node_distances[node_id]; - } - else - { -#if 0 - if (node_cords.find(node_id) == node_cords.end()) - { - diskann::cout << "Warning: node " << node_id << " not found in node_cords" << std::endl; - diskann::cout << "Are you using deferred fetch for detached graph?" << std::endl; - assert(false); - } - // ! As for DEBUG mode and partition_read = True, we are overriding the node_disk_buf - // ! with our graph-structure only reading. So we need to use node_cords to get the correct - // ! coordinates. - T *node_fp_coords = reinterpret_cast(node_cords[node_id].data()); - // T *node_fp_coords = offset_to_node_coords(node_disk_buf); -#endif - T *node_fp_coords = offset_to_node_coords(node_disk_buf); - memcpy(data_buf, node_fp_coords, _disk_bytes_per_point); - if (!_use_disk_index_pq) - { - cur_expanded_dist = _dist_cmp->compare(aligned_query_T, data_buf, (uint32_t)_aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = _disk_pq_table.inner_product(query_float, (uint8_t *)data_buf); - else - cur_expanded_dist = _disk_pq_table.l2_distance(query_float, (uint8_t *)data_buf); - } - } - full_retset.push_back(Neighbor(node_id, cur_expanded_dist)); - -#if 0 - T *node_fp_coords = offset_to_node_coords(node_disk_buf); - memcpy(data_buf, node_fp_coords, _disk_bytes_per_point); - float exact_expanded_dist = 0; - if (!_use_disk_index_pq) - { - exact_expanded_dist = _dist_cmp->compare(aligned_query_T, data_buf, (uint32_t)_aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - exact_expanded_dist = _disk_pq_table.inner_product(query_float, (uint8_t *)data_buf); - else - exact_expanded_dist = _disk_pq_table.l2_distance(query_float, (uint8_t *)data_buf); - } - exact_dist_retset.push_back(Neighbor(node_id, exact_expanded_dist)); - exact_embeddings.push_back(std::vector(data_buf, data_buf + _aligned_dim)); -#endif - - uint32_t *node_nbrs; - uint64_t nnbrs; - - if (!_use_partition) - { - auto node_buf = offset_to_node_nhood(node_disk_buf); - nnbrs = (uint64_t)(*node_buf); - node_nbrs = (node_buf + 1); - } - -#if 0 - auto node_nbrs_vec = node_nbrs_ori[node_id]; - nnbrs = node_nbrs_vec.size(); - node_nbrs = node_nbrs_vec.data(); -#endif - if (_use_partition) - { - char *sector_buffer = frontier_nhood.second; - int j = node_offsets[node_id]; - uint64_t node_offset = j * _graph_node_len; - if (node_offset + 4 > defaults::SECTOR_LEN) - { - diskann::cerr << "Error: node offset out of range: " << node_offset << " (+4) > " - << defaults::SECTOR_LEN << " for node " << node_id << std::endl; - assert(false); - } - - char *adjacency_ptr = sector_buffer + node_offset; - uint32_t neighbor_count = *reinterpret_cast(adjacency_ptr); - - if (neighbor_count > 10000) - { - diskann::cerr << "Error: suspicious neighbor count: " << neighbor_count << " for node " << node_id - << std::endl; - assert(false); - } - - size_t needed = neighbor_count * sizeof(uint32_t); - if (node_offset + 4 + needed > defaults::SECTOR_LEN) - { - diskann::cerr << "Error: neighbor data out of range: " << (node_offset + 4 + needed) << " > " - << defaults::SECTOR_LEN << " for node " << node_id << std::endl; - assert(false); - } - -#if 0 - if (neighbor_count != nnbrs) - { - diskann::cout << "Warning: neighbor_count != nnbrs: " << neighbor_count << " != " << nnbrs - << std::endl; - assert(false); - } -#endif - - nnbrs = neighbor_count; - -#if 0 - uint32_t *our_node_nbrs = (uint32_t *)(adjacency_ptr + 4); - for (uint32_t i = 0; i < nnbrs; i++) - { - if (our_node_nbrs[i] != node_nbrs[i]) - { - diskann::cout << "Warning: our_node_nbrs[" << i << "] != node_nbrs[" << i - << "]: " << our_node_nbrs[i] << " != " << node_nbrs[i] << std::endl; - assert(false); - } - } -#endif - - node_nbrs = reinterpret_cast(adjacency_ptr + 4); - } - - // compute node_nbrs <-> query dist in PQ space - cpu_timer.reset(); - // have a function to prune the node_nbrs and nnbrs - - // prune_node_nbrs(node_nbrs, nnbrs); - - if (!batch_recompute) - { - prune_node_nbrs(node_nbrs, nnbrs); - compute_dists(node_nbrs, nnbrs, dist_scratch); - if (stats != nullptr) - { - stats->n_cmps += (uint32_t)nnbrs; - stats->cpu_us += (float)cpu_timer.elapsed(); - } - - cpu_timer.reset(); - // process prefetch-ed nhood - for (uint64_t m = 0; m < nnbrs; ++m) - { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) - { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) - continue; - - if (use_filter && !(point_has_label(id, filter_label)) && - (!_use_universal_label || !point_has_label(id, _universal_filter_label))) - continue; - cmps++; - float dist = dist_scratch[m]; - if (stats != nullptr) - { - stats->n_cmps++; - } - - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - if (stats != nullptr) - { - stats->cpu_us += (float)cpu_timer.elapsed(); - } - } - else - { - // add all the node_nbrs to the batch_requests - batched_node_ids.insert(batched_node_ids.end(), node_nbrs, node_nbrs + nnbrs); - } - } - - if (batch_recompute) - { - auto nnbrs = batched_node_ids.size(); - uint32_t *batched_data_ptr = batched_node_ids.data(); // Get pointer to data - prune_node_nbrs(batched_data_ptr, nnbrs); // Prune using the pointer, nnbrs is updated - - compute_dists(batched_data_ptr, nnbrs, batched_dists); // Compute dists for the pruned set - // ! Not sure if dist_scratch has enough space - - // process prefetch-ed nhood - for (uint64_t m = 0; m < nnbrs; ++m) - { - uint32_t id = batched_node_ids[m]; - if (visited.insert(id).second) - { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) - continue; - - if (use_filter && !(point_has_label(id, filter_label)) && - (!_use_universal_label || !point_has_label(id, _universal_filter_label))) - continue; - cmps++; - float dist = batched_dists[m]; - if (stats != nullptr) - { - stats->n_cmps++; - } - - Neighbor nn(id, dist); - retset.insert(nn); - } - } - } - // } - // } - hops++; - } - - delete[] batched_dists; - - diskann::cout << "Graph traversal completed, hops: " << hops << std::endl; - - if (USE_DEFERRED_FETCH) - { - diskann::cout << "hops: " << hops << std::endl; - - std::vector node_ids; - node_ids.reserve(full_retset.size()); - for (auto &nr : full_retset) - { - node_ids.push_back(nr.id); - } - - Timer fetch_timer; - std::vector> real_embeddings; - bool success = fetch_embeddings_http(node_ids, real_embeddings); - if (!success) - { - throw ANNException("Failed to fetch embeddings", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::cout << "Fetched " << real_embeddings.size() << " embeddings in " << fetch_timer.elapsed() << " us" - << std::endl; - - // compute real-dist - Timer compute_timer; - // preprocess the real embedding to match the format of nomarlized version of diskann - preprocess_fetched_embeddings(real_embeddings, metric, _max_base_norm, this->_data_dim); - -#if 0 - assert(real_embeddings.size() == full_retset.size()); - assert(real_embeddings.size() == exact_dist_retset.size()); - assert(real_embeddings.size() == exact_embeddings.size()); -#endif - - for (int i = 0; i < real_embeddings.size(); i++) - { - // padding real_embeddings[i] to _aligned_dim - real_embeddings[i].resize(_aligned_dim, 0); -#if 0 - // compare real_embeddings[i] with exact_embeddings[i] - if (real_embeddings[i].size() != exact_embeddings[i].size()) - { - diskann::cout << "real_embeddings[i].size(): " << real_embeddings[i].size() << std::endl; - diskann::cout << "exact_embeddings[i].size(): " << exact_embeddings[i].size() << std::endl; - - // dumping to files - std::ofstream diff_file("./diff_embeddings.txt"); - diff_file << "real_embeddings[i].size(): " << real_embeddings[i].size() << std::endl; - diff_file << "exact_embeddings[i].size(): " << exact_embeddings[i].size() << std::endl; - for (int j = 0; j < real_embeddings[i].size(); j++) - { - diff_file << real_embeddings[i][j] << " "; - } - diff_file << std::endl; - for (int j = 0; j < exact_embeddings[i].size(); j++) - { - diff_file << exact_embeddings[i][j] << " "; - } - diff_file << std::endl; - assert(false); - } - for (int j = 0; j < real_embeddings[i].size(); j++) - { - if (abs(real_embeddings[i][j] - exact_embeddings[i][j]) > 5e-4) - { - diskann::cout << "Difference found at node_id: " << full_retset[i].id << " and dimension: " << j - << std::endl; - diskann::cout << "real_embeddings[i][j]: " << real_embeddings[i][j] << std::endl; - diskann::cout << "exact_embeddings[i][j]: " << exact_embeddings[i][j] << std::endl; - assert(false); - } - } -#endif - - float dist; - assert(!_use_disk_index_pq); - memcpy(data_buf, real_embeddings[i].data(), real_embeddings[0].size() * sizeof(T)); - dist = _dist_cmp->compare(aligned_query_T, data_buf, (uint32_t)_aligned_dim); - - full_retset[i].distance = dist; - -#if 0 - if (abs(dist - exact_dist_retset[i].distance) > 5e-4) - { - diskann::cout << "Difference found at node_id: " << full_retset[i].id << std::endl; - diskann::cout << "dist: " << dist << std::endl; - diskann::cout << "exact_dist_retset[i].distance: " << exact_dist_retset[i].distance << std::endl; - assert(false); - } -#endif - } - diskann::cout << "compute_timer.elapsed(): " << compute_timer.elapsed() << std::endl; - } - - std::sort(full_retset.begin(), full_retset.end()); - -// Compare PQ results with exact results when skip_search_reorder is true -#if 0 - if (skip_search_reorder) - { - // Sort the exact distance results - std::sort(exact_dist_retset.begin(), exact_dist_retset.end()); - - // Create a map to find positions of IDs in the PQ-sorted list - std::unordered_map pq_positions; - for (size_t i = 0; i < full_retset.size(); i++) - { - pq_positions[full_retset[i].id] = i; - } - - int current_search_id = search_counter.fetch_add(1); - int thread_id = omp_get_thread_num(); - - std::lock_guard lock(log_file_mutex); - - std::ofstream log_file("./top3_positions_log.txt", std::ios::app); - // Write header if file is empty - log_file.seekp(0, std::ios::end); - if (log_file.tellp() == 0) - { - diskann::cout << "Saved top3 distributions to " << std::filesystem::canonical("./top3_positions_log.txt") - << std::endl; - log_file << "Search#,ThreadID,FullSetSize,Rank,ID,PQ_Rank,PQ_Distance,Exact_Distance" << std::endl; - } - - // Log the top-k results from exact distance sorting and their positions in PQ-sorted list - size_t top_k = std::min((size_t)k_search, exact_dist_retset.size()); - for (size_t i = 0; i < top_k; i++) - { - uint32_t id = exact_dist_retset[i].id; - float exact_dist = exact_dist_retset[i].distance; - - // Find this ID's position in the PQ-sorted list - size_t pq_pos = pq_positions.count(id) ? pq_positions[id] : full_retset.size(); - float pq_dist = (pq_pos < full_retset.size()) ? full_retset[pq_pos].distance : -1; - - log_file << current_search_id << "," << thread_id << "," << full_retset.size() << "," << i + 1 << "," << id - << "," << pq_pos + 1 << "," << pq_dist << "," << exact_dist << std::endl; - } - - log_file.close(); - } -#endif - - if (use_reorder_data) - { - if (!(this->_reorder_data_exists)) - { - throw ANNException("Requested use of reordering data which does " - "not exist in index " - "file", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::vector vec_read_reqs; - - if (full_retset.size() > k_search * FULL_PRECISION_REORDER_MULTIPLIER) - full_retset.erase(full_retset.begin() + k_search * FULL_PRECISION_REORDER_MULTIPLIER, full_retset.end()); - - for (size_t i = 0; i < full_retset.size(); ++i) - { - // MULTISECTORFIX - vec_read_reqs.emplace_back(VECTOR_SECTOR_NO(((size_t)full_retset[i].id)) * defaults::SECTOR_LEN, - defaults::SECTOR_LEN, sector_scratch + i * defaults::SECTOR_LEN); - - if (stats != nullptr) - { - stats->n_4k++; - stats->n_ios++; - } - } - - io_timer.reset(); -#ifdef USE_BING_INFRA - reader->read(vec_read_reqs, ctx, true); // async reader windows. -#else - reader->read(vec_read_reqs, ctx); // synchronous IO linux -#endif - if (stats != nullptr) - { - stats->io_us += io_timer.elapsed(); - } - - for (size_t i = 0; i < full_retset.size(); ++i) - { - auto id = full_retset[i].id; - // MULTISECTORFIX - auto location = (sector_scratch + i * defaults::SECTOR_LEN) + VECTOR_SECTOR_OFFSET(id); - full_retset[i].distance = _dist_cmp->compare(aligned_query_T, (T *)location, (uint32_t)this->_data_dim); - } - - std::sort(full_retset.begin(), full_retset.end()); - } - - // copy k_search values - for (uint64_t i = 0; i < k_search; i++) - { - indices[i] = full_retset[i].id; - auto key = (uint32_t)indices[i]; - if (_dummy_pts.find(key) != _dummy_pts.end()) - { - indices[i] = _dummy_to_real_map[key]; - } - - if (distances != nullptr) - { - distances[i] = full_retset[i].distance; - if (metric == diskann::Metric::INNER_PRODUCT) - { - // flip the sign to convert min to max - distances[i] = (-distances[i]); - // rescale to revert back to original norms (cancelling the - // effect of base and query pre-processing) - if (_max_base_norm != 0) - distances[i] *= (_max_base_norm * query_norm); - } - } - } - -#ifdef USE_BING_INFRA - ctx.m_completeCount = 0; -#endif - - if (stats != nullptr) - { - stats->total_us = (float)query_timer.elapsed(); - } - - // After search is complete, print cache hit rate statistics - if (recompute_beighbor_embeddings && dedup_node_dis && total_nodes_requested > 0) - { - float cache_hit_rate = static_cast(total_nodes_from_cache) / total_nodes_requested * 100.0f; - diskann::cout << "Node distance cache statistics:" << std::endl; - diskann::cout << " Total nodes requested: " << total_nodes_requested << std::endl; - diskann::cout << " Nodes served from cache: " << total_nodes_from_cache << std::endl; - diskann::cout << " Cache hit rate: " << cache_hit_rate << "%" << std::endl; - } -} - -// range search returns results of all neighbors within distance of range. -// indices and distances need to be pre-allocated of size l_search and the -// return value is the number of matching hits. -template -uint32_t PQFlashIndex::range_search(const T *query1, const double range, const uint64_t min_l_search, - const uint64_t max_l_search, std::vector &indices, - std::vector &distances, const uint64_t min_beam_width, - QueryStats *stats) -{ - uint32_t res_count = 0; - - bool stop_flag = false; - - uint32_t l_search = (uint32_t)min_l_search; // starting size of the candidate list - while (!stop_flag) - { - indices.resize(l_search); - distances.resize(l_search); - uint64_t cur_bw = min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; - cur_bw = (cur_bw > 100) ? 100 : cur_bw; - for (auto &x : distances) - x = std::numeric_limits::max(); - this->cached_beam_search(query1, l_search, l_search, indices.data(), distances.data(), cur_bw, false, stats); - for (uint32_t i = 0; i < l_search; i++) - { - if (distances[i] > (float)range) - { - res_count = i; - break; - } - else if (i == l_search - 1) - res_count = l_search; - } - if (res_count < (uint32_t)(l_search / 2.0)) - stop_flag = true; - l_search = l_search * 2; - if (l_search > max_l_search) - stop_flag = true; - } - indices.resize(res_count); - distances.resize(res_count); - return res_count; -} - -template uint64_t PQFlashIndex::get_data_dim() -{ - return _data_dim; -} - -template diskann::Metric PQFlashIndex::get_metric() -{ - return this->metric; -} - -#ifdef EXEC_ENV_OLS -template char *PQFlashIndex::getHeaderBytes() -{ - IOContext &ctx = reader->get_ctx(); - AlignedRead readReq; - readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; - readReq.len = PQFlashIndex::HEADER_SIZE; - readReq.offset = 0; - - std::vector readReqs; - readReqs.push_back(readReq); - - reader->read(readReqs, ctx, false); - - return (char *)readReq.buf; -} -#endif - -template -std::vector PQFlashIndex::get_pq_vector(std::uint64_t vid) -{ - std::uint8_t *pqVec = &this->data[vid * this->_n_chunks]; - return std::vector(pqVec, pqVec + this->_n_chunks); -} - -template std::uint64_t PQFlashIndex::get_num_points() -{ - return _num_points; -} - -// instantiations -template class PQFlashIndex; -template class PQFlashIndex; -template class PQFlashIndex; -template class PQFlashIndex; -template class PQFlashIndex; -template class PQFlashIndex; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_l2_distance.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/pq_l2_distance.cpp deleted file mode 100644 index 9bd5311..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/pq_l2_distance.cpp +++ /dev/null @@ -1,284 +0,0 @@ - -#include "pq.h" -#include "pq_l2_distance.h" -#include "pq_scratch.h" - -// block size for reading/processing large files and matrices in blocks -#define BLOCK_SIZE 5000000 - -namespace diskann -{ - -template -PQL2Distance::PQL2Distance(uint32_t num_chunks, bool use_opq) : _num_chunks(num_chunks), _is_opq(use_opq) -{ -} - -template PQL2Distance::~PQL2Distance() -{ -#ifndef EXEC_ENV_OLS - if (_tables != nullptr) - delete[] _tables; - if (_chunk_offsets != nullptr) - delete[] _chunk_offsets; - if (_centroid != nullptr) - delete[] _centroid; - if (_rotmat_tr != nullptr) - delete[] _rotmat_tr; -#endif - if (_tables_tr != nullptr) - delete[] _tables_tr; -} - -template bool PQL2Distance::is_opq() const -{ - return this->_is_opq; -} - -template -std::string PQL2Distance::get_quantized_vectors_filename(const std::string &prefix) const -{ - if (_num_chunks == 0) - { - throw diskann::ANNException("Must set num_chunks before calling get_quantized_vectors_filename", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - return diskann::get_quantized_vectors_filename(prefix, _is_opq, (uint32_t)_num_chunks); -} -template std::string PQL2Distance::get_pivot_data_filename(const std::string &prefix) const -{ - if (_num_chunks == 0) - { - throw diskann::ANNException("Must set num_chunks before calling get_pivot_data_filename", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - return diskann::get_pivot_data_filename(prefix, _is_opq, (uint32_t)_num_chunks); -} -template -std::string PQL2Distance::get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const -{ - return diskann::get_rotation_matrix_suffix(pq_pivots_filename); -} - -#ifdef EXEC_ENV_OLS -template -void PQL2Distance::load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, - size_t num_chunks) -{ -#else -template -void PQL2Distance::load_pivot_data(const std::string &pq_table_file, size_t num_chunks) -{ -#endif - size_t nr, nc; - // std::string rotmat_file = get_opq_rot_matrix_filename(pq_table_file, - // false); - -#ifdef EXEC_ENV_OLS - size_t *file_offset_data; // since load_bin only sets the pointer, no need - // to delete. - diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); -#else - std::unique_ptr file_offset_data; - diskann::load_bin(pq_table_file, file_offset_data, nr, nc); -#endif - - bool use_old_filetype = false; - - if (nr != 4 && nr != 5) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 - << " or " << 5; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (nr == 4) - { - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << std::endl; - } - else if (nr == 5) - { - use_old_filetype = true; - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << file_offset_data[4] << std::endl; - } - else - { - throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); - } - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); -#else - diskann::load_bin(pq_table_file, _tables, nr, nc, file_offset_data[0]); -#endif - - if ((nr != NUM_PQ_CENTROIDS)) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". file_num_centers = " << nr - << " but expecting " << NUM_PQ_CENTROIDS << " centers"; - throw diskann::ANNException("Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - this->_ndims = nc; - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); -#else - diskann::load_bin(pq_table_file, _centroid, nr, nc, file_offset_data[1]); -#endif - - if ((nr != this->_ndims) || (nc != 1)) - { - diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". file_dim = " << nr - << ", file_cols = " << nc << " but expecting " << this->_ndims << " entries in 1 dimension."; - throw diskann::ANNException("Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - int chunk_offsets_index = 2; - if (use_old_filetype) - { - chunk_offsets_index = 3; - } -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#else - diskann::load_bin(pq_table_file, _chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#endif - - if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) - { - diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr - << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; - throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - this->_num_chunks = nr - 1; - diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->_ndims - << ", #chunks: " << this->_num_chunks << std::endl; - - // For OPQ there will be a rotation matrix to load. - if (this->_is_opq) - { - std::string rotmat_file = get_rotation_matrix_suffix(pq_table_file); -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); -#else - diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); -#endif - if (nr != this->_ndims || nc != this->_ndims) - { - diskann::cerr << "Error loading rotation matrix file" << std::endl; - throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - - // alloc and compute transpose - _tables_tr = new float[256 * this->_ndims]; - for (size_t i = 0; i < 256; i++) - { - for (size_t j = 0; j < this->_ndims; j++) - { - _tables_tr[j * 256 + i] = _tables[i * this->_ndims + j]; - } - } -} - -template uint32_t PQL2Distance::get_num_chunks() const -{ - return static_cast(_num_chunks); -} - -// REFACTOR: Instead of doing half the work in the caller and half in this -// function, we let this function -// do all of the work, making it easier for the caller. -template -void PQL2Distance::preprocess_query(const data_t *aligned_query, uint32_t dim, PQScratch &scratch) -{ - // Copy query vector to float and then to "rotated" query - for (size_t d = 0; d < dim; d++) - { - scratch.aligned_query_float[d] = (float)aligned_query[d]; - } - scratch.initialize(dim, aligned_query); - - for (uint32_t d = 0; d < _ndims; d++) - { - scratch.rotated_query[d] -= _centroid[d]; - } - std::vector tmp(_ndims, 0); - if (_is_opq) - { - for (uint32_t d = 0; d < _ndims; d++) - { - for (uint32_t d1 = 0; d1 < _ndims; d1++) - { - tmp[d] += scratch.rotated_query[d1] * _rotmat_tr[d1 * _ndims + d]; - } - } - std::memcpy(scratch.rotated_query, tmp.data(), _ndims * sizeof(float)); - } - this->prepopulate_chunkwise_distances(scratch.rotated_query, scratch.aligned_pqtable_dist_scratch); -} - -template -void PQL2Distance::preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, float *dists_out) -{ - pq_dist_lookup(pq_scratch.aligned_pq_coord_scratch, n_ids, _num_chunks, pq_scratch.aligned_pqtable_dist_scratch, - dists_out); -} - -template -void PQL2Distance::preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, - std::vector &dists_out) -{ - pq_dist_lookup(pq_scratch.aligned_pq_coord_scratch, n_ids, _num_chunks, pq_scratch.aligned_pqtable_dist_scratch, - dists_out); -} - -template float PQL2Distance::brute_force_distance(const float *query_vec, uint8_t *base_vec) -{ - float res = 0; - for (size_t chunk = 0; chunk < _num_chunks; chunk++) - { - for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = _tables_tr + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); - res += diff * diff; - } - } - return res; -} - -template -void PQL2Distance::prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec) -{ - memset(dist_vec, 0, 256 * _num_chunks * sizeof(float)); - // chunk wise distance computation - for (size_t chunk = 0; chunk < _num_chunks; chunk++) - { - // sum (q-c)^2 for the dimensions associated with this chunk - float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = _tables_tr + (256 * j); - for (size_t idx = 0; idx < 256; idx++) - { - double diff = centers_dim_vec[idx] - (query_vec[j]); - chunk_dists[idx] += (float)(diff * diff); - } - } - } -} - -template DISKANN_DLLEXPORT class PQL2Distance; -template DISKANN_DLLEXPORT class PQL2Distance; -template DISKANN_DLLEXPORT class PQL2Distance; - -} // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/search_wrapper.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/search_wrapper.cpp deleted file mode 100644 index 001e36d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/search_wrapper.cpp +++ /dev/null @@ -1,217 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include - -#include "utils.h" -#include - -#ifndef _WINDOWS -#include -#include -#include -#include "linux_aligned_file_reader.h" -#else -#ifdef USE_BING_INFRA -#include "bing_aligned_file_reader.h" -#else -#include "windows_aligned_file_reader.h" -#endif -#endif - -namespace diskann -{ -const unsigned int DEFAULT_W = 1; - -SearchResult::SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, - const float *const distances, const std::string *const tags, - const unsigned *const partitions) - : _K(K), _search_time_in_ms(elapsed_time_in_ms) -{ - for (unsigned i = 0; i < K; ++i) - { - this->_indices.push_back(indices[i]); - this->_distances.push_back(distances[i]); - if (tags != NULL) - this->_tags.push_back(tags[i]); - if (partitions != NULL) - this->_partitions.push_back(partitions[i]); - } - if (tags != nullptr) - this->_tags_enabled = true; - else - this->_tags_enabled = false; - - if (partitions != nullptr) - this->_partitions_enabled = true; - else - this->_partitions_enabled = false; -} - -BaseSearch::BaseSearch(const std::string &tagsFile) -{ - if (tagsFile.size() != 0) - { - std::ifstream in(tagsFile); - - if (!in.is_open()) - { - std::cerr << "Could not open " << tagsFile << std::endl; - } - - std::string tag; - while (std::getline(in, tag)) - { - _tags_str.push_back(tag); - } - - _tags_enabled = true; - - std::cout << "Loaded " << _tags_str.size() << " tags from " << tagsFile << std::endl; - } - else - { - _tags_enabled = false; - } -} - -void BaseSearch::lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags) -{ - if (_tags_enabled == false) - throw std::runtime_error("Can not look up tags as they are not enabled."); - else - { - for (unsigned k = 0; k < K; ++k) - { - if (indices[k] > _tags_str.size()) - throw std::runtime_error("In tag lookup, index exceeded the number of tags"); - else - ret_tags[k] = _tags_str[indices[k]]; - } - } -} - -template -InMemorySearch::InMemorySearch(const std::string &baseFile, const std::string &indexFile, - const std::string &tagsFile, Metric m, uint32_t num_threads, uint32_t search_l) - : BaseSearch(tagsFile) -{ - size_t dimensions, total_points = 0; - diskann::get_bin_metadata(baseFile, total_points, dimensions); - auto search_params = diskann::IndexSearchParams(search_l, num_threads); - _index = std::unique_ptr>( - new diskann::Index(m, dimensions, total_points, nullptr, search_params, 0, false)); - - _index->load(indexFile.c_str(), num_threads, search_l); -} - -template -SearchResult InMemorySearch::search(const T *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) -{ - unsigned int *indices = new unsigned int[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->search(query, K, Ls, indices, distances); - auto duration = - std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) - .count(); - - std::string *tags = nullptr; - if (_tags_enabled) - { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - - delete[] indices; - delete[] distances; - return result; -} - -template InMemorySearch::~InMemorySearch() -{ -} - -template -PQFlashSearch::PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, - const unsigned num_threads, const std::string &tagsFile, Metric m) - : BaseSearch(tagsFile) -{ -#ifdef _WINDOWS -#ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); -#else - reader.reset(new diskann::BingAlignedFileReader()); -#endif -#else - auto ptr = new LinuxAlignedFileReader(); - reader.reset(ptr); -#endif - - std::string index_prefix_path(indexPrefix); - std::string disk_index_file = index_prefix_path + "_disk.index"; - std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - - _index = std::unique_ptr>(new diskann::PQFlashIndex(reader, m)); - - int res = _index->load(num_threads, index_prefix_path.c_str()); - - if (res != 0) - { - std::cerr << "Unable to load index. Status code: " << res << "." << std::endl; - } - - std::vector node_list; - std::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _index->cache_bfs_levels(num_nodes_to_cache, node_list); - _index->load_cache_list(node_list); - omp_set_num_threads(num_threads); -} - -template -SearchResult PQFlashSearch::search(const T *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) -{ - uint64_t *indices_u64 = new uint64_t[K]; - unsigned *indices = new unsigned[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->cached_beam_search(query, K, Ls, indices_u64, distances, DEFAULT_W); - auto duration = - std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) - .count(); - for (unsigned k = 0; k < K; ++k) - indices[k] = indices_u64[k]; - - std::string *tags = nullptr; - if (_tags_enabled) - { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - delete[] indices_u64; - delete[] indices; - delete[] distances; - return result; -} - -template PQFlashSearch::~PQFlashSearch() -{ -} - -template class InMemorySearch; -template class InMemorySearch; -template class InMemorySearch; - -template class PQFlashSearch; -template class PQFlashSearch; -template class PQFlashSearch; -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/server.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/server.cpp deleted file mode 100644 index f79b0af..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/restapi/server.cpp +++ /dev/null @@ -1,271 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include -#include -#include -#include -#include -#include - -#include - -namespace diskann -{ - -Server::Server(web::uri &uri, std::vector> &multi_searcher, - const std::string &typestring) - : _multi_search(multi_searcher.size() > 1 ? true : false) -{ - for (auto &searcher : multi_searcher) - _multi_searcher.push_back(std::move(searcher)); - - _listener = std::unique_ptr( - new web::http::experimental::listener::http_listener(uri)); - if (typestring == std::string("float")) - { - _listener->support(std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else if (typestring == std::string("int8_t")) - { - _listener->support(web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else if (typestring == std::string("uint8_t")) - { - _listener->support(web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else - { - throw "Unsupported type in server constuctor"; - } -} - -Server::~Server() -{ -} - -pplx::task Server::open() -{ - return _listener->open(); -} -pplx::task Server::close() -{ - return _listener->close(); -} - -diskann::SearchResult Server::aggregate_results(const unsigned K, const std::vector &results) -{ - if (_multi_search) - { - auto best_indices = new unsigned[K]; - auto best_distances = new float[K]; - auto best_partitions = new unsigned[K]; - auto best_tags = results[0].tags_enabled() ? new std::string[K] : nullptr; - - auto numsearchers = _multi_searcher.size(); - std::vector pos(numsearchers, 0); - - for (size_t k = 0; k < K; ++k) - { - float best_distance = std::numeric_limits::max(); - unsigned best_partition = 0; - - for (size_t i = 0; i < numsearchers; ++i) - { - if (results[i].get_distances()[pos[i]] < best_distance) - { - best_distance = results[i].get_distances()[pos[i]]; - best_partition = i; - } - } - best_distances[k] = best_distance; - best_indices[k] = results[best_partition].get_indices()[pos[best_partition]]; - best_partitions[k] = best_partition; - if (results[best_partition].tags_enabled()) - best_tags[k] = results[best_partition].get_tags()[pos[best_partition]]; - std::cout << best_partition << " " << pos[best_partition] << std::endl; - pos[best_partition]++; - } - - unsigned int total_time = 0; - for (size_t i = 0; i < numsearchers; ++i) - total_time += results[i].get_time(); - diskann::SearchResult result = - SearchResult(K, total_time, best_indices, best_distances, best_tags, best_partitions); - - delete[] best_indices; - delete[] best_distances; - delete[] best_partitions; - delete[] best_tags; - - return result; - } - else - { - return results[0]; - } -} - -template void Server::handle_post(web::http::http_request message) -{ - message.extract_string(true) - .then([=](utility::string_t body) { - int64_t queryId = -1; - unsigned int K = 0; - try - { - T *queryVector = nullptr; - unsigned int dimensions = 0; - unsigned int Ls; - parseJson(body, K, queryId, queryVector, dimensions, Ls); - - auto startTime = std::chrono::high_resolution_clock::now(); - std::vector results; - - for (auto &searcher : _multi_searcher) - results.push_back(searcher->search(queryVector, dimensions, (unsigned int)K, Ls)); - diskann::SearchResult result = aggregate_results(K, results); - diskann::aligned_free(queryVector); - web::json::value response = prepareResponse(queryId, K); - response[INDICES_KEY] = idsToJsonArray(result); - response[DISTANCES_KEY] = distancesToJsonArray(result); - if (result.tags_enabled()) - response[TAGS_KEY] = tagsToJsonArray(result); - if (result.partitions_enabled()) - response[PARTITION_KEY] = partitionsToJsonArray(result); - - response[TIME_TAKEN_KEY] = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - startTime) - .count(); - - std::cout << "Responding to: " << queryId << std::endl; - return std::make_pair(web::http::status_codes::OK, response); - } - catch (const std::exception &ex) - { - std::cerr << "Exception while processing query: " << queryId << ":" << ex.what() << std::endl; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(ex.what()); - return std::make_pair(web::http::status_codes::InternalError, response); - } - catch (...) - { - std::cerr << "Uncaught exception while processing query: " << queryId; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(UNKNOWN_ERROR); - return std::make_pair(web::http::status_codes::InternalError, response); - } - }) - .then([=](std::pair response_status) { - try - { - message.reply(response_status.first, response_status.second).wait(); - } - catch (const std::exception &ex) - { - std::cerr << "Exception while processing reply: " << ex.what() << std::endl; - }; - }); -} - -web::json::value Server::prepareResponse(const int64_t &queryId, const int k) -{ - web::json::value response = web::json::value::object(); - response[QUERY_ID_KEY] = queryId; - response[K_KEY] = k; - - return response; -} - -template -void Server::parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, - unsigned int &dimensions, unsigned &Ls) -{ - std::cout << body << std::endl; - web::json::value val = web::json::value::parse(body); - web::json::array queryArr = val.at(VECTOR_KEY).as_array(); - queryId = val.has_field(QUERY_ID_KEY) ? val.at(QUERY_ID_KEY).as_number().to_int64() : -1; - Ls = val.has_field(L_KEY) ? val.at(L_KEY).as_number().to_uint32() : DEFAULT_L; - k = val.at(K_KEY).as_integer(); - - if (k <= 0 || k > Ls) - { - throw new std::invalid_argument("Num of expected NN (k) must be greater than zero and less than or " - "equal to Ls."); - } - if (queryArr.size() == 0) - { - throw new std::invalid_argument("Query vector has zero elements."); - } - - dimensions = static_cast(queryArr.size()); - unsigned new_dim = ROUND_UP(dimensions, 8); - diskann::alloc_aligned((void **)&queryVector, new_dim * sizeof(T), 8 * sizeof(T)); - memset(queryVector, 0, new_dim * sizeof(float)); - for (size_t i = 0; i < queryArr.size(); i++) - { - queryVector[i] = (float)queryArr[i].as_double(); - } -} - -template -web::json::value Server::toJsonArray(const std::vector &v, std::function valConverter) -{ - web::json::value rslts = web::json::value::array(); - for (size_t i = 0; i < v.size(); i++) - { - auto jsonVal = valConverter(v[i]); - rslts[i] = jsonVal; - } - return rslts; -} - -web::json::value Server::idsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value idArray = web::json::value::array(); - auto ids = result.get_indices(); - for (size_t i = 0; i < ids.size(); i++) - { - auto idVal = web::json::value::number(ids[i]); - idArray[i] = idVal; - } - std::cout << "Vector size: " << ids.size() << std::endl; - return idArray; -} - -web::json::value Server::distancesToJsonArray(const diskann::SearchResult &result) -{ - web::json::value distArray = web::json::value::array(); - auto distances = result.get_distances(); - for (size_t i = 0; i < distances.size(); i++) - { - distArray[i] = web::json::value::number(distances[i]); - } - return distArray; -} - -web::json::value Server::tagsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value tagArray = web::json::value::array(); - auto tags = result.get_tags(); - for (size_t i = 0; i < tags.size(); i++) - { - tagArray[i] = web::json::value::string(tags[i]); - } - return tagArray; -} - -web::json::value Server::partitionsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value partitionArray = web::json::value::array(); - auto partitions = result.get_partitions(); - for (size_t i = 0; i < partitions.size(); i++) - { - partitionArray[i] = web::json::value::number(partitions[i]); - } - return partitionArray; -} -}; // namespace diskann \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/scratch.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/scratch.cpp deleted file mode 100644 index 1f8a34b..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/scratch.cpp +++ /dev/null @@ -1,182 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include -#include - -#include "scratch.h" -#include "pq_scratch.h" - -namespace diskann -{ -// -// Functions to manage scratch space for in-memory index based search -// -template -InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, - size_t aligned_dim, size_t alignment_factor, bool init_pq_scratch) - : _L(0), _R(r), _maxc(maxc) -{ - if (search_l == 0 || indexing_l == 0 || r == 0 || dim == 0) - { - std::stringstream ss; - ss << "In InMemQueryScratch, one of search_l = " << search_l << ", indexing_l = " << indexing_l - << ", dim = " << dim << " or r = " << r << " is zero." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - alloc_aligned(((void **)&this->_aligned_query_T), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); - memset(this->_aligned_query_T, 0, aligned_dim * sizeof(T)); - - if (init_pq_scratch) - this->_pq_scratch = new PQScratch(defaults::MAX_GRAPH_DEGREE, aligned_dim); - else - this->_pq_scratch = nullptr; - - _occlude_factor.reserve(maxc); - _inserted_into_pool_bs = new boost::dynamic_bitset<>(); - _id_scratch.reserve((size_t)std::ceil(1.5 * defaults::GRAPH_SLACK_FACTOR * _R)); - _dist_scratch.reserve((size_t)std::ceil(1.5 * defaults::GRAPH_SLACK_FACTOR * _R)); - - resize_for_new_L(std::max(search_l, indexing_l)); -} - -template void InMemQueryScratch::clear() -{ - _pool.clear(); - _best_l_nodes.clear(); - _occlude_factor.clear(); - - _inserted_into_pool_rs.clear(); - _inserted_into_pool_bs->reset(); - - _id_scratch.clear(); - _dist_scratch.clear(); - - _expanded_nodes_set.clear(); - _expanded_nghrs_vec.clear(); - _occlude_list_output.clear(); -} - -template void InMemQueryScratch::resize_for_new_L(uint32_t new_l) -{ - if (new_l > _L) - { - _L = new_l; - _pool.reserve(3 * _L + _R); - _best_l_nodes.reserve(_L); - - _inserted_into_pool_rs.reserve(20 * _L); - } -} - -template InMemQueryScratch::~InMemQueryScratch() -{ - if (this->_aligned_query_T != nullptr) - { - aligned_free(this->_aligned_query_T); - this->_aligned_query_T = nullptr; - } - - delete this->_pq_scratch; - delete _inserted_into_pool_bs; -} - -// -// Functions to manage scratch space for SSD based search -// -template void SSDQueryScratch::reset() -{ - sector_idx = 0; - visited.clear(); - retset.clear(); - full_retset.clear(); -} - -template SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, size_t visited_reserve) -{ - size_t coord_alloc_size = ROUND_UP(sizeof(T) * aligned_dim, 256); - - diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); - diskann::alloc_aligned((void **)§or_scratch, defaults::MAX_N_SECTOR_READS * defaults::SECTOR_LEN, - defaults::SECTOR_LEN); - diskann::alloc_aligned((void **)&this->_aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); - - this->_pq_scratch = new PQScratch(defaults::MAX_GRAPH_DEGREE, aligned_dim); - - memset(coord_scratch, 0, coord_alloc_size); - memset(this->_aligned_query_T, 0, aligned_dim * sizeof(T)); - - visited.reserve(visited_reserve); - full_retset.reserve(visited_reserve); -} - -template SSDQueryScratch::~SSDQueryScratch() -{ - diskann::aligned_free((void *)coord_scratch); - diskann::aligned_free((void *)sector_scratch); - diskann::aligned_free((void *)this->_aligned_query_T); - - delete this->_pq_scratch; -} - -template -SSDThreadData::SSDThreadData(size_t aligned_dim, size_t visited_reserve) : scratch(aligned_dim, visited_reserve) -{ -} - -template void SSDThreadData::clear() -{ - scratch.reset(); -} - -template PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) -{ - diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, - (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); - diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); - - memset(aligned_query_float, 0, aligned_dim * sizeof(float)); - memset(rotated_query, 0, aligned_dim * sizeof(float)); -} - -template PQScratch::~PQScratch() -{ - diskann::aligned_free((void *)aligned_pq_coord_scratch); - diskann::aligned_free((void *)aligned_pqtable_dist_scratch); - diskann::aligned_free((void *)aligned_dist_scratch); - diskann::aligned_free((void *)aligned_query_float); - diskann::aligned_free((void *)rotated_query); -} - -template void PQScratch::initialize(size_t dim, const T *query, const float norm) -{ - for (size_t d = 0; d < dim; ++d) - { - if (norm != 1.0f) - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; - else - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); - } -} - -template DISKANN_DLLEXPORT class InMemQueryScratch; -template DISKANN_DLLEXPORT class InMemQueryScratch; -template DISKANN_DLLEXPORT class InMemQueryScratch; - -template DISKANN_DLLEXPORT class SSDQueryScratch; -template DISKANN_DLLEXPORT class SSDQueryScratch; -template DISKANN_DLLEXPORT class SSDQueryScratch; - -template DISKANN_DLLEXPORT class PQScratch; -template DISKANN_DLLEXPORT class PQScratch; -template DISKANN_DLLEXPORT class PQScratch; - -template DISKANN_DLLEXPORT class SSDThreadData; -template DISKANN_DLLEXPORT class SSDThreadData; -template DISKANN_DLLEXPORT class SSDThreadData; - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/utils.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/utils.cpp deleted file mode 100644 index 3773cda..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/utils.cpp +++ /dev/null @@ -1,477 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include "utils.h" - -#include - -#ifdef EXEC_ENV_OLS -#include "aligned_file_reader.h" -#endif - -const uint32_t MAX_REQUEST_SIZE = 1024 * 1024 * 1024; // 64MB -const uint32_t MAX_SIMULTANEOUS_READ_REQUESTS = 128; - -#ifdef _WINDOWS -#include - -// Taken from: -// https://insufficientlycomplicated.wordpress.com/2011/11/07/detecting-intel-advanced-vector-extensions-avx-in-visual-studio/ -bool cpuHasAvxSupport() -{ - bool avxSupported = false; - - // Checking for AVX requires 3 things: - // 1) CPUID indicates that the OS uses XSAVE and XRSTORE - // instructions (allowing saving YMM registers on context - // switch) - // 2) CPUID indicates support for AVX - // 3) XGETBV indicates the AVX registers will be saved and - // restored on context switch - // - // Note that XGETBV is only available on 686 or later CPUs, so - // the instruction needs to be conditionally run. - int cpuInfo[4]; - __cpuid(cpuInfo, 1); - - bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27) || false; - bool cpuAVXSuport = cpuInfo[2] & (1 << 28) || false; - - if (osUsesXSAVE_XRSTORE && cpuAVXSuport) - { - // Check if the OS will save the YMM registers - unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK); - avxSupported = (xcrFeatureMask & 0x6) || false; - } - - return avxSupported; -} - -bool cpuHasAvx2Support() -{ - int cpuInfo[4]; - __cpuid(cpuInfo, 0); - int n = cpuInfo[0]; - if (n >= 7) - { - __cpuidex(cpuInfo, 7, 0); - static int avx2Mask = 0x20; - return (cpuInfo[1] & avx2Mask) > 0; - } - return false; -} - -bool AvxSupportedCPU = cpuHasAvxSupport(); -bool Avx2SupportedCPU = cpuHasAvx2Support(); - -#else - -bool Avx2SupportedCPU = true; -bool AvxSupportedCPU = false; -#endif - -namespace diskann -{ - -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, size_t npts, size_t ndims) -{ - readr.read((char *)read_buf, npts * ndims * sizeof(float)); - uint32_t ndims_u32 = (uint32_t)ndims; -#pragma omp parallel for - for (int64_t i = 0; i < (int64_t)npts; i++) - { - float norm_pt = std::numeric_limits::epsilon(); - for (uint32_t dim = 0; dim < ndims_u32; dim++) - { - norm_pt += *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); - } - norm_pt = std::sqrt(norm_pt); - for (uint32_t dim = 0; dim < ndims_u32; dim++) - { - *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; - } - } - writr.write((char *)read_buf, npts * ndims * sizeof(float)); -} - -void normalize_data_file(const std::string &inFileName, const std::string &outFileName) -{ - std::ifstream readr(inFileName, std::ios::binary); - std::ofstream writr(outFileName, std::ios::binary); - - int npts_s32, ndims_s32; - readr.read((char *)&npts_s32, sizeof(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - - writr.write((char *)&npts_s32, sizeof(int32_t)); - writr.write((char *)&ndims_s32, sizeof(int32_t)); - - size_t npts = (size_t)npts_s32; - size_t ndims = (size_t)ndims_s32; - diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName << std::endl; - diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - diskann::cout << "# blks: " << nblks << std::endl; - - float *read_buf = new float[npts * ndims]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writr, readr, read_buf, cblk_size, ndims); - } - delete[] read_buf; - - diskann::cout << "Wrote normalized points to file: " << outFileName << std::endl; -} - -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, - uint32_t *our_results, uint32_t dim_or, uint32_t recall_at) -{ - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) - { - gt.clear(); - res.clear(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *res_vec = our_results + dim_or * i; - size_t tie_breaker = recall_at; - if (gs_dist != nullptr) - { - tie_breaker = recall_at - 1; - float *gt_dist_vec = gs_dist + dim_gs * i; - while (tie_breaker < dim_gs && gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) - tie_breaker++; - } - - gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, - res_vec + recall_at); // change to recall_at for recall k@k - // or dim_or for k@dim_or - uint32_t cur_recall = 0; - for (auto &v : gt) - { - if (res.find(v) != res.end()) - { - cur_recall++; - } - } - total_recall += cur_recall; - } - return total_recall / (num_queries) * (100.0 / recall_at); -} - -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, - uint32_t *our_results, uint32_t dim_or, uint32_t recall_at, - const tsl::robin_set &active_tags) -{ - double total_recall = 0; - std::set gt, res; - bool printed = false; - for (size_t i = 0; i < num_queries; i++) - { - gt.clear(); - res.clear(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *res_vec = our_results + dim_or * i; - size_t tie_breaker = recall_at; - uint32_t active_points_count = 0; - uint32_t cur_counter = 0; - while (active_points_count < recall_at && cur_counter < dim_gs) - { - if (active_tags.find(*(gt_vec + cur_counter)) != active_tags.end()) - { - active_points_count++; - } - cur_counter++; - } - if (active_tags.empty()) - cur_counter = recall_at; - - if ((active_points_count < recall_at && !active_tags.empty()) && !printed) - { - diskann::cout << "Warning: Couldn't find enough closest neighbors " << active_points_count << "/" - << recall_at - << " from " - "truthset for query # " - << i << ". Will result in under-reported value of recall." << std::endl; - printed = true; - } - if (gs_dist != nullptr) - { - tie_breaker = cur_counter - 1; - float *gt_dist_vec = gs_dist + dim_gs * i; - while (tie_breaker < dim_gs && gt_dist_vec[tie_breaker] == gt_dist_vec[cur_counter - 1]) - tie_breaker++; - } - - gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, res_vec + recall_at); - uint32_t cur_recall = 0; - for (auto &v : res) - { - if (gt.find(v) != gt.end()) - { - cur_recall++; - } - } - total_recall += cur_recall; - } - return ((double)(total_recall / (num_queries))) * ((double)(100.0 / recall_at)); -} - -double calculate_range_search_recall(uint32_t num_queries, std::vector> &groundtruth, - std::vector> &our_results) -{ - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) - { - gt.clear(); - res.clear(); - - gt.insert(groundtruth[i].begin(), groundtruth[i].end()); - res.insert(our_results[i].begin(), our_results[i].end()); - uint32_t cur_recall = 0; - for (auto &v : gt) - { - if (res.find(v) != res.end()) - { - cur_recall++; - } - } - if (gt.size() != 0) - total_recall += ((100.0 * cur_recall) / gt.size()); - else - total_recall += 100; - } - return total_recall / (num_queries); -} - -#ifdef EXEC_ENV_OLS -void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset) -{ - std::vector readReqs; - AlignedRead readReq; - uint32_t buf[2]; // npts/ndim are uint32_ts. - - readReq.buf = buf; - readReq.offset = offset; - readReq.len = 2 * sizeof(uint32_t); - readReqs.push_back(readReq); - - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); // synchronous - if ((*(ctx.m_pRequestsStatus))[0] == IOContext::READ_SUCCESS) - { - npts = buf[0]; - ndim = buf[1]; - diskann::cout << "File has: " << npts << " points, " << ndim << " dimensions at offset: " << offset - << std::endl; - } - else - { - std::stringstream str; - str << "Could not read binary metadata from index file at offset: " << offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset) -{ - // Code assumes that the reader is already setup correctly. - get_bin_metadata(reader, npts, ndim, offset); - data = new T[npts * ndim]; - - size_t data_size = npts * ndim * sizeof(T); - size_t write_offset = 0; - size_t read_start = offset + 2 * sizeof(uint32_t); - - // BingAlignedFileReader can only read uint32_t bytes of data. So, - // we limit ourselves even more to reading 1GB at a time. - std::vector readReqs; - while (data_size > 0) - { - AlignedRead readReq; - readReq.buf = data + write_offset; - readReq.offset = read_start + write_offset; - readReq.len = data_size > MAX_REQUEST_SIZE ? MAX_REQUEST_SIZE : data_size; - readReqs.push_back(readReq); - // in the corner case, the loop will not execute - data_size -= readReq.len; - write_offset += readReq.len; - } - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); - for (int i = 0; i < readReqs.size(); i++) - { - // Since we are making sync calls, no request will be in the - // READ_WAIT state. - if ((*(ctx.m_pRequestsStatus))[i] != IOContext::READ_SUCCESS) - { - std::stringstream str; - str << "Could not read binary data from index file at offset: " << readReqs[i].offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - } -} -template -void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, size_t offset) -{ - T *ptr = nullptr; - load_bin(reader, ptr, npts, ndim, offset); - data.reset(ptr); -} - -template -void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, - const size_t &rounded_dim, size_t offset) -{ - if (data == nullptr) - { - diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file()", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - size_t pts, dim; - get_bin_metadata(reader, pts, dim, offset); - - if (ndim != dim || npts != pts) - { - std::stringstream ss; - ss << "Either file dimension: " << dim << " is != passed dimension: " << ndim << " or file #pts: " << pts - << " is != passed #pts: " << npts << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // Instead of reading one point of ndim size and setting (rounded_dim - dim) - // values to zero We'll set everything to zero and read in chunks of data at - // the appropriate locations. - size_t read_offset = offset + 2 * sizeof(uint32_t); - memset(data, 0, npts * rounded_dim * sizeof(T)); - int i = 0; - std::vector read_requests; - - while (i < npts) - { - int j = 0; - read_requests.clear(); - while (j < MAX_SIMULTANEOUS_READ_REQUESTS && i < npts) - { - AlignedRead read_req; - read_req.buf = data + i * rounded_dim; - read_req.len = dim * sizeof(T); - read_req.offset = read_offset + i * dim * sizeof(T); - read_requests.push_back(read_req); - i++; - j++; - } - IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - for (int k = 0; k < read_requests.size(); k++) - { - if ((*ctx.m_pRequestsStatus)[k] != IOContext::READ_SUCCESS) - { - throw diskann::ANNException("Load data from file using AlignedReader failed.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - } - } -} - -// Unlike load_bin, assumes that data is already allocated 'size' entries -template void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset) -{ - if (data == nullptr) - { - throw diskann::ANNException("read_array requires an allocated buffer.", -1); - } - - if (size * sizeof(T) > MAX_REQUEST_SIZE) - { - std::stringstream ss; - ss << "Cannot read more than " << MAX_REQUEST_SIZE << " bytes. Current request size: " << std::to_string(size) - << " sizeof(T): " << sizeof(T) << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - std::vector read_requests; - AlignedRead read_req; - read_req.buf = data; - read_req.len = size * sizeof(T); - read_req.offset = offset; - read_requests.push_back(read_req); - IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - - if ((*(ctx.m_pRequestsStatus))[0] != IOContext::READ_SUCCESS) - { - std::stringstream ss; - ss << "Failed to read_array() of size: " << size * sizeof(T) << " at offset: " << offset << " from reader. " - << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template void read_value(AlignedFileReader &reader, T &value, size_t offset) -{ - read_array(reader, &value, 1, offset); -} - -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); - -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint8_t *&data, size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int64_t *&data, size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint64_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint32_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int32_t *&data, size_t &npts, size_t &ndim, - size_t offset); - -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, uint8_t *&data, - size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, int8_t *&data, - size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, float *&data, - size_t &npts, size_t &dim, const size_t &rounded_dim, - size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, char *data, size_t size, size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint8_t *data, size_t size, - size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, int8_t *data, size_t size, size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint32_t *data, size_t size, - size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, float *data, size_t size, size_t offset); - -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint8_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, int8_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, float &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint32_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint64_t &value, size_t offset); - -#endif - -} // namespace diskann diff --git a/packages/leann-backend-diskann/third_party/DiskANN/src/windows_aligned_file_reader.cpp b/packages/leann-backend-diskann/third_party/DiskANN/src/windows_aligned_file_reader.cpp deleted file mode 100644 index 3650b92..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/src/windows_aligned_file_reader.cpp +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#ifdef _WINDOWS -#ifndef USE_BING_INFRA -#include "windows_aligned_file_reader.h" -#include -#include "utils.h" -#include - -#define SECTOR_LEN 4096 - -void WindowsAlignedFileReader::open(const std::string &fname) -{ -#ifdef UNICODE - m_filename = std::wstring(fname.begin(), fname.end()); -#else - m_filename = fname; -#endif - - this->register_thread(); -} - -void WindowsAlignedFileReader::close() -{ - for (auto &k_v : ctx_map) - { - IOContext ctx = ctx_map[k_v.first]; - CloseHandle(ctx.fhandle); - } -} - -void WindowsAlignedFileReader::register_thread() -{ - std::unique_lock lk(this->ctx_mut); - if (this->ctx_map.find(std::this_thread::get_id()) != ctx_map.end()) - { - diskann::cout << "Warning:: Duplicate registration for thread_id : " << std::this_thread::get_id() << std::endl; - } - - IOContext ctx; - ctx.fhandle = CreateFile( - m_filename.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, - FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, NULL); - if (ctx.fhandle == INVALID_HANDLE_VALUE) - { - const size_t c_max_filepath_len = 256; - size_t actual_len = 0; - char filePath[c_max_filepath_len]; - if (wcstombs_s(&actual_len, filePath, c_max_filepath_len, m_filename.c_str(), m_filename.length()) == 0) - { - diskann::cout << "Error opening " << filePath << " -- error=" << GetLastError() << std::endl; - } - else - { - diskann::cout << "Error converting wchar to char -- error=" << GetLastError() << std::endl; - } - } - - // create IOCompletionPort - ctx.iocp = CreateIoCompletionPort(ctx.fhandle, ctx.iocp, 0, 0); - - // create MAX_DEPTH # of reqs - for (uint64_t i = 0; i < MAX_IO_DEPTH; i++) - { - OVERLAPPED os; - memset(&os, 0, sizeof(OVERLAPPED)); - // os.hEvent = CreateEventA(NULL, TRUE, FALSE, NULL); - ctx.reqs.push_back(os); - } - this->ctx_map.insert(std::make_pair(std::this_thread::get_id(), ctx)); -} - -IOContext &WindowsAlignedFileReader::get_ctx() -{ - std::unique_lock lk(this->ctx_mut); - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) - { - std::stringstream stream; - stream << "unable to find IOContext for thread_id : " << std::this_thread::get_id() << "\n"; - throw diskann::ANNException(stream.str(), -2, __FUNCSIG__, __FILE__, __LINE__); - } - IOContext &ctx = ctx_map[std::this_thread::get_id()]; - lk.unlock(); - return ctx; -} - -void WindowsAlignedFileReader::read(std::vector &read_reqs, IOContext &ctx, bool async) -{ - using namespace std::chrono_literals; - // execute each request sequentially - size_t n_reqs = read_reqs.size(); - uint64_t n_batches = ROUND_UP(n_reqs, MAX_IO_DEPTH) / MAX_IO_DEPTH; - for (uint64_t i = 0; i < n_batches; i++) - { - // reset all OVERLAPPED objects - for (auto &os : ctx.reqs) - { - // HANDLE evt = os.hEvent; - memset(&os, 0, sizeof(os)); - // os.hEvent = evt; - - /* - if (ResetEvent(os.hEvent) == 0) { - diskann::cerr << "ResetEvent failed" << std::endl; - exit(-3); - } - */ - } - - // batch start/end - uint64_t batch_start = MAX_IO_DEPTH * i; - uint64_t batch_size = std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); - - // fill OVERLAPPED and issue them - for (uint64_t j = 0; j < batch_size; j++) - { - AlignedRead &req = read_reqs[batch_start + j]; - OVERLAPPED &os = ctx.reqs[j]; - - uint64_t offset = req.offset; - uint64_t nbytes = req.len; - char *read_buf = (char *)req.buf; - assert(IS_ALIGNED(read_buf, SECTOR_LEN)); - assert(IS_ALIGNED(offset, SECTOR_LEN)); - assert(IS_ALIGNED(nbytes, SECTOR_LEN)); - - // fill in OVERLAPPED struct - os.Offset = offset & 0xffffffff; - os.OffsetHigh = (offset >> 32); - - BOOL ret = ReadFile(ctx.fhandle, read_buf, (DWORD)nbytes, NULL, &os); - if (ret == FALSE) - { - auto error = GetLastError(); - if (error != ERROR_IO_PENDING) - { - diskann::cerr << "Error queuing IO -- " << error << "\n"; - } - } - else - { - diskann::cerr << "Error queueing IO -- ReadFile returned TRUE" << std::endl; - } - } - DWORD n_read = 0; - uint64_t n_complete = 0; - ULONG_PTR completion_key = 0; - OVERLAPPED *lp_os; - while (n_complete < batch_size) - { - if (GetQueuedCompletionStatus(ctx.iocp, &n_read, &completion_key, &lp_os, INFINITE) != 0) - { - // successfully dequeued a completed I/O - n_complete++; - } - else - { - // failed to dequeue OR dequeued failed I/O - if (lp_os == NULL) - { - DWORD error = GetLastError(); - if (error != WAIT_TIMEOUT) - { - diskann::cerr << "GetQueuedCompletionStatus() failed " - "with error = " - << error << std::endl; - throw diskann::ANNException("GetQueuedCompletionStatus failed with error: ", error, __FUNCSIG__, - __FILE__, __LINE__); - } - // no completion packet dequeued ==> sleep for 5us and try - // again - std::this_thread::sleep_for(5us); - } - else - { - // completion packet for failed IO dequeued - auto op_idx = lp_os - ctx.reqs.data(); - std::stringstream stream; - stream << "I/O failed , offset: " << read_reqs[op_idx].offset - << "with error code: " << GetLastError() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - } - } -} -#endif -#endif diff --git a/packages/leann-backend-diskann/third_party/DiskANN/tests/CMakeLists.txt b/packages/leann-backend-diskann/third_party/DiskANN/tests/CMakeLists.txt deleted file mode 100644 index 6af8405..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/tests/CMakeLists.txt +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT license. - -set(CMAKE_COMPILE_WARNING_AS_ERROR ON) - -find_package(Boost COMPONENTS unit_test_framework) - -# For Windows, fall back to nuget version if find_package didn't find it. -if (MSVC AND NOT Boost_FOUND) - set(DISKANN_BOOST_INCLUDE "${DISKANN_MSVC_PACKAGES}/boost/lib/native/include") - # Multi-threaded static library. - set(UNIT_TEST_FRAMEWORK_LIB_PATTERN "${DISKANN_MSVC_PACKAGES}/boost_unit_test_framework-vc${MSVC_TOOLSET_VERSION}/lib/native/libboost_unit_test_framework-vc${MSVC_TOOLSET_VERSION}-mt-x64-*.lib") - file(GLOB DISKANN_BOOST_UNIT_TEST_FRAMEWORK_LIB ${UNIT_TEST_FRAMEWORK_LIB_PATTERN}) - - set(UNIT_TEST_FRAMEWORK_DLIB_PATTERN "${DISKANN_MSVC_PACKAGES}/boost_unit_test_framework-vc${MSVC_TOOLSET_VERSION}/lib/native/libboost_unit_test_framework-vc${MSVC_TOOLSET_VERSION}-mt-gd-x64-*.lib") - file(GLOB DISKANN_BOOST_UNIT_TEST_FRAMEWORK_DLIB ${UNIT_TEST_FRAMEWORK_DLIB_PATTERN}) - - if (EXISTS ${DISKANN_BOOST_INCLUDE} AND EXISTS ${DISKANN_BOOST_UNIT_TEST_FRAMEWORK_LIB} AND EXISTS ${DISKANN_BOOST_UNIT_TEST_FRAMEWORK_DLIB}) - set(Boost_FOUND ON) - set(Boost_INCLUDE_DIR ${DISKANN_BOOST_INCLUDE}) - add_library(Boost::unit_test_framework STATIC IMPORTED) - set_target_properties(Boost::unit_test_framework PROPERTIES IMPORTED_LOCATION_RELEASE "${DISKANN_BOOST_UNIT_TEST_FRAMEWORK_LIB}") - set_target_properties(Boost::unit_test_framework PROPERTIES IMPORTED_LOCATION_DEBUG "${DISKANN_BOOST_UNIT_TEST_FRAMEWORK_DLIB}") - message(STATUS "Falling back to using Boost from the nuget package") - else() - message(WARNING "Couldn't find Boost. Was looking for ${DISKANN_BOOST_INCLUDE} and ${UNIT_TEST_FRAMEWORK_LIB_PATTERN}") - endif() -endif() - -if (NOT Boost_FOUND) - message(FATAL_ERROR "Couldn't find Boost dependency") -endif() - - -set(DISKANN_UNIT_TEST_SOURCES main.cpp index_write_parameters_builder_tests.cpp) - -add_executable(${PROJECT_NAME}_unit_tests ${DISKANN_SOURCES} ${DISKANN_UNIT_TEST_SOURCES}) -target_link_libraries(${PROJECT_NAME}_unit_tests ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::unit_test_framework) - -add_test(NAME ${PROJECT_NAME}_unit_tests COMMAND ${PROJECT_NAME}_unit_tests) - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/tests/README.md b/packages/leann-backend-diskann/third_party/DiskANN/tests/README.md deleted file mode 100644 index 113c998..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/tests/README.md +++ /dev/null @@ -1,11 +0,0 @@ -# Unit Test project - -This unit test project is based on the [boost unit test framework](https://www.boost.org/doc/libs/1_78_0/libs/test/doc/html/index.html). Below are the simple steps to add new unit test, you could find more usage from the [boost unit test document](https://www.boost.org/doc/libs/1_78_0/libs/test/doc/html/index.html). - -## How to add unit test - -- Create new [BOOST_AUTO_TEST_SUITE](https://www.boost.org/doc/libs/1_78_0/libs/test/doc/html/boost_test/utf_reference/test_org_reference/test_org_boost_auto_test_suite.html) for each class in an individual cpp file - -- Add [BOOST_AUTO_TEST_CASE](https://www.boost.org/doc/libs/1_78_0/libs/test/doc/html/boost_test/utf_reference/test_org_reference/test_org_boost_auto_test_case.html) for each test case in the [BOOST_AUTO_TEST_SUITE](https://www.boost.org/doc/libs/1_78_0/libs/test/doc/html/boost_test/utf_reference/test_org_reference/test_org_boost_auto_test_suite.html) - -- Update the [CMakeLists.txt](CMakeLists.txt) file to add the new cpp file to the test project \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/tests/index_write_parameters_builder_tests.cpp b/packages/leann-backend-diskann/third_party/DiskANN/tests/index_write_parameters_builder_tests.cpp deleted file mode 100644 index 0aa798d..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/tests/index_write_parameters_builder_tests.cpp +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#include - -#include "parameters.h" - -BOOST_AUTO_TEST_SUITE(IndexWriteParametersBuilder_tests) - -BOOST_AUTO_TEST_CASE(test_build) -{ - uint32_t search_list_size = rand(); - uint32_t max_degree = rand(); - float alpha = (float)rand(); - uint32_t filter_list_size = rand(); - uint32_t max_occlusion_size = rand(); - bool saturate_graph = true; - - diskann::IndexWriteParametersBuilder builder(search_list_size, max_degree); - - builder.with_alpha(alpha) - .with_filter_list_size(filter_list_size) - .with_max_occlusion_size(max_occlusion_size) - .with_num_threads(0) - .with_saturate_graph(saturate_graph); - - { - auto parameters = builder.build(); - - BOOST_TEST(search_list_size == parameters.search_list_size); - BOOST_TEST(max_degree == parameters.max_degree); - BOOST_TEST(alpha == parameters.alpha); - BOOST_TEST(filter_list_size == parameters.filter_list_size); - BOOST_TEST(max_occlusion_size == parameters.max_occlusion_size); - BOOST_TEST(saturate_graph == parameters.saturate_graph); - - BOOST_TEST(parameters.num_threads > (uint32_t)0); - } - - { - uint32_t num_threads = rand() + 1; - saturate_graph = false; - builder.with_num_threads(num_threads).with_saturate_graph(saturate_graph); - - auto parameters = builder.build(); - - BOOST_TEST(search_list_size == parameters.search_list_size); - BOOST_TEST(max_degree == parameters.max_degree); - BOOST_TEST(alpha == parameters.alpha); - BOOST_TEST(filter_list_size == parameters.filter_list_size); - BOOST_TEST(max_occlusion_size == parameters.max_occlusion_size); - BOOST_TEST(saturate_graph == parameters.saturate_graph); - - BOOST_TEST(num_threads == parameters.num_threads); - } -} - -BOOST_AUTO_TEST_SUITE_END() diff --git a/packages/leann-backend-diskann/third_party/DiskANN/tests/main.cpp b/packages/leann-backend-diskann/third_party/DiskANN/tests/main.cpp deleted file mode 100644 index 53440a1..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/tests/main.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -#define BOOST_TEST_MODULE diskann_unit_tests - -#include diff --git a/packages/leann-backend-diskann/third_party/DiskANN/windows/packages.config.in b/packages/leann-backend-diskann/third_party/DiskANN/windows/packages.config.in deleted file mode 100644 index f8eecf0..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/windows/packages.config.in +++ /dev/null @@ -1,11 +0,0 @@ - - - - - - - - - - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/windows/packages_restapi.config.in b/packages/leann-backend-diskann/third_party/DiskANN/windows/packages_restapi.config.in deleted file mode 100644 index 6d1a60c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/windows/packages_restapi.config.in +++ /dev/null @@ -1,4 +0,0 @@ - - - - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/SSD_index.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/SSD_index.md deleted file mode 100644 index 3144528..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/SSD_index.md +++ /dev/null @@ -1,74 +0,0 @@ -**Usage for SSD-based indices** -=============================== - -To generate an SSD-friendly index, use the `apps/build_disk_index` program. ----------------------------------------------------------------------------- - -The arguments are as follows: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **--dist_fn**: Three distance functions are supported: cosine distance, minimum Euclidean distance (l2) and maximum inner product (mips). -3. **--data_file**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as an integer. The next 4 bytes represent the dimension of data as an integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. `sizeof(T)` is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **--index_path_prefix**: the index will span a few files, all beginning with the specified prefix path. For example, if you provide `~/index_test` as the prefix path, build generates files such as `~/index_test_pq_pivots.bin, ~/index_test_pq_compressed.bin, ~/index_test_disk.index, ...`. There may be between 8 and 10 files generated with this prefix depending on how the index is constructed. -5. **-R (--max_degree)** (default is 64): the degree of the graph index, typically between 60 and 150. Larger R will result in larger indices and longer indexing times, but better search quality. -6. **-L (--Lbuild)** (default is 100): the size of search list during index build. Typical values are between 75 to 200. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Use a value for L value that is at least the value of R unless you need to build indices really quickly and can somewhat compromise on quality. -7. **-B (--search_DRAM_budget)**: bound on the memory footprint of the index at search time in GB. Once built, the index will use up only the specified RAM limit, the rest will reside on disk. This will dictate how aggressively we compress the data vectors to store in memory. Larger will yield better performance at search time. For an n point index, to use b byte PQ compressed representation in memory, use `B = ((n * b) / 2^30 + (250000*(4*R + sizeof(T)*ndim)) / 2^30)`. The second term in the summation is to allow some buffer for caching about 250,000 nodes from the graph in memory while serving. If you are not sure about this term, add 0.25GB to the first term. -8. **-M (--build_DRAM_budget)**: Limit on the memory allowed for building the index in GB. If you specify a value less than what is required to build the index in one pass, the index is built using a divide and conquer approach so that sub-graphs will fit in the RAM budget. The sub-graphs are overlayed to build the overall index. This approach can be upto 1.5 times slower than building the index in one shot. Allocate as much memory as your RAM allows. -9. **-T (--num_threads)** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -10. **--PQ_disk_bytes** (default is 0): Use 0 to store uncompressed data on SSD. This allows the index to asymptote to 100% recall. If your vectors are too large to store in SSD, this parameter provides the option to compress the vectors using PQ for storing on SSD. This will trade off recall. You would also want this to be greater than the number of bytes used for the PQ compressed data stored in-memory -11. **--build_PQ_bytes** (default is 0): Set to a positive value less than the dimensionality of the data to enable faster index build with PQ based distance comparisons. -12. **--use_opq**: use the flag to use OPQ rather than PQ compression. OPQ is more space efficient for some high dimensional datasets, but also needs a bit more build time. - -To search the SSD-index, use the `apps/search_disk_index` program. -------------------------------------------------------------------- - -The arguments are as follows: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. Use the same data type as in arg (1) above used in building the index. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). Use the same distance as in arg (2) above used in building the index. -3. **--index_path_prefix**: same as the prefix used in building the index (see arg 4 above). -4. **--num_nodes_to_cache** (default is 0): While serving the index, the entire graph is stored on SSD. For faster search performance, you can cache a few frequently accessed nodes in memory. -5. **-T (--num_threads)** (default is to get_omp_num_procs()): The number of threads used for searching. Threads run in parallel and one thread handles one query at a time. More threads will result in higher aggregate query throughput, but will also use more IOs/second across the system, which may lead to higher per-query latency. So find the balance depending on the maximum number of IOPs supported by the SSD. -6. **-W (--beamwidth)** (default is 2): The beamwidth to be used for search. This is the maximum number of IO requests each query will issue per iteration of search code. Larger beamwidth will result in fewer IO round-trips per query, but might result in slightly higher total number of IO requests to SSD per query. For the highest query throughput with a fixed SSD IOps rating, use `W=1`. For best latency, use `W=4,8` or higher complexity search. Specifying 0 will optimize the beamwidth depending on the number of threads performing search, but will involve some tuning overhead. -7. **--query_file**: The queries to be searched on in same binary file format as the data file in arg (2) above. The query file must be the same type as argument (1). -8. **--gt_file**: The ground truth file for the queries in arg (7) and data file used in index construction. The binary file must start with *n*, the number of queries (4 bytes), followed by *d*, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the d closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. The groundtruth file, if not available, can be calculated using the program `apps/utils/compute_groundtruth`. Use "null" if you do not have this file and if you do not want to compute recall. -9. **K**: search for *K* neighbors and measure *K*-recall@*K*, meaning the intersection between the retrieved top-*K* nearest neighbors and ground truth *K* nearest neighbors. -10. **result_output_prefix**: Search results will be stored in files with specified prefix, in bin format. -11. **-L (--search_list)**: A list of search_list sizes to perform search with. Larger parameters will result in slower latencies, but higher accuracies. Must be at least the value of *K* in arg (9). - - -Example with BIGANN: --------------------- - -This example demonstrates the use of the commands above on a 100K slice of the [BIGANN dataset](http://corpus-texmex.irisa.fr/) with 128 dimensional SIFT descriptors applied to images. - -Download the base and query set and convert the data to binary format -```bash -mkdir -p DiskANN/build/data && cd DiskANN/build/data -wget ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz -tar -xf sift.tar.gz -cd .. -./apps/utils/fvecs_to_bin float data/sift/sift_learn.fvecs data/sift/sift_learn.fbin -./apps/utils/fvecs_to_bin float data/sift/sift_query.fvecs data/sift/sift_query.fbin -``` - -Now build and search the index and measure the recall using ground truth computed using brutefoce. -```bash -./apps/utils/compute_groundtruth --data_type float --dist_fn l2 --base_file data/sift/sift_learn.fbin --query_file data/sift/sift_query.fbin --gt_file data/sift/sift_query_learn_gt100 --K 100 -# Using 0.003GB search memory budget for 100K vectors implies 32 byte PQ compression -./apps/build_disk_index --data_type float --dist_fn l2 --data_path data/sift/sift_learn.fbin --index_path_prefix data/sift/disk_index_sift_learn_R32_L50_A1.2 -R 32 -L50 -B 0.003 -M 1 - ./apps/search_disk_index --data_type float --dist_fn l2 --index_path_prefix data/sift/disk_index_sift_learn_R32_L50_A1.2 --query_file data/sift/sift_query.fbin --gt_file data/sift/sift_query_learn_gt100 -K 10 -L 10 20 30 40 50 100 --result_path data/sift/res --num_nodes_to_cache 10000 - ``` - -The search might be slower on machine with remote SSDs. The output lists the query throughput, the mean and 99.9pc latency in microseconds and mean number of 4KB IOs to disk for each `L` parameter provided. - -``` - L Beamwidth QPS Mean Latency 99.9 Latency Mean IOs CPU (s) Recall@10 -====================================================================================================================== - 10 2 27723.95 2271.92 4700.00 8.81 40.47 81.79 - 20 2 15369.23 4121.04 7576.00 15.93 61.60 96.42 - 30 2 10335.75 6147.14 11424.00 23.30 74.96 98.78 - 40 2 7684.18 8278.83 14714.00 30.78 94.27 99.40 - 50 2 6421.66 9913.28 16550.00 38.35 116.86 99.63 - 100 2 3337.98 19107.81 29292.00 76.59 226.88 99.91 -``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/dynamic_index.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/dynamic_index.md deleted file mode 100644 index 17c3fb3..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/dynamic_index.md +++ /dev/null @@ -1,187 +0,0 @@ - - -**Usage for dynamic indices** -================================ - -A "dynamic" index refers to an index which supports insertion of new points into a (possibly previously built) index as well as deletions of points. -While eager deletes can be supported by DiskANN, `lazy_deletes` are the preferred method. -A sequence of lazy deletions must be followed by an invocation of the `consolidate_deletes` method that frees up slots in the index and edits the graph to maintain good recall. - - -The program `apps/test_insert_deletes_consolidate` demonstrates this functionality. It allows the user to specify which points from the data file will be used -to initially build the index, which points will be deleted from the index, and which points will be inserted into the index. -Insertions, searches and lazy deletions can be performed concurrently. -Conslolidation of lazy deletes can be performed synchnronously or concurrently with insertions and deletions. -When modifying the index sequentially, the user has the ability to take *snapshots*-- -that is, save the index to memory for every *m* insertions or deletions instead of only at the end of the build. - -The program `apps/test_streaming_scenario` simulates a scenario where the index actively maintains a sliding window of active points from a larger dataset. -The program starts with an index build over the first `active_window` set of points from a data file. -The program then simultaneously inserts newer points drawn from the file and deletes older points from the index -in chunks of `consolidate_interval` points so that the number of active points in the index is approximately `active_window`. -It terminates when the end of data file is reached, and the final index has `active_window + consolidate_interval` number of points. - -The index also supports filters on steaming index, you can use `insert_point` function overloads to either insert points as before or insert points with labels. -Additional options are added to support this in `apps/test_streaming_scenario` and `apps/test_streaming_scenario` please refer to program arguments for more details. - ---- -> Note -* The index does not support mixed points, that is, either all points do not have labels or all points have labels. -* You can search the built filter index (one built with filters) without filters as well. - -> WARNING: Deleting points in case of filtered build may cause the quality of Index to degrade and affect recall. ---- - -`apps/test_insert_deletes_consolidate` to try inserting, lazy deletes and consolidate_delete ---------------------------------------------------------------------------------------------- - -The arguments are as follows: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). -3. **--data_file**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **--index_path_prefix**: The constructed index components will be saved to this path prefix. -5. **-R (--max_degree)** (default is 64): the degree of the graph index, typically between 32 and 150. Larger R will result in larger indices and longer indexing times, but might yield better search quality. -6. **-L (--Lbuild)** (default is 100): the size of search list we maintain during index building. Typical values are between 75 to 400. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Ensure that value of L is at least that of R value unless you need to build indices really quickly and can somewhat compromise on quality. -7. **--alpha** (default is 1.2): A float value between 1.0 and 1.5 which determines the diameter of the graph, which will be approximately *log n* to the base alpha. Typical values are between 1 to 1.5. 1 will yield the sparsest graph, 1.5 will yield denser graphs. -8. **T (--num_threads)** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -9. **--points_to_skip**: number of points to skip from the beginning of the data file. -10. **--max_points_to_insert**: the maximum size of the index. -11. **--beginning_index_size**: how many points to build the initial index with. The number of points inserted dynamically will be max_points_to_insert - beginning_index_size. -12. **--points_per_checkpoint**: when inserting and deleting sequentially, each update is handled in points_per_checkpoint batches. When updating concurrently, insertions are handled in points_per_checkpoint batches but deletions are always processed in a single batch. -13. **--checkpoints_per_snapshot**: when inserting and deleting sequentially, the graph is saved to memory every checkpoints_per_snapshot checkpoints. This is not currently supported for concurrent updates. -14. **--points_to_delete_from_beginning**: how many points to delete from the index, starting in order of insertion. If deletions are concurrent with insertions, points_to_delete_from_beginning cannot be larger than beginning_index_size. -15. **--start_point_norm**: Set the starting node to a random point on a sphere of this radius. A reasonable choice is to set this to the average norm of the data set. Use when starting an index with zero points. -16. **--do_concurrent** (default false): whether to perform conslidate_deletes and other updates concurrently or sequentially. If concurrent is specified, half the threads are used for insertions and half the threads are used for processing deletes. Note that insertions are performed before deletions if this flag is set to false, so in this case is possible to delete more than beginning_index_size points. - -`apps/test_streaming_scenario` to try inserting, lazy deletes and consolidate_delete ---------------------------------------------------------------------------------------------- - -The arguments are as follows: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). -3. **--data_file**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **--index_path_prefix**: The constructed index components will be saved to this path prefix. -5. **-R (--max_degree)** (default is 64): the degree of the graph index, typically between 32 and 150. Larger R will result in larger indices and longer indexing times, but might yield better search quality. -6. **-L (--Lbuild)** (default is 100): the size of search list we maintain during index building. Typical values are between 75 to 400. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Ensure that value of L is at least that of R value unless you need to build indices really quickly and can somewhat compromise on quality. -7. **--alpha** (default is 1.2): A float value between 1.0 and 1.5 which determines the diameter of the graph, which will be approximately *log n* to the base alpha. Typical values are between 1 to 1.5. 1 will yield the sparsest graph, 1.5 will yield denser graphs. -8. **--insert_threads**: number of threads used for inserting points in to the index. -9. **--consolidate_threads**: number of threads used for consolidating deletes to the index. -10. **--max_points_to_insert**: Maximum number of points from the data file to insert in to the index. -11. **--active_window**: Approximate number of points in the index at any point. -12. **--consolidate_interval**: Granularity at which insert and delete functions are called. -13. **--start_point_norm**: Set the starting node to a random point on a sphere of this radius. A reasonable choice is to set this to the average norm of the data stream. - -** To build with filters add these optional parameters. - -14. **--label_file**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of labels corresponding to point `i` in the file passed via `--data_file`. -15. **--FilteredLbuild**: If building a filtered index, we maintain a separate search list from the one provided by `--Lbuild/-L`. -16. **--num_start_points**: number of frozen points in this case should be more then number of unique labels. -17. **--universal_label**: Optionally, the label data may contain a special "universal" label. A point with the universal label can be matched against a query with any label. Note that if a point has the universal label, then the filter data must only have the universal label on the line corresponding. -18. **--label_type**: Optionally, type of label to be use its either uint or short, defaulted to `uint`. - -To search the generated index, use the `apps/search_memory_index` program: ---------------------------------------------------------------------------- - - -The arguments are as follows: - -1. **data_type**: The type of dataset you built the index on. float(32 bit), signed int8 and unsigned uint8 are supported. Use the same data type as in arg (1) above used in building the index. -2. **dist_fn**: There are two distance functions supported: l2 and mips. There is an additional *fast_l2* implementation that could provide faster results for small (about a million-sized) indices. Use the same distance as in arg (2) above used in building the index. -3. **memory_index_path**: index built above in argument (4). -4. **T**: The number of threads used for searching. Threads run in parallel and one thread handles one query at a time. More threads will result in higher aggregate query throughput, but may lead to higher per-query latency, especially if the DRAM bandwidth is a bottleneck. So find the balance depending on throughput and latency required for your application. -5. **query_bin**: The queries to be searched on in same binary file format as the data file (ii) above. The query file must be the same type as in argument (1). -6. **truthset.bin**: The ground truth file for the queries in arg (7) and data file used in index construction. The binary file must start with *n*, the number of queries (4 bytes), followed by *d*, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the d closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. The groundtruth file, if not available, can be calculated using the program `apps/utils/compute_groundtruth`. Use "null" if you do not have this file and if you do not want to compute recall. -7. **K**: search for *K* neighbors and measure *K*-recall@*K*, meaning the intersection between the retrieved top-*K* nearest neighbors and ground truth *K* nearest neighbors. -8. **result_output_prefix**: search results will be stored in files, one per L value (see next arg), with specified prefix, in binary format. -9. **-L (--search_list)**: A list of search_list sizes to perform search with. Larger parameters will result in slower latencies, but higher accuracies. Must be at least the value of *K* in (7). -10. **--dynamic** (default false): whether the index being searched is dynamic or not. -11. **--tags** (default false): whether to search with tags. This should be used if point *i* in the ground truth file does not correspond the point in the *i*th position in the loaded index. - -** to search with filters add these - -12. **--filter_label**: Filter for each query. For each query, a search is performed with this filter. - -Example with BIGANN: --------------------- - -This example demonstrates the use of the commands above on a 100K slice of the [BIGANN dataset](http://corpus-texmex.irisa.fr/) with 128 dimensional SIFT descriptors applied to images. - -Download the base and query set and convert the data to binary format -```bash -mkdir -p DiskANN/build/data && cd DiskANN/build/data -wget ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz -tar -xf sift.tar.gz -cd .. -./apps/utils/fvecs_to_bin float data/sift/sift_learn.fvecs data/sift/sift_learn.fbin -./apps/utils/fvecs_to_bin float data/sift/sift_query.fvecs data/sift/sift_query.fbin -``` - -The example below tests the following scenario: using a file with 100000 points, the index is incrementally constructed point by point. After the first 50000 ponts are inserted, another concurrent job deletes the first 25000 points from the index and consolidates the index (edit the graph and cleans up resources). At the same time an additional 25000 points (i.e. points 50001 to 75000) are concurrently inserted into the index. Note that the index should be built **before** calculating the ground truth, since the memory index returns the slice of the sift100K dataset that was used to build the final graph (that is, points 25001-75000 in the original index). -```bash -type='float' -data='data/sift/sift_learn.fbin' -query='data/sift/sift_query.fbin' -index_prefix='data/sift/index' -result='data/sift/res' -deletes=25000 -inserts=75000 -deletes_after=50000 -pts_per_checkpoint=10000 -begin=0 -thr=64 -index=${index_prefix}.after-concurrent-delete-del${deletes}-${inserts} -gt_file=data/sift/gt100_learn-conc-${deletes}-${inserts} - - ~/DiskANN/build/apps/test_insert_deletes_consolidate --data_type ${type} --dist_fn l2 --data_path ${data} --index_path_prefix ${index_prefix} -R 64 -L 300 --alpha 1.2 -T ${thr} --points_to_skip 0 --max_points_to_insert ${inserts} --beginning_index_size ${begin} --points_per_checkpoint ${pts_per_checkpoint} --checkpoints_per_snapshot 0 --points_to_delete_from_beginning ${deletes} --start_deletes_after ${deletes_after} --do_concurrent true; - - ~/DiskANN/build/apps/utils/compute_groundtruth --data_type ${type} --dist_fn l2 --base_file ${index}.data --query_file ${query} --K 100 --gt_file ${gt_file} --tags_file ${index}.tags - -~/DiskANN/build/apps/search_memory_index --data_type ${type} --dist_fn l2 --index_path_prefix ${index} --result_path ${result} --query_file ${query} --gt_file ${gt_file} -K 10 -L 20 40 60 80 100 -T ${thr} --dynamic true --tags 1 - ``` - - The example below tests the following scenario: using a file with 100000 points, insert 10000 points at a time. After the first 40000 -are inserted, start deleting the first 10000 points while inserting points 40000--50000. Then delete points 10000--20000 while inserting -points 50000--60000 and so until the index is left with points 60000-100000. - - -Generate labels for filtered build like this. Generating 50 unique labels zipf's distributed for 100K point dataset. -``` -~/DiskANN/build/apps/utils/generate_synthetic_labels --num_labels 50 --num_points 100000 --output_file data/zipf_labels_50_100K.txt --distribution_type zipf -``` - -```bash -type='float' -data='data/sift/sift_learn.fbin' -query='data/sift/sift_query.fbin' -index_prefix='data/sift/idx_learn_str' -result='data/sift/res' -ins_thr=16 -cons_thr=16 -inserts=100000 -active=20000 -cons_int=10000 -index=${index_prefix}.after-streaming-act${active}-cons${cons_int}-max${inserts} -gt=data/sift/gt100_learn-act${active}-cons${cons_int}-max${inserts} -filter_label=1 - -## filter options -universal_label = '0' -label_file = 'data/zipf_labels_50_100K.txt' -num_start_points = 50 -gt_filtered= data/sift/gt100_learn-act${active}-cons${cons_int}-max${inserts}_wlabel_${filter_label} - - -# Without Filters (build and search) -./apps/test_streaming_scenario --data_type ${type} --dist_fn l2 --data_path ${data} --index_path_prefix ${index_prefix} -R 64 -L 600 --alpha 1.2 --insert_threads ${ins_thr} --consolidate_threads ${cons_thr} --max_points_to_insert ${inserts} --active_window ${active} --consolidate_interval ${cons_int} --start_point_norm 508; -./apps/utils/compute_groundtruth --data_type ${type} --dist_fn l2 --base_file ${index}.data --query_file ${query} --K 100 --gt_file ${gt} --tags_file ${index}.tags -./apps/search_memory_index --data_type ${type} --dist_fn l2 --index_path_prefix ${index} --result_path ${result} --query_file ${query} --gt_file ${gt} -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 - -# With filters (build and search) - -./apps/test_streaming_scenario --data_type ${type} --num_start_points ${num_start_points} --label_file ${label_file} --universal_label {universal_label} --dist_fn l2 --data_path ${data} --index_path_prefix ${index_prefix} -R 64 -L 600 --alpha 1.2 --insert_threads ${ins_thr} --consolidate_threads ${cons_thr} --max_points_to_insert ${inserts} --active_window ${active} --consolidate_interval ${cons_int} --start_point_norm 508; -./apps/utils/compute_groundtruth_for_filters --data_type ${type} --dist_fn l2 --base_file ${index}.data --query_file ${query} --K 100 --gt_file ${gt_filtered} --label_file ${label_file} --universal_label {universal_label} --filter_label {filter_label} -./apps/search_memory_index --data_type ${type} --filter_label {filter_label} --dist_fn l2 --index_path_prefix ${index} --result_path ${result} --query_file ${query} --gt_file ${gt_filtered} -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 -``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_in_memory.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_in_memory.md deleted file mode 100644 index fe34b80..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_in_memory.md +++ /dev/null @@ -1,126 +0,0 @@ -**Usage for filtered indices** -================================ -## Building a filtered Index -DiskANN provides two algorithms for building an index with filters support: filtered-vamana and stitched-vamana. Here, we describe the parameters for building both. `apps/build_memory_index.cpp` and `apps/build_stitched_index.cpp` are respectively used to build each kind of index. - -### 1. filtered-vamana - -1. **`--data_type`**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **`--dist_fn`**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). -3. **`--data_file`**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **`--index_path_prefix`**: The constructed index components will be saved to this path prefix. -5. **`-R (--max_degree)`** (default is 64): the degree of the graph index, typically between 32 and 150. Larger R will result in larger indices and longer indexing times, but might yield better search quality. -6. **`-L (--Lbuild)`** (default is 100): the size of search list we maintain during index building. Typical values are between 75 to 400. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Ensure that value of L is at least that of R value unless you need to build indices really quickly and can somewhat compromise on quality. Note that this is to be used only for building an unfiltered index. The corresponding search list parameter for a filtered index is managed by `--FilteredLbuild`. -7. **`--alpha`** (default is 1.2): A float value between 1.0 and 1.5 which determines the diameter of the graph, which will be approximately *log n* to the base alpha. Typical values are between 1 to 1.5. 1 will yield the sparsest graph, 1.5 will yield denser graphs. -8. **`-T (--num_threads)`** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -9. **`--build_PQ_bytes`** (default is 0): Set to a positive value less than the dimensionality of the data to enable faster index build with PQ based distance comparisons. Defaults to using full precision vectors for distance comparisons. -10. **`--use_opq`**: use the flag to use OPQ rather than PQ compression. OPQ is more space efficient for some high dimensional datasets, but also needs a bit more build time. -11. **`--label_file`**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of filters corresponding to point `i` in the file passed via `--data_file`. -12. **`--universal_label`**: Optionally, the the filter data may contain a "wild-card" filter corresponding to all filters. This is referred to as a universal label. Note that if a point has the universal label, then the filter data must only have the universal label on the line corresponding to said point. -13. **`--FilteredLbuild`**: If building a filtered index, we maintain a separate search list from the one provided by `--Lbuild`. - -### 2. stitched-vamana -1. **`--data_type`**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **`--data_path`**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -3. **`--index_path_prefix`**: The constructed index components will be saved to this path prefix. -4. **`-R (--max_degree)`** (default is 64): Recall that stitched-vamana first builds a sub-index for each filter. This parameter sets the max degree for each sub-index. -5. **`-L (--Lbuild)`** (default is 100): the size of search list we maintain during sub-index building. Typical values are between 75 to 400. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Ensure that value of L is at least that of R value unless you need to build indices really quickly and can somewhat compromise on quality. -6. **`--alpha`** (default is 1.2): A float value between 1.0 and 1.5 which determines the diameter of the graph, which will be approximately *log n* to the base alpha. Typical values are between 1 to 1.5. 1 will yield the sparsest graph, 1.5 will yield denser graphs. -7. **`-T (--num_threads)`** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -8. **`--label_file`**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of filters corresponding to point `i` in the file passed via `--data_file`. -9. **`--universal_label`**: Optionally, the the filter data may contain a "wild-card" filter corresponding to all filters. This is referred to as a universal label. Note that if a point has the universal label, then the filter data must only have the universal label on the line corresponding to said point. -10. **`--Stitched_R`**: Once all sub-indices are "stitched" together, we prune the resulting graph down to the degree given by this parameter. - -## Computing a groundtruth file for a filtered index -In order to evaluate the performance of our algorithms, we can compare its results (i.e. the top `k` neighbors found for each query) against the results found by an exact nearest neighbor search. We provide the program `apps/utils/compute_groundtruth.cpp` to provide the results for the latter: - -1. **`--data_type`** The type of dataset you built an index with. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **`--dist_fn`**: There are two distance functions supported: l2 and mips. -3. **`--base_file`**: The input data over which to build an index, in .bin format. Corresponds to the `--data_path` argument from above. -4. **`--query_file`**: The queries to be searched on, which are stored in the same .bin format. -5. **`--label_file`**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of filters corresponding to point `i` in the file passed via `--data_file`. -6. **`--filter_label`**: Filter for each query. For each query, a search is performed with this filter. -7. **`--universal_label`**: Corresponds to the universal label passed when building an index with filter support. -8. **`--gt_file`**: File to output results to. The binary file starts with `n`, the number of queries (4 bytes), followed by `d`, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the `d` closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. -9. **`-K`**: The number of nearest neighbors to compute for each query. - - - -## Searching a Filtered Index - -Searching a filtered index uses the `apps/search_memory_index.cpp`: - -1. **`--data_type`**: The type of dataset you built the index on. float(32 bit), signed int8 and unsigned uint8 are supported. Use the same data type as in arg (1) above used in building the index. -2. **`--dist_fn`**: There are two distance functions supported: l2 and mips. There is an additional *fast_l2* implementation that could provide faster results for small (about a million-sized) indices. Use the same distance as in arg (2) above used in building the index. Note that stitched-vamana only supports l2. -3. **`--index_path_prefix`**: index built above in argument (4). -4. **`--result_path`**: search results will be stored in files, one per L value (see last arg), with specified prefix, in binary format. -5. **`-T (--num_threads)`**: The number of threads used for searching. Threads run in parallel and one thread handles one query at a time. More threads will result in higher aggregate query throughput, but may lead to higher per-query latency, especially if the DRAM bandwidth is a bottleneck. So find the balance depending on throughput and latency required for your application. -6. **`--query_file`**: The queries to be searched on in same binary file format as the data file (ii) above. The query file must be the same type as in argument (1). -7. **`--filter_label`**: The filter to be used when searching an index with filters. For each query, a search is performed with this filter. -8. **`--gt_file`**: The ground truth file for the queries and data file used in index construction. Use "null" if you do not have this file and if you do not want to compute recall. Note that if building a filtered index, a special groundtruth must be computed, as described above. -9. **`-K`**: search for *K* neighbors and measure *K*-recall@*K*, meaning the intersection between the retrieved top-*K* nearest neighbors and ground truth *K* nearest neighbors. -10. **`-L (--search_list)`**: A list of search_list sizes to perform search with. Larger parameters will result in slower latencies, but higher accuracies. Must be atleast the value of *K* in (7). - -Example with SIFT10K: --------------------- -We demonstrate how to work through this pipeline using the SIFT10K dataset (http://corpus-texmex.irisa.fr/). Before starting, make sure you have compiled diskANN according to the instructions in the README and can see the following binaries (paths with respect to repository root): -- `build/apps/utils/compute_groundtruth` -- `build/apps/utils/fvecs_to_bin` -- `build/apps/build_memory_index` -- `build/apps/build_stitched_index` -- `build/apps/search_memory_index` - -Now, download the base and query set and convert the data to binary format: -```bash -wget ftp://ftp.irisa.fr/local/texmex/corpus/siftsmall.tar.gz -tar -zxvf siftsmall.tar.gz -build/apps/utils/fvecs_to_bin float siftsmall/siftsmall_base.fvecs siftsmall/siftsmall_base.bin -build/apps/utils/fvecs_to_bin float siftsmall/siftsmall_query.fvecs siftsmall/siftsmall_query.bin -``` - -We now need to make label file for our vectors. For convenience, we've included a synthetic label generator through which we can generate label file as follow -```bash - build/apps/utils/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file ./rand_labels_50_10K.txt --distribution_type zipf -``` -Note : `distribution_type` can be `rand` or `zipf` - -This will genearate label file with 10000 data points with 50 distinct labels, ranging from 1 to 50 assigned using zipf distribution (0 is the universal label). - -Label count for each unique label in the generated label file can be printed with help of following command -```bash - build/apps/utils/stats_label_data.exe --labels_file ./rand_labels_50_10K.txt --universal_label 0 -``` - -Note that neither approach is designed for use with random synthetic labels, which will lead to unpredictable accuracy at search time. - -Now build and search the index and measure the recall using ground truth computed using bruteforce. We search for results with the filter 35. -```bash -build/apps/utils/compute_groundtruth --data_type float --dist_fn l2 --base_file siftsmall/siftsmall_base.bin --query_file siftsmall/siftsmall_query.bin --gt_file siftsmall/siftsmall_gt_35.bin --K 100 --label_file ./rand_labels_50_10K.txt --filter_label 35 --universal_label 0 -build/apps/build_memory_index --data_type float --dist_fn l2 --data_path siftsmall/siftsmall_base.bin --index_path_prefix siftsmall/siftsmall_R32_L50_filtered_index -R 32 --FilteredLbuild 50 --alpha 1.2 --label_file ./rand_labels_50_10K.txt --universal_label 0 -build/apps/build_stitched_index --data_type float --data_path siftsmall/siftsmall_base.bin --index_path_prefix siftsmall/siftsmall_R20_L40_SR32_stitched_index -R 20 -L 40 --stitched_R 32 --alpha 1.2 --label_file ./rand_labels_50_10K.txt --universal_label 0 -build/apps/search_memory_index --data_type float --dist_fn l2 --index_path_prefix data/sift/siftsmall_R20_L40_SR32_filtered_index --query_file siftsmall/siftsmall_query.bin --gt_file siftsmall/siftsmall_gt_35.bin --filter_label 35 -K 10 -L 10 20 30 40 50 100 --result_path siftsmall/filtered_search_results -build/apps/search_memory_index --data_type float --dist_fn l2 --index_path_prefix data/sift/siftsmall_R20_L40_SR32_stitched_index --query_file siftsmall/siftsmall_query.bin --gt_file siftsmall/siftsmall_gt_35.bin --filter_label 35 -K 10 -L 10 20 30 40 50 100 --result_path siftsmall/stitched_search_results -``` - - The output of both searches is listed below. The throughput (Queries/sec) as well as mean and 99.9 latency in microseconds for each `L` parameter provided. (Measured on a physical machine with a Intel(R) Xeon(R) W-2145 CPU and 64 GB RAM) - ``` - Stitched Index - Ls QPS Avg dist cmps Mean Latency (mus) 99.9 Latency Recall@10 -================================================================================= - 10 31324.39 37.33 116.79 311.90 17.80 - 20 91357.57 44.36 193.06 1042.30 17.90 - 30 69314.48 49.89 258.09 1398.00 18.20 - 40 61421.29 60.52 289.08 1515.00 18.60 - 50 54203.48 70.27 294.26 685.10 19.40 - 100 52904.45 79.00 336.26 1018.80 19.50 - -Filtered Index - Ls QPS Avg dist cmps Mean Latency (mus) 99.9 Latency Recall@10 -================================================================================= - 10 69671.84 21.48 45.25 146.20 11.60 - 20 168577.20 38.94 100.54 547.90 18.20 - 30 127129.41 52.95 126.83 768.40 19.70 - 40 106349.04 62.38 167.23 899.10 20.90 - 50 89952.33 70.95 189.12 1070.80 22.10 - 100 56899.00 112.26 304.67 636.60 23.80 - ``` diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_ssd_index.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_ssd_index.md deleted file mode 100644 index 7457d8c..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/filtered_ssd_index.md +++ /dev/null @@ -1,103 +0,0 @@ -**Usage for filtered indices** -================================ - -To generate an SSD-friendly index, use the `apps/build_disk_index` program. ----------------------------------------------------------------------------- - -## Building a SSD based filtered Index - -### filtered-vamana SSD Index - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). -3. **--data_file**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as an integer. The next 4 bytes represent the dimension of data as an integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. `sizeof(T)` is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **--index_path_prefix**: the index will span a few files, all beginning with the specified prefix path. For example, if you provide `~/index_test` as the prefix path, build generates files such as `~/index_test_pq_pivots.bin, ~/index_test_pq_compressed.bin, ~/index_test_disk.index, ...`. There may be between 8 and 10 files generated with this prefix depending on how the index is constructed. -5. **-R (--max_degree)** (default is 64): the degree of the graph index, typically between 60 and 150. Larger R will result in larger indices and longer indexing times, but better search quality. -6. **-L (--Lbuild)** (default is 100): the size of search listduring index build. Typical values are between 75 to 200. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Use a value for L value that is at least the value of R unless you need to build indices really quickly and can somewhat compromise on quality. Note that this is to be used only for building an unfiltered index. The corresponding search list parameter for a filtered index is managed by `--FilteredLbuild`. -7. **-B (--search_DRAM_budget)**: bound on the memory footprint of the index at search time in GB. Once built, the index will use up only the specified RAM limit, the rest will reside on disk. This will dictate how aggressively we compress the data vectors to store in memory. Larger will yield better performance at search time. For an n point index, to use b byte PQ compressed representation in memory, use `B = ((n * b) / 2^30 + (250000*(4*R + sizeof(T)*ndim)) / 2^30)`. The second term in the summation is to allow some buffer for caching about 250,000 nodes from the graph in memory while serving. If you are not sure about this term, add 0.25GB to the first term. -8. **-M (--build_DRAM_budget)**: Limit on the memory allowed for building the index in GB. If you specify a value less than what is required to build the index in one pass, the index is built using a divide and conquer approach so that sub-graphs will fit in the RAM budget. The sub-graphs are overlayed to build the overall index. This approach can be upto 1.5 times slower than building the index in one shot. Allocate as much memory as your RAM allows. -9. **-T (--num_threads)** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -10. **--PQ_disk_bytes** (default is 0): Use 0 to store uncompressed data on SSD. This allows the index to asymptote to 100% recall. If your vectors are too large to store in SSD, this parameter provides the option to compress the vectors using PQ for storing on SSD. This will trade off recall. You would also want this to be greater than the number of bytes used for the PQ compressed data stored in-memory -11. **--build_PQ_bytes** (default is 0): Set to a positive value less than the dimensionality of the data to enable faster index build with PQ based distance comparisons. -12. **--use_opq**: use the flag to use OPQ rather than PQ compression. OPQ is more space efficient for some high dimensional datasets, but also needs a bit more build time. -13. **--label_file**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of filters corresponding to point `i` in the file passed via `--data_file`. -14. **--universal_label**: Optionally, the label data may contain a special "universal" label. A point with the universal label can be matched against a query with any label. Note that if a point has the universal label, then the filter data must only have the universal label on the line corresponding. -15. **--FilteredLbuild**: If building a filtered index, we maintain a separate search list from the one provided by `--Lbuild`. -16. **--filter_threshold**: Threshold to break up the existing nodes to generate new graph internally by breaking dense points where each node will have a maximum F labels. Default value is zero where no break up happens for the dense points. - - -## Computing a groundtruth file for a filtered index -In order to evaluate the performance of our algorithms, we can compare its results (i.e. the top `k` neighbors found for each query) against the results found by an exact nearest neighbor search. We provide the program `apps/utils/compute_groundtruth.cpp` to provide the results for the latter: - -1. **`--data_type`** The type of dataset you built an index with. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **`--dist_fn`**: There are two distance functions supported: l2 and mips. -3. **`--base_file`**: The input data over which to build an index, in .bin format. Corresponds to the `--data_path` argument from above. -4. **`--query_file`**: The queries to be searched on, which are stored in the same .bin format. -5. **`--label_file`**: Filter data for each point, in `.txt` format. Line `i` of the file consists of a comma-separated list of filters corresponding to point `i` in the file passed via `--data_file`. -6. **`--filter_label`**: Filter for each query. For each query, a search is performed with this filter. -7. **`--universal_label`**: Corresponds to the universal label passed when building an index with filter support. -8. **`--gt_file`**: File to output results to. The binary file starts with `n`, the number of queries (4 bytes), followed by `d`, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the `d` closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. -9. **`-K`**: The number of nearest neighbors to compute for each query. - -## Searching a Filtered Index - -Searching a filtered index uses the `apps/search_disk_index.cpp`: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. Use the same data type as in arg (1) above used in building the index. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). Use the same distance as in arg (2) above used in building the index. -3. **--index_path_prefix**: same as the prefix used in building the index (see arg 4 above). -4. **--num_nodes_to_cache** (default is 0): While serving the index, the entire graph is stored on SSD. For faster search performance, you can cache a few frequently accessed nodes in memory. -5. **-T (--num_threads)** (default is to get_omp_num_procs()): The number of threads used for searching. Threads run in parallel and one thread handles one query at a time. More threads will result in higher aggregate query throughput, but will also use more IOs/second across the system, which may lead to higher per-query latency. So find the balance depending on the maximum number of IOPs supported by the SSD. -6. **-W (--beamwidth)** (default is 2): The beamwidth to be used for search. This is the maximum number of IO requests each query will issue per iteration of search code. Larger beamwidth will result in fewer IO round-trips per query, but might result in slightly higher total number of IO requests to SSD per query. For the highest query throughput with a fixed SSD IOps rating, use `W=1`. For best latency, use `W=4,8` or higher complexity search. Specifying 0 will optimize the beamwidth depending on the number of threads performing search, but will involve some tuning overhead. -7. **--query_file**: The queries to be searched on in same binary file format as the data file in arg (2) above. The query file must be the same type as argument (1). -8. **--gt_file**: The ground truth file for the queries in arg (7) and data file used in index construction. The binary file must start with *n*, the number of queries (4 bytes), followed by *d*, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the d closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. The groundtruth file, if not available, can be calculated using the program `apps/utils/compute_groundtruth`. Use "null" if you do not have this file and if you do not want to compute recall. -9. **-K**: search for *K* neighbors and measure *K*-recall@*K*, meaning the intersection between the retrieved top-*K* nearest neighbors and ground truth *K* nearest neighbors. -10. **--result_path**: Search results will be stored in files with specified prefix, in bin format. -11. **-L (--search_list)**: A list of search_list sizes to perform search with. Larger parameters will result in slower latencies, but higher accuracies. Must be atleast the value of *K* in arg (9). -12. **--filter_label**: The filter to be used when searching an index with filters. For each query, a search is performed with this filter. - - -Example with SIFT10K: --------------------- -We demonstrate how to work through this pipeline using the SIFT10K dataset (http://corpus-texmex.irisa.fr/). Before starting, make sure you have compiled diskANN according to the instructions in the README and can see the following binaries (paths with respect to repository root): -- `build/apps/utils/compute_groundtruth` -- `build/apps/utils/fvecs_to_bin` -- `build/apps/build_disk_index` -- `build/apps/search_disk_index` - -Now, download the base and query set and convert the data to binary format: -```bash -wget ftp://ftp.irisa.fr/local/texmex/corpus/siftsmall.tar.gz -tar -zxvf siftsmall.tar.gz -build/apps/utils/fvecs_to_bin float siftsmall/siftsmall_base.fvecs siftsmall/siftsmall_base.bin -build/apps/utils/fvecs_to_bin float siftsmall/siftsmall_query.fvecs siftsmall/siftsmall_query.bin -``` - -We now need to make label file for our vectors. For convenience, we've included a synthetic label generator through which we can generate label file as follow -```bash - build/apps/utils/generate_synthetic_labels --num_labels 50 --num_points 10000 --output_file ./rand_labels_50_10K.txt --distribution_type zipf -``` -Note : `distribution_type` can be `rand` or `zipf` - -This will genearate label file with 10000 data points with 50 distinct labels, ranging from 1 to 50 assigned using zipf distribution (0 is the universal label). - -Now build and search the index and measure the recall using ground truth computed using bruteforce. We search for results with the filter 35. -```bash -build/apps/utils/compute_groundtruth --data_type float --dist_fn l2 --base_file siftsmall/siftsmall_base.bin --query_file siftsmall/siftsmall_query.bin --gt_file siftsmall_gt_35.bin --K 100 --label_file rand_labels_50_10K.txt --filter_label 35 --universal_label 0 -build/apps/build_disk_index --data_type float --dist_fn l2 --data_path siftsmall/siftsmall_base.bin --index_path_prefix data/sift/siftsmall_R32_L50_filtered -R 32 --FilteredLbuild 50 -B 1 -M 1 --label_file rand_labels_50_10K.txt --universal_label 0 -F 0 -build/apps/search_disk_index --data_type float --dist_fn l2 --index_path_prefix data/sift/siftsmall_R32_L50_filtered --result_path siftsmall/search_35 --query_file siftsmall/siftsmall_query.bin --gt_file siftsmall_gt_35.bin -K 10 -L 10 20 30 40 50 100 --filter_label 35 -W 4 -T 8 -``` - - The output of both searches is listed below. The throughput (Queries/sec) as well as mean and 99.9 latency in microseconds for each `L` parameter provided. (Measured on a physical machine with a 11th Gen Intel(R) Core(TM) i7-1185G7 CPU and 32 GB RAM) - - ``` -Filtered Disk Index - L Beamwidth QPS Mean Latency 99.9 Latency Mean IOs CPU (s) Recall@10 -================================================================================================================== - 10 4 1922.02 4062.19 12849.00 15.49 66.19 11.80 - 20 4 4609.91 1618.68 3438.00 30.66 140.48 17.20 - 30 4 3377.83 2250.22 4631.00 42.70 202.39 20.70 - 40 4 2707.77 2817.21 4889.00 51.46 267.03 22.00 - 50 4 2191.56 3509.43 5943.00 60.80 349.10 23.50 -100 4 1257.92 6113.45 7321.00 109.08 609.42 23.90 -``` \ No newline at end of file diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/in_memory_index.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/in_memory_index.md deleted file mode 100644 index 6d78320..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/in_memory_index.md +++ /dev/null @@ -1,73 +0,0 @@ -**Usage for in-memory indices** -================================ - -To generate index, use the `apps/build_memory_index` program. --------------------------------------------------------------- - -The arguments are as follows: - -1. **--data_type**: The type of dataset you wish to build an index on. float(32 bit), signed int8 and unsigned uint8 are supported. -2. **--dist_fn**: There are two distance functions supported: minimum Euclidean distance (l2) and maximum inner product (mips). -3. **--data_file**: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following `n*d*sizeof(T)` bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. -4. **--index_path_prefix**: The constructed index components will be saved to this path prefix. -5. **-R (--max_degree)** (default is 64): the degree of the graph index, typically between 32 and 150. Larger R will result in larger indices and longer indexing times, but might yield better search quality. -6. **-L (--Lbuild)** (default is 100): the size of search list we maintain during index building. Typical values are between 75 to 400. Larger values will take more time to build but result in indices that provide higher recall for the same search complexity. Ensure that value of L is at least that of R value unless you need to build indices really quickly and can somewhat compromise on quality. -7. **--alpha** (default is 1.2): A float value between 1.0 and 1.5 which determines the diameter of the graph, which will be approximately *log n* to the base alpha. Typical values are between 1 to 1.5. 1 will yield the sparsest graph, 1.5 will yield denser graphs. -8. **T (--num_threads)** (default is to get_omp_num_procs()): number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine and DRAM bandwidth). -9. **--build_PQ_bytes** (default is 0): Set to a positive value less than the dimensionality of the data to enable faster index build with PQ based distance comparisons. Defaults to using full precision vectors for distance comparisons. -10.**--use_opq**: use the flag to use OPQ rather than PQ compression. OPQ is more space efficient for some high dimensional datasets, but also needs a bit more build time. - - -To search the generated index, use the `apps/search_memory_index` program: ---------------------------------------------------------------------------- - - -The arguments are as follows: - -1. **data_type**: The type of dataset you built the index on. float(32 bit), signed int8 and unsigned uint8 are supported. Use the same data type as in arg (1) above used in building the index. -2. **dist_fn**: There are two distance functions supported: l2 and mips. There is an additional *fast_l2* implementation that could provide faster results for small (about a million-sized) indices. Use the same distance as in arg (2) above used in building the index. -3. **memory_index_path**: index built above in argument (4). -4. **T**: The number of threads used for searching. Threads run in parallel and one thread handles one query at a time. More threads will result in higher aggregate query throughput, but may lead to higher per-query latency, especially if the DRAM bandwidth is a bottleneck. So find the balance depending on throughput and latency required for your application. -5. **query_bin**: The queries to be searched on in same binary file format as the data file (ii) above. The query file must be the same type as in argument (1). -6. **truthset.bin**: The ground truth file for the queries in arg (7) and data file used in index construction. The binary file must start with *n*, the number of queries (4 bytes), followed by *d*, the number of ground truth elements per query (4 bytes), followed by `n*d` entries per query representing the d closest IDs per query in integer format, followed by `n*d` entries representing the corresponding distances (float). Total file size is `8 + 4*n*d + 4*n*d` bytes. The groundtruth file, if not available, can be calculated using the program `apps/utils/compute_groundtruth`. Use "null" if you do not have this file and if you do not want to compute recall. -7. **K**: search for *K* neighbors and measure *K*-recall@*K*, meaning the intersection between the retrieved top-*K* nearest neighbors and ground truth *K* nearest neighbors. -8. **result_output_prefix**: search results will be stored in files, one per L value (see next arg), with specified prefix, in binary format. -9. **-L (--search_list)**: A list of search_list sizes to perform search with. Larger parameters will result in slower latencies, but higher accuracies. Must be atleast the value of *K* in (7). - - -Example with BIGANN: --------------------- - -This example demonstrates the use of the commands above on a 100K slice of the [BIGANN dataset](http://corpus-texmex.irisa.fr/) with 128 dimensional SIFT descriptors applied to images. - -Download the base and query set and convert the data to binary format -```bash -mkdir -p DiskANN/build/data && cd DiskANN/build/data -wget ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz -tar -xf sift.tar.gz -cd .. -./apps/utils/fvecs_to_bin float data/sift/sift_learn.fvecs data/sift/sift_learn.fbin -./apps/utils/fvecs_to_bin float data/sift/sift_query.fvecs data/sift/sift_query.fbin -``` - -Now build and search the index and measure the recall using ground truth computed using brutefoce. -```bash -./apps/utils/compute_groundtruth --data_type float --dist_fn l2 --base_file data/sift/sift_learn.fbin --query_file data/sift/sift_query.fbin --gt_file data/sift/sift_query_learn_gt100 --K 100 -./apps/build_memory_index --data_type float --dist_fn l2 --data_path data/sift/sift_learn.fbin --index_path_prefix data/sift/index_sift_learn_R32_L50_A1.2 -R 32 -L 50 --alpha 1.2 - ./apps/search_memory_index --data_type float --dist_fn l2 --index_path_prefix data/sift/index_sift_learn_R32_L50_A1.2 --query_file data/sift/sift_query.fbin --gt_file data/sift/sift_query_learn_gt100 -K 10 -L 10 20 30 40 50 100 --result_path data/sift/res - ``` - - - The output of search lists the throughput (Queries/sec) as well as mean and 99.9 latency in microseconds for each `L` parameter provided. (We measured on a 32-core 64-vCPU D-series Azure VM) - ``` - Ls QPS Avg dist cmps Mean Latency (mus) 99.9 Latency Recall@10 -================================================================================= - 10 319901.78 348.93 174.51 4943.35 97.80 - 20 346572.72 525.85 183.36 376.60 98.93 - 30 292060.12 688.86 217.73 421.60 99.30 - 40 248945.22 841.74 255.41 476.80 99.45 - 50 215888.81 986.67 294.62 542.21 99.56 - 100 129711.39 1631.94 490.58 848.61 99.88 - ``` - - diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/python.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/python.md deleted file mode 100644 index d009cd7..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/python.md +++ /dev/null @@ -1,133 +0,0 @@ -# `diskannpy` - -We publish (sporadic) builds of `diskann` with python bindings to `pypi.org`, which you can install via `pip install diskannpy`. - -#### Caveats -Native python modules with cffi need to be built for *every* version of Python and *every* OS and *every* native-integration-library. - -This makes for a complicated build matrix that only `(ana)conda` is properly fit to solve. However, we do build wheels -for python 3.9-3.11, across linux, Windows, and macOS (x86_64). These versions are also built against `numpy` 1.25 - -which makes for a hard runtime requirement that can be challenging to use if you are using older or newer versions of numpy. - -There *are* instructions for building against other versions of numpy -[documented in this issue response](https://github.com/microsoft/DiskANN/issues/544#issuecomment-2103437976) if you require a different build. - -# Basic Usage - -`diskannpy` provides access to both building and reading `DiskANN` indices. In all cases, the _lingua franca_ is numpy -ndarrays. Currently, the only supported dtypes are `np.float32`, `np.int8`, and `np.uint8`. - -`diskannpy` provides a number of helpful functions, like reading or writing `diskann` style vector binary files via the -`vectors_to_file` and `vectors_from_file` functions. For a full suite of python functions and their documentation, -please be sure to read the latest documentation @ [https://microsoft.github.io/](https://microsoft.github.io/DiskANN/docs/python/latest/diskannpy.html). - - -## Scenarios -The following scenarios are supported via the `diskannpy` api. - - -### Commonalities -```python -my_dtype = np.float32 # or np.uint8 or np.int8 ONLY -my_set_of_vectors: np.typing.NDArray[my_dtype] = ... # your vectors come from somewhere - you need to bring these! -index_to_identifiers_map: np.typing.NDArray[str] = ... # your vectors likely have some kind of external identifier - -# you need to keep track of the external identifier -> index relationship somehow -identifiers_to_index_map: dict[str, np.uint32|np.uint.64] = ... # your map of your external id to the `diskannpy` internal id -# diskannpy `query` responses will contain the _internal id only_, and if you don't have these maps you won't be able to -# know what this relates to -``` - -### Build Disk Index -A disk index is a memory mapped, [vamana](https://proceedings.neurips.cc/paper_files/paper/2019/file/09853c7fb1d3f8ee67a61b6bf4a7f8e6-Paper.pdf) -index that heavily leans into the hardware speeds of modern NVMe based solid state storage. - -This means you can build performant ANN indices that overflow plausibly available system memory! - -```python -import numpy as np -import diskannpy as dap - -vecs = my_set_of_vectors / np.linalg.norm(my_set_of_vectors, axis=1) # useful if your intention is to rank by a directionless -# cosine angle distance - -dap.build_disk_index( - data=vecs, - distance_metric="l2", # can also be cosine, especially if you don't normalize your vectors like above - index_directory="/tmp/my_index", - complexity=128, # the larger this is, the more candidate points we consider when ranking - graph_degree=64, # the beauty of a vamana index is it's ability to shard and be able to transfer long distances across the grpah without navigating the whole thing. the larger this value is, the higher quality your results, but the longer it will take to build - search_memory_maximum=16.0, # a floating point number to represent how much memory in GB we want to optimize for @ query time - build_memory_maximum=100.0, # a floating point number to represent how much memory in GB we are allocating for the index building process - num_threads=0, # 0 means use all available threads - but if you are in a shared environment you may need to restrict how greedy you are - vector_dtype=my_dtype, # we specified this in the Commonalities section above - index_prefix="ann", # ann is the default anyway. all files generated will have the prefix `ann_`, in the form of `f"{index_prefix}_"` - pq_disk_bytes=0 # using product quantization of your vectors can still achieve excellent recall characteristics at a fraction of the latency, but we'll do it without PQ for now -) -``` - -### Search Disk Index - -Now we want to search our disk index - using a completely different set of vectors that aren't necessarily guaranteed to -be in our index. We will call this set of vectors `q`, and it is *critical* that they are the same dtype and -dimensionality as the disk index we have just built. - -**Note**: If you manually normalized your indexed vectors prior to building the index, you will *also* need to normalize -them prior to query! - -#### Common index query setup - -```python -index = dap.StaticDiskIndex( - index_directory="/tmp/my_index", - num_threads=0, - num_nodes_to_cache=1_000_000, - index_prefix="ann" -) -``` - -#### Individual Vectors -```python -some_index: np.uint32 = ... # the index in our `q` array of points that we will be using to query on an individual basis -my_query_vector: np.typing.NDArray[my_dtype] = q[some_index] # make sure this is a 1-d array of the same dimensionality as your index! -# normalize if required by my_query_vector /= np.linalg.norm(my_query_vector) -internal_indices, distances = index.search( - query=my_query_vector, - k_neighbors=25, - complexity=50, # must be as big or bigger than `k_neighbors` -) -``` - -#### Mapping to our External Ids -The internal IDs that diskann returns via query aren't necessarily directly useful to you, and the onus is on you -to figure out what they actually link to via your `index_to_identifiers_map` map. -```python -actual_identifiers = index_to_identifiers_map[internal_indices] # using np fancy indexing (advanced indexing?) to map them all to ids you actually understand -``` - -#### Batch Vectors -```python -import multiprocessing - -internal_indices, distances = index.batch_search( - queries=q, - k_neighbors=25, - complexity=50, - num_threads=multiprocessing.cpu_count(), # there's a current bug where this is not handling the value 0 properly - beam_width=8 # beamwidth is the parameter that indicates our parallelism of individual searches, whereas num_threads - # indicates the number of threads *per* query item in the batch -) -# note that in batch_query form, our internal_indices and distances are 2d arrays -``` - -#### Mapping to our External Ids -Unlike the previous entry, I have yet to get the fancy awesome advanced indexing to work in one shot, we will have -to do this the not-numpy-paragon way. - -```python -actual_neighbors = np.full(shape=internal_indices.shape, dtype=str, fill_value="") -for row in range(internal_indices.shape[0]): - actual_neighbors[row] = index_to_identifiers_map[internal_indices[row]] -``` - -This is only scratching the surface of what `diskannpy` can offer. Please read the API documentation @ [https://microsoft.github.io/](https://microsoft.github.io/DiskANN/docs/python/latest/diskannpy.html) -for more details. diff --git a/packages/leann-backend-diskann/third_party/DiskANN/workflows/rest_api.md b/packages/leann-backend-diskann/third_party/DiskANN/workflows/rest_api.md deleted file mode 100644 index b735fbe..0000000 --- a/packages/leann-backend-diskann/third_party/DiskANN/workflows/rest_api.md +++ /dev/null @@ -1,72 +0,0 @@ - -**REST service set up for serving DiskANN indices and query interface** -======================================================================= - -Install dependencies on Ubuntu and compile ------------------------------------------- -In addition to the common dependencies in the [README](/README.md), install [Microsoft C++ REST SDK](https://github.com/Microsoft/cpprestsdk). - -```bash -sudo apt install libcpprest-dev -mkdir -p build && cd build -cmake -DRESTAPI=True -DCMAKE_BUILD_TYPE=Release .. -make -j -``` - -Starting an index hosting service ---------------------------------- -Follow the instructions for [building an in-memory DiskANN index](/workflows/in_memory_index.md) or [building an SSD DiskANN index](/workflows/SSD_index.md). Then start a service bound at the appropriate IP:port. For querying from the local machine, you may want to use `http://127.0.0.1:port`. For serving queries originating from remote machines, you may want to use `http://0.0.0.0:port`. - -```bash -# To start serving an in-memory index -./apps/restapi/inmem_server --address --data_type --data_file --index_path_prefix --num_threads --l_search --tags_file [tags_file] - -# To start serving an SSD-based index. -./apps/restapi/ssd_server --address --data_type --index_path_prefix --num_nodes_to_cache --num_threads --tags_file [tags_file] -``` -The `data_type` and the `data_file` should be the same as those used in the construction of the index. The server returns the ids and distances of the closests vector in the index to the query. The ids are implicitly defined by the order of the vector in the data file. If you wish to assign a different numbering or GUID or URL to the vectors in the index, use the optional `tags_file`. This should be a file which lists a "tag" string for each vector in the index. The file should contain one string per line. The string on the line `n` is considered the tag corresponding to the vector `n` in the index (in the implicit order defined in the `data_file`). - -For an SSD-based index, specify the number of nodes to cache in-memory to make queries faster. For large indices with over 100 million vectors, a typical value for `num_nodes_to_cache` could be 500000. Increase or decrease based on DRAM footprint desired. - -For an SSD-based index, also specify the number of threads used for search by setting the `num_threads` parameter. - -You can also query multiple SSD based indices using the following command by listing the prefix of each index in a file (one prefix per line) and passing it through the `index_prefix_paths` parameter to the following command. -```bash -multiple_ssdserver --address --data_type --index_prefix_paths --num_nodes_to_cache --num_threads --tags_file [tags_file] -``` -The service searches each of the indices and aggregate the results based on distances to find the closest neighbors across all indices. - -Querying the service --------------------- -Issue a json query with the following fields -- "k" : The number of nearest neighbors needed -- "query" : The query vector with a listing of co-ordinates. -- "query_id" : An id to track the query. Use a unique number to keep track of queries, or "0" if you do not want to keep track. -- "Ls" : query complexity. Higher Ls takes more milliseconds to process but offers higher recall. Default to 256 if you don't want to tune this. - -**Post a json query using python** - -```python -import requests -jsonquery = {"Ls": 256, - "query_id": 1234, - "query": [0.00407, 0.01534, 0.02498, ...], - "k": 10} - -response = requests.post('http://ip_addr:port', json=jsonquery) -print(response.text) -``` - -The response might look like the following. The partition array indicates the ID of index from which the result was found in the case of a multi-index set up. For a single index set up, the response would not contain the information on partitions. The response may or may not contain `tags` based on whether the server was started with a `tags_file`. -```json -{"distances":[1.6947,1.6954,1.6972,1.6985,1.6991,1.7003,1.7008,1.7014,1.7021,1.7039],"indices":[8976853,8221762,30909336,13100282,30514543,11537860,7133262,34074869,50512601,17983301],"k":10,"partition":[20,7,20,20,6,6,11,6,6,20],"query_id":1234,"tags":["https://xyz1", "https://xyz2", "https://xyz3", "https://xyz4", "https://xyz5", "https://xyz6", "https://xyz7", "https://xyz8", "https://xyz9", "https://xyz10"],"time_taken_in_us":3245} -``` - -**Command line interface to issue multiple queries from a file** - -To issue `num_queries` queries from `query_file`, run the following command -```bash -client ip_addr:port data_type query_file num_queries Ls" -``` - diff --git a/packages/leann-backend-hnsw/third_party/faiss b/packages/leann-backend-hnsw/third_party/faiss new file mode 160000 index 0000000..b906cee --- /dev/null +++ b/packages/leann-backend-hnsw/third_party/faiss @@ -0,0 +1 @@ +Subproject commit b906ceeb8f93c589545b47bef697a993ca9ef9a0 diff --git a/packages/leann-backend-hnsw/third_party/faiss/.clang-format b/packages/leann-backend-hnsw/third_party/faiss/.clang-format deleted file mode 100644 index 1fe6508..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.clang-format +++ /dev/null @@ -1,88 +0,0 @@ ---- -AccessModifierOffset: -1 -AlignAfterOpenBracket: AlwaysBreak -AlignConsecutiveAssignments: false -AlignConsecutiveDeclarations: false -AlignEscapedNewlinesLeft: true -AlignOperands: false -AlignTrailingComments: true -AllowAllParametersOfDeclarationOnNextLine: false -AllowShortBlocksOnASingleLine: false -AllowShortCaseLabelsOnASingleLine: false -AllowShortFunctionsOnASingleLine: Empty -AllowShortIfStatementsOnASingleLine: false -AllowShortLoopsOnASingleLine: false -AlwaysBreakAfterReturnType: None -AlwaysBreakBeforeMultilineStrings: true -AlwaysBreakTemplateDeclarations: true -BinPackArguments: false # at some point, set this to true -BinPackParameters: false # at some point, set this to true -BraceWrapping: - AfterClass: false - AfterControlStatement: false - AfterEnum: false - AfterFunction: false - AfterNamespace: false - AfterObjCDeclaration: false - AfterStruct: false - AfterUnion: false - BeforeCatch: false - BeforeElse: false - IndentBraces: false -BreakBeforeBinaryOperators: None -BreakBeforeBraces: Attach -BreakBeforeTernaryOperators: true -BreakConstructorInitializersBeforeComma: false -BreakAfterJavaFieldAnnotations: false -BreakStringLiterals: false -ColumnLimit: 80 -CommentPragmas: '^ IWYU pragma:' -CompactNamespaces: false -ConstructorInitializerAllOnOneLineOrOnePerLine: true -ConstructorInitializerIndentWidth: 8 -ContinuationIndentWidth: 8 -Cpp11BracedListStyle: true -DerivePointerAlignment: false -DisableFormat: false -ForEachMacros: [ FOR_EACH_RANGE, FOR_EACH, ] -IncludeCategories: - - Regex: '^<.*\.h(pp)?>' - Priority: 1 - - Regex: '^<.*' - Priority: 2 - - Regex: '.*' - Priority: 3 -IndentCaseLabels: true -IndentWidth: 4 -IndentWrappedFunctionNames: false -KeepEmptyLinesAtTheStartOfBlocks: false -MacroBlockBegin: '' -MacroBlockEnd: '' -MaxEmptyLinesToKeep: 1 -NamespaceIndentation: None -ObjCBlockIndentWidth: 4 -ObjCSpaceAfterProperty: false -ObjCSpaceBeforeProtocolList: false -PenaltyBreakBeforeFirstCallParameter: 1 -PenaltyBreakComment: 300 -PenaltyBreakFirstLessLess: 120 -PenaltyBreakString: 1000 -PenaltyExcessCharacter: 1000000 -PenaltyReturnTypeOnItsOwnLine: 2000000 -PointerAlignment: Left -ReflowComments: true -SortIncludes: true -SpaceAfterCStyleCast: false -SpaceBeforeAssignmentOperators: true -SpaceBeforeParens: ControlStatements -SpaceInEmptyParentheses: false -SpacesBeforeTrailingComments: 1 -SpacesInAngles: false -SpacesInContainerLiterals: true -SpacesInCStyleCastParentheses: false -SpacesInParentheses: false -SpacesInSquareBrackets: false -Standard: Cpp11 -TabWidth: 8 -UseTab: Never -... diff --git a/packages/leann-backend-hnsw/third_party/faiss/.dockerignore b/packages/leann-backend-hnsw/third_party/faiss/.dockerignore deleted file mode 100644 index 7763a51..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.dockerignore +++ /dev/null @@ -1 +0,0 @@ -sift1M \ No newline at end of file diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/ISSUE_TEMPLATE.md b/packages/leann-backend-hnsw/third_party/faiss/.github/ISSUE_TEMPLATE.md deleted file mode 100644 index 132be64..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/ISSUE_TEMPLATE.md +++ /dev/null @@ -1,33 +0,0 @@ -# Summary - - - -# Platform - - - -OS: - -Faiss version: - -Installed from: - -Faiss compilation options: - -Running on: -- [ ] CPU -- [ ] GPU - -Interface: -- [ ] C++ -- [ ] Python - -# Reproduction instructions - - - - diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_cmake/action.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_cmake/action.yml deleted file mode 100644 index 6251519..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_cmake/action.yml +++ /dev/null @@ -1,189 +0,0 @@ -name: Build cmake -inputs: - opt_level: - description: 'Compile options / optimization level.' - required: false - default: generic - gpu: - description: 'Enable GPU support.' - required: false - default: OFF - cuvs: - description: 'Enable cuVS support.' - required: false - default: OFF - rocm: - description: 'Enable ROCm support.' - required: false - default: OFF -runs: - using: composite - steps: - - name: Setup miniconda - uses: conda-incubator/setup-miniconda@v3 - with: - python-version: '3.11' - miniforge-version: latest # ensures conda-forge channel is used. - channels: conda-forge - conda-remove-defaults: 'true' - # Set to aarch64 if we're on arm64 because there's no miniforge ARM64 package, just aarch64. - # They are the same thing, just named differently. - architecture: ${{ runner.arch == 'ARM64' && 'aarch64' || runner.arch }} - - name: Configure build environment - shell: bash - run: | - # initialize Conda - conda config --set solver libmamba - # Ensure starting packages are from conda-forge. - conda list --show-channel-urls - conda update -y -q conda - echo "$CONDA/bin" >> $GITHUB_PATH - - conda install -y -q python=3.11 cmake=3.26 make=4.2 swig=4.0 "numpy<2" scipy=1.14 pytest=7.4 gflags=2.2 - - # install base packages for ARM64 - if [ "${{ runner.arch }}" = "ARM64" ]; then - conda install -y -q -c conda-forge openblas=0.3.29 gxx_linux-aarch64=14.2 sysroot_linux-aarch64=2.17 - fi - - # install base packages for X86_64 - if [ "${{ runner.arch }}" = "X64" ]; then - # TODO: merge this with ARM64 - conda install -y -q -c conda-forge gxx_linux-64=14.2 sysroot_linux-64=2.17 - conda install -y -q mkl=2022.2.1 mkl-devel=2022.2.1 - fi - - # no CUDA needed for ROCm so skip this - if [ "${{ inputs.rocm }}" = "ON" ]; then - : - # regular CUDA for GPU builds - elif [ "${{ inputs.gpu }}" = "ON" ] && [ "${{ inputs.cuvs }}" = "OFF" ]; then - conda install -y -q cuda-toolkit=12.4 -c "nvidia/label/cuda-12.4.0" - # and CUDA from cuVS channel for cuVS builds - elif [ "${{ inputs.cuvs }}" = "ON" ]; then - conda install -y -q libcuvs=24.12 'cuda-version>=12.0,<=12.5' cuda-toolkit=12.4.1 gxx_linux-64=12.4 -c rapidsai -c conda-forge - fi - - # install test packages - if [ "${{ inputs.rocm }}" = "ON" ]; then - : # skip torch install via conda, we need to install via pip to get - # ROCm-enabled version until it's supported in conda by PyTorch - elif [ "${{ inputs.gpu }}" = "ON" ]; then - conda install -y -q "pytorch<2.5" pytorch-cuda=12.4 -c pytorch -c "nvidia/label/cuda-12.4.0" - else - conda install -y -q "pytorch<2.5" -c pytorch - fi - - name: ROCm - Install dependencies - if: inputs.rocm == 'ON' - shell: bash - run: | - # Update repos and install kmod, wget, gpg - sudo apt-get -qq update >/dev/null - sudo apt-get -qq install -y kmod wget gpg >/dev/null - - # Get UBUNTU version name - UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'` - - # Set ROCm version - ROCM_VERSION="6.2" - - # Download, prepare, and install the package signing key - mkdir --parents --mode=0755 /etc/apt/keyrings - wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null - - # Add rocm repository - wget -qO - http://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add - - rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}" - echo "deb [arch=amd64] ${rocm_baseurl} ${UBUNTU_VERSION_NAME} main" | sudo tee /etc/apt/sources.list.d/rocm.list - sudo apt-get -qq update --allow-insecure-repositories >/dev/null - sudo apt-get -qq install -y --allow-unauthenticated \ - "rocm-dev${ROCM_VERSION}" "rocm-utils${ROCM_VERSION}" \ - "rocm-libs${ROCM_VERSION}" >/dev/null - - # Fake presence of MI200-class accelerators - echo "gfx90a" | sudo tee /opt/rocm/bin/target.lst - - # Cleanup - sudo apt-get -qq autoclean >/dev/null - sudo apt-get -qq clean >/dev/null - sudo rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* - - name: Symblink system dependencies - if: inputs.rocm == 'ON' - shell: bash - run: | - # symblink system libraries for HIP compiler - sudo ln -s /lib/x86_64-linux-gnu/libc.so.6 /lib64/libc.so.6 - sudo ln -s /lib/x86_64-linux-gnu/libc_nonshared.a /usr/lib64/libc_nonshared.a - sudo ln -s /usr/lib/x86_64-linux-gnu/libpthread.so.0 /lib64/libpthread.so.0 - sudo ln -s $HOME/miniconda3/x86_64-conda-linux-gnu/sysroot/usr/lib64/libpthread_nonshared.a /usr/lib64/libpthread_nonshared.a - - name: Build all targets - shell: bash - run: | - eval "$(conda shell.bash hook)" - conda activate - cmake -B build \ - -DBUILD_TESTING=ON \ - -DBUILD_SHARED_LIBS=ON \ - -DFAISS_ENABLE_GPU=${{ inputs.gpu }} \ - -DFAISS_ENABLE_CUVS=${{ inputs.cuvs }} \ - -DFAISS_ENABLE_ROCM=${{ inputs.rocm }} \ - -DFAISS_OPT_LEVEL=${{ inputs.opt_level }} \ - -DFAISS_ENABLE_C_API=ON \ - -DPYTHON_EXECUTABLE=$CONDA/bin/python \ - -DCMAKE_BUILD_TYPE=Release \ - -DBLA_VENDOR=${{ runner.arch == 'X64' && 'Intel10_64_dyn' || '' }} \ - -DCMAKE_CUDA_FLAGS=${{ runner.arch == 'X64' && '"-gencode arch=compute_75,code=sm_75"' || '' }} \ - . - make -k -C build -j$(nproc) - - name: C++ tests - shell: bash - run: | - export GTEST_OUTPUT="xml:$(realpath .)/test-results/googletest/" - make -C build test - - name: C++ perf benchmarks - shell: bash - if: inputs.rocm == 'OFF' - run: | - find ./build/perf_tests/ -executable -type f -name "bench*" -exec '{}' -v \; - - name: Install Python extension - shell: bash - working-directory: build/faiss/python - run: | - $CONDA/bin/python setup.py install - - name: ROCm - install ROCm-enabled torch via pip - if: inputs.rocm == 'ON' - shell: bash - run: | - pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1 - - name: Python tests (CPU only) - if: inputs.gpu == 'OFF' - shell: bash - run: | - pytest --junitxml=test-results/pytest/results.xml tests/test_*.py - pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py - - name: Python tests (CPU + GPU) - if: inputs.gpu == 'ON' - shell: bash - run: | - pytest --junitxml=test-results/pytest/results.xml tests/test_*.py - pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py - cp tests/common_faiss_tests.py faiss/gpu/test - pytest --junitxml=test-results/pytest/results-gpu.xml faiss/gpu/test/test_*.py - pytest --junitxml=test-results/pytest/results-gpu-torch.xml faiss/gpu/test/torch_*.py - - name: Test avx2 loading - if: inputs.opt_level == 'avx2' - shell: bash - run: | - FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss.so - LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss_avx2.so - - name: Upload test results - if: always() - uses: actions/upload-artifact@v4 - with: - name: test-results-arch=${{ runner.arch }}-opt=${{ inputs.opt_level }}-gpu=${{ inputs.gpu }}-cuvs=${{ inputs.cuvs }}-rocm=${{ inputs.rocm }} - path: test-results - - name: Check installed packages channel - shell: bash - run: | - # Shows that all installed packages are from conda-forge. - conda list --show-channel-urls diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_conda/action.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_conda/action.yml deleted file mode 100644 index 14c2270..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/actions/build_conda/action.yml +++ /dev/null @@ -1,107 +0,0 @@ -name: Conda build -description: Builds Faiss inside a Conda environment and uploads to repository when label is provided. -inputs: - label: - description: "The label to be used for uploads to Conda." - default: "" - required: false - cuda: - description: "CUDA toolkit version to use." - default: "" - required: false - cuvs: - description: "Enable cuVS support." - default: "" - required: false -runs: - using: composite - steps: - - name: Choose shell - shell: bash - id: choose_shell - run: | - # Use pwsh on Windows; bash everywhere else - if [ "${{ runner.os }}" != "Windows" ]; then - echo "shell=bash" >> "$GITHUB_OUTPUT" - else - echo "shell=pwsh" >> "$GITHUB_OUTPUT" - fi - - name: Setup miniconda - uses: conda-incubator/setup-miniconda@v3 - with: - python-version: '3.11' - miniforge-version: latest # ensures conda-forge channel is used. - channels: conda-forge - conda-remove-defaults: 'true' - # Set to runner.arch=aarch64 if we're on arm64 because - # there's no miniforge ARM64 package, just aarch64. - # They are the same thing, just named differently. - # However there is an ARM64 for macOS, so exclude that. - architecture: ${{ (runner.arch == 'ARM64' && runner.os != 'macOS') && 'aarch64' || runner.arch }} - - name: Install conda build tools - shell: ${{ steps.choose_shell.outputs.shell }} - run: | - # Ensure starting packages are from conda-forge. - conda list --show-channel-urls - conda install -y -q "conda!=24.11.0" - conda install -y -q "conda-build!=24.11.0" "liblief=0.14.1" - conda list --show-channel-urls - - name: Enable anaconda uploads - if: inputs.label != '' - shell: ${{ steps.choose_shell.outputs.shell }} - env: - PACKAGE_TYPE: ${{ inputs.label }} - run: | - conda install -y -q anaconda-client - conda config --set anaconda_upload yes - - name: Conda build (CPU) - if: inputs.label == '' && inputs.cuda == '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - run: | - conda build faiss --python 3.11 -c pytorch - - name: Conda build (CPU) w/ anaconda upload - if: inputs.label != '' && inputs.cuda == '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - env: - PACKAGE_TYPE: ${{ inputs.label }} - run: | - conda build faiss --user pytorch --label ${{ inputs.label }} -c pytorch - - name: Conda build (GPU) - if: inputs.label == '' && inputs.cuda != '' && inputs.cuvs == '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - run: | - conda build faiss-gpu --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ - -c pytorch -c nvidia/label/cuda-${{ inputs.cuda }} -c nvidia - - name: Conda build (GPU) w/ anaconda upload - if: inputs.label != '' && inputs.cuda != '' && inputs.cuvs == '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - env: - PACKAGE_TYPE: ${{ inputs.label }} - run: | - conda build faiss-gpu --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ - --user pytorch --label ${{ inputs.label }} -c pytorch -c nvidia/label/cuda-${{ inputs.cuda }} -c nvidia - - name: Conda build (GPU w/ cuVS) - if: inputs.label == '' && inputs.cuda != '' && inputs.cuvs != '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - run: | - conda build faiss-gpu-cuvs --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ - -c pytorch -c rapidsai -c rapidsai-nightly -c conda-forge -c nvidia - - name: Conda build (GPU w/ cuVS) w/ anaconda upload - if: inputs.label != '' && inputs.cuda != '' && inputs.cuvs != '' - shell: ${{ steps.choose_shell.outputs.shell }} - working-directory: conda - env: - PACKAGE_TYPE: ${{ inputs.label }} - run: | - conda build faiss-gpu-cuvs --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \ - --user pytorch --label ${{ inputs.label }} -c pytorch -c rapidsai -c rapidsai-nightly -c conda-forge -c nvidia - - name: Check installed packages channel - shell: ${{ steps.choose_shell.outputs.shell }} - run: | - # Shows that all installed packages are from conda-forge. - conda list --show-channel-urls diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/autoclose.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/autoclose.yml deleted file mode 100644 index 41a5827..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/autoclose.yml +++ /dev/null @@ -1,23 +0,0 @@ -name: Close Inactive Issues -on: - schedule: - - cron: "30 1 * * *" - -jobs: - close-issues: - runs-on: ubuntu-latest - permissions: - issues: write - pull-requests: write - steps: - - uses: actions/stale@v5 - with: - only-labels: autoclose - days-before-issue-stale: 7 - days-before-issue-close: 7 - stale-issue-label: "stale" - stale-issue-message: "This issue is stale because it has been open for 7 days with no activity." - close-issue-message: "This issue was closed because it has been inactive for 7 days since being marked as stale." - days-before-pr-stale: -1 - days-before-pr-close: -1 - repo-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-pull-request.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-pull-request.yml deleted file mode 100644 index bc0d2d6..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-pull-request.yml +++ /dev/null @@ -1,169 +0,0 @@ -on: - workflow_call: -env: - OMP_NUM_THREADS: '10' - MKL_THREADING_LAYER: GNU -jobs: - format: - name: Format - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Install clang-format - run: | - sudo apt-get update -y - sudo apt-get install -y wget - sudo apt install -y lsb-release wget software-properties-common gnupg - wget https://apt.llvm.org/llvm.sh - chmod u+x llvm.sh - sudo ./llvm.sh 18 - sudo apt-get install -y git-core clang-format-18 - - name: Verify clang-format - run: | - git ls-files | grep -E '\.(cpp|h|cu|cuh)$' | xargs clang-format-18 -i - if git diff --quiet; then - echo "Formatting OK!" - else - echo "Formatting not OK!" - echo "------------------" - git --no-pager diff --color - exit 1 - fi - linux-x86_64-cmake: - name: Linux x86_64 (cmake) - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - linux-x86_64-AVX2-cmake: - name: Linux x86_64 AVX2 (cmake) - needs: linux-x86_64-cmake - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - opt_level: avx2 - linux-x86_64-AVX512-cmake: - name: Linux x86_64 AVX512 (cmake) - needs: linux-x86_64-cmake - runs-on: faiss-aws-m7i.large - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - opt_level: avx512 - linux-x86_64-AVX512_SPR-cmake: - name: Linux x86_64 AVX512_SPR (cmake) - needs: linux-x86_64-cmake - runs-on: faiss-aws-m7i.large - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - opt_level: avx512_spr - linux-x86_64-GPU-cmake: - name: Linux x86_64 GPU (cmake) - needs: linux-x86_64-cmake - runs-on: 4-core-ubuntu-gpu-t4 - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - gpu: ON - linux-x86_64-GPU-w-CUVS-cmake: - name: Linux x86_64 GPU w/ cuVS (cmake) - needs: linux-x86_64-cmake - runs-on: 4-core-ubuntu-gpu-t4 - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - gpu: ON - cuvs: ON - linux-x86_64-GPU-w-ROCm-cmake: - name: Linux x86_64 GPU w/ ROCm (cmake) - needs: linux-x86_64-cmake - runs-on: faiss-amd-MI200 - container: - image: ubuntu:22.04 - options: --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE --cap-add=SYS_ADMIN - steps: - - name: Container setup - run: | - if [ -f /.dockerenv ]; then - apt-get update && apt-get install -y sudo && apt-get install -y git - git config --global --add safe.directory '*' - else - echo 'Skipping. Current job is not running inside a container.' - fi - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - gpu: ON - rocm: ON - linux-arm64-SVE-cmake: - name: Linux arm64 SVE (cmake) - needs: linux-x86_64-cmake - runs-on: faiss-aws-r8g.large - steps: - - name: Checkout - uses: actions/checkout@v4 - - name: Build and Test (cmake) - uses: ./.github/actions/build_cmake - with: - opt_level: sve - env: - # Context: https://github.com/facebookresearch/faiss/wiki/Troubleshooting#surprising-faiss-openmp-and-openblas-interaction - OPENBLAS_NUM_THREADS: '1' - linux-x86_64-conda: - name: Linux x86_64 (conda) - needs: linux-x86_64-cmake - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - windows-x86_64-conda: - name: Windows x86_64 (conda) - needs: linux-x86_64-cmake - runs-on: windows-2019 - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - linux-arm64-conda: - name: Linux arm64 (conda) - needs: linux-x86_64-cmake - runs-on: 2-core-ubuntu-arm - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-release.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-release.yml deleted file mode 100644 index b5b02f2..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build-release.yml +++ /dev/null @@ -1,144 +0,0 @@ -on: - workflow_call: - secrets: - ANACONDA_API_TOKEN: - required: true -env: - OMP_NUM_THREADS: '10' - MKL_THREADING_LAYER: GNU -jobs: - linux-x86_64-packages: - name: Linux x86_64 packages - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - linux-x86_64-GPU-packages-CUDA-11-4-4: - name: Linux x86_64 GPU packages (CUDA 11.4.4) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "60-real;61-real;62-real;70-real;72-real;75-real;80;86-real" - FAISS_FLATTEN_CONDA_INCLUDES: "1" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - cuda: "11.4.4" - linux-x86_64-GPU-CUVS-packages-CUDA11-8-0: - name: Linux x86_64 GPU w/ cuVS packages (CUDA 11.8.0) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - cuvs: "ON" - cuda: "11.8.0" - linux-x86_64-GPU-packages-CUDA-12-1-1: - name: Linux x86_64 GPU packages (CUDA 12.1.1) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - cuda: "12.1.1" - linux-x86_64-GPU-CUVS-packages-CUDA12-4-0: - name: Linux x86_64 GPU w/ cuVS packages (CUDA 12.4.0) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - cuvs: "ON" - cuda: "12.4.0" - windows-x86_64-packages: - name: Windows x86_64 packages - runs-on: windows-2019 - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - osx-arm64-packages: - name: OSX arm64 packages - runs-on: macos-14 - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main - linux-arm64-packages: - name: Linux arm64 packages - runs-on: 2-core-ubuntu-arm - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - name: Build and Package (conda) - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: main diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build.yml deleted file mode 100644 index 82792cb..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/build.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: Build -on: - workflow_dispatch: - pull_request: - branches: - - main - push: - tags: - - 'v*' -jobs: - build-pull-request: - uses: ./.github/workflows/build-pull-request.yml - build-release: - uses: ./.github/workflows/build-release.yml - secrets: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/nightly.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/nightly.yml deleted file mode 100644 index ef1e8d2..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/nightly.yml +++ /dev/null @@ -1,148 +0,0 @@ -name: Nightly -on: - schedule: - - cron: '10 6 * * *' -env: - OMP_NUM_THREADS: '10' - MKL_THREADING_LAYER: GNU -jobs: - linux-x86_64-nightly: - name: Linux x86_64 nightlies - runs-on: 4-core-ubuntu - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - linux-x86_64-GPU-CUDA-11-4-4-nightly: - name: Linux x86_64 GPU nightlies (CUDA 11.4.4) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "60-real;61-real;62-real;70-real;72-real;75-real;80;86-real" - FAISS_FLATTEN_CONDA_INCLUDES: "1" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - cuda: "11.4.4" - linux-x86_64-GPU-CUVS-CUDA11-8-0-nightly: - name: Linux x86_64 GPU w/ cuVS nightlies (CUDA 11.8.0) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - cuvs: "ON" - cuda: "11.8.0" - linux-x86_64-GPU-CUDA-12-1-1-nightly: - name: Linux x86_64 GPU nightlies (CUDA 12.1.1) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - cuda: "12.1.1" - linux-x86_64-GPU-CUVS-CUDA12-4-0-nightly: - name: Linux x86_64 GPU w/ cuVS nightlies (CUDA 12.4.0) - runs-on: 4-core-ubuntu-gpu-t4 - env: - CUDA_ARCHS: "70-real;72-real;75-real;80;86-real" - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - cuvs: "ON" - cuda: "12.4.0" - windows-x86_64-nightly: - name: Windows x86_64 nightlies - runs-on: windows-2019 - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - osx-arm64-nightly: - name: OSX arm64 nightlies - runs-on: macos-14 - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - linux-arm64-nightly: - name: Linux arm64 nightlies - runs-on: 2-core-ubuntu-arm - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - fetch-depth: 0 - fetch-tags: true - - uses: ./.github/actions/build_conda - env: - ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} - with: - label: nightly - auto-retry: - name: Auto retry on failure - if: fromJSON(github.run_attempt) < 2 - runs-on: ubuntu-latest - steps: - - name: Start rerun workflow - env: - GH_REPO: ${{ github.repository }} - GH_TOKEN: ${{ github.token }} - GH_DEBUG: api - run: | - gh workflow run retry_build.yml \ - -F run_id=${{ github.run_id }} diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/publish-docs.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/publish-docs.yml deleted file mode 100644 index a75c485..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/publish-docs.yml +++ /dev/null @@ -1,44 +0,0 @@ -name: Publish Docs -on: - page_build: - branches: - - gh-pages - paths-ignore: - - 'docs/**' - workflow_run: - workflows: [update-doxygen] - types: - - completed -jobs: - build_and_publish: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - name: Set up Python - uses: actions/setup-python@v2 - with: - python-version: 3.8 - - name: Checkout gh-pages - run: | - git fetch origin gh-pages - git checkout gh-pages - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -r requirements.txt - - name: Generate html - run: | - make html - git rm -rf docs - mv _build/html docs - touch docs/.nojekyll - - name: Push changes - run: | - git config --global user.email "$GITHUB_ACTOR@users.noreply.github.com" - git config --global user.name "$GITHUB_ACTOR" - git add docs - if [ -n "$(git status --porcelain)" ] - then - git commit docs -m "Sphinx rebuild ($(git rev-parse --short gh-pages))." - git push origin gh-pages - fi diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/retry_build.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/retry_build.yml deleted file mode 100644 index 45c07ff..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/retry_build.yml +++ /dev/null @@ -1,33 +0,0 @@ -name: Retry Build -on: - workflow_dispatch: - inputs: - run_id: - required: true -jobs: - rerun-on-failure: - permissions: write-all - runs-on: ubuntu-latest - steps: - - name: rerun ${{ inputs.run_id }} - env: - GH_REPO: ${{ github.repository }} - GH_TOKEN: ${{ github.token }} - GH_DEBUG: api - run: | - # status can be one of "queued", "in_progress", "completed", "waiting", "requested", "pending" - # https://docs.github.com/en/rest/checks/runs - # while not completed, sleep for 10 minutes - while gh run view ${{ inputs.run_id }} --json status | grep -v completed - do - echo Workflow in progress - sleeping for 10 minutes then checking again - sleep 10m - done - - # Only retry if there are failed jobs - if gh run view ${{ inputs.run_id }} --exit-status; then - echo Workflow succeeded - no retry necessary. - else - echo Workflow failed - initiating retry. - gh run rerun ${{ inputs.run_id }} --failed - fi diff --git a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/update-doxygen.yml b/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/update-doxygen.yml deleted file mode 100644 index 64d9435..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.github/workflows/update-doxygen.yml +++ /dev/null @@ -1,40 +0,0 @@ -name: Update Doxygen -on: - push: - branches: - - main - paths: - - 'faiss/**' -jobs: - doxygen: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - name: Set up Python - uses: actions/setup-python@v2 - with: - python-version: 3.8 - - name: Install dependencies - run: | - sudo apt-get install -y doxygen - python -m pip install --upgrade pip - pip install breathe - - name: Generate doxygen xml - run: doxygen - - name: Push changes - run: | - git config --global user.email "$GITHUB_ACTOR@users.noreply.github.com" - git config --global user.name "$GITHUB_ACTOR" - mkdir ./tmp - mv xml ./tmp/xml - git fetch origin gh-pages - git checkout gh-pages - git rm -rf xml cpp_api - mv ./tmp/xml ./xml - breathe-apidoc -o cpp_api xml - git add xml cpp_api - if [ -n "$(git status --porcelain)" ] - then - git commit -m "Update API docs ($(git rev-parse --short main))." - git push origin gh-pages - fi diff --git a/packages/leann-backend-hnsw/third_party/faiss/.gitignore b/packages/leann-backend-hnsw/third_party/faiss/.gitignore deleted file mode 100644 index 2d5a8dc..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.gitignore +++ /dev/null @@ -1,26 +0,0 @@ -*.swp -*.swo -*.o -*.a -*.dSYM -*.so -*.dylib -*.pyc -*~ -/build/ -/config.* -/aclocal.m4 -/autom4te.cache/ -/makefile.inc -/bin/ -/c_api/bin/ -/c_api/gpu/bin/ -/tests/test -/tests/gtest/ -faiss/python/swigfaiss_avx2.swig -faiss/python/swigfaiss_avx512.swig -faiss/python/swigfaiss_avx512_spr.swig -faiss/python/swigfaiss_sve.swig -.cache/ -compile_commands.json -sift/ \ No newline at end of file diff --git a/packages/leann-backend-hnsw/third_party/faiss/.vscode/launch.json b/packages/leann-backend-hnsw/third_party/faiss/.vscode/launch.json deleted file mode 100644 index d6087f1..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/.vscode/launch.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - // Use IntelliSense to learn about possible attributes. - // Hover to view descriptions of existing attributes. - // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 - "version": "0.2.0", - "configurations": [ - { - "name": "Build Demo", - "type": "lldb", - "request": "launch", - "program": "${workspaceFolder}/../.venv/bin/python", - "console": "integratedTerminal", - "cwd": "${workspaceFolder}", - "args": [ - "${workspaceFolder}/demo/build_demo.py" - ], - }, - ] -} \ No newline at end of file diff --git a/packages/leann-backend-hnsw/third_party/faiss/CHANGELOG.md b/packages/leann-backend-hnsw/third_party/faiss/CHANGELOG.md deleted file mode 100644 index c1771f2..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/CHANGELOG.md +++ /dev/null @@ -1,482 +0,0 @@ -# Changelog -All notable changes to this project will be documented in this file. - -## [Unreleased] - -## [1.10.0] - 2025-01-30 - - -Added -- Add desc_name to dataset descriptor (#3935) -- implement ST_norm_from_LUT for the ResidualQuantizer (#3917) -- Add example of how to build, link, and test an external SWIG module (#3922) -- add copyright header (#3948) -- Add some SVE implementations (#3933) -- Enable linting: lint config changes plus arc lint command (#3966) -- Re-add example of how to build, link, and test an external SWIG module (#3981) -- demo: IndexPQ: separate codes from codebook (#3987) -- add all wrapped indexes to the index_read (#3988) -- add validity check AlignedTableTightAlloc clear method (#3997) -- Add index binary to telemetry (#4001) -- Add VectorTransform read from filename to the C API (#3970) -- Added IndexLSH to the demo (#4009) -- write distributed_kmeans centroids and assignments to hive tables (#4017) -- introduce data splits in dataset descriptor (#4012) -- Faiss GPU: bfloat16 brute-force kNN support (#4018) -- ROCm support for bfloat16 (#4039) -- Unit tests for distances_simd.cpp (#4058) -- add cuda-toolkit for GPU (#4057) -- Add more unit testing for IndexHNSW [1/n] (#4054) -- Add more unit testing for IndexHNSW [2/n] (#4056) -- Add more unit testing for HNSW [3/n] (#4059) -- Add more unit testing for HNSW [4/n] (#4061) -- Add more unit tests for index_read and index_write (#4068) -- Add testing for utils/hamming.cpp (#4079) -- Test sa_decode methd on IndexIVFFlat (#4098) -- Conditionally compile extras like benchmarks and demos (#4094) -- Add a new architecture mode: 'avx512_spr'. (#4025) -- Use _mm512_popcnt_epi64 to speedup hamming distance evaluation. (#4020) -- PQ with pytorch (#4116) -- add range_search() to IndexRefine (#4022) -- Expose accumulate_to_mem from faiss interface (#4099) -- Windows Arm64 support (#4087) -- add test to cover GPU (#4130) -- Added support for building without MKL (#4147) - -Changed -- Move train, build and search to their respective operators (#3934) -- PQFS into Index trainer (#3941) -- Place a useful cmake function 'link_to_faiss_lib' into a separate file (#3939) -- Cache device major version value to avoid multiple calls of getCudaDeviceProperties (#3950) -- Consolidate set_target_properties() calls in faiss/CMakeLists.txt (#3973) -- Removing Manual Hipify Build Step (#3962) -- Allow to replace graph structure for NSG graphs (#3975) -- Adjust nightly build (#3978) -- Update RAFT CI with pytorch 2.4.1 (#3980) -- Moved add_sa_codes, sa_code_size to Index, IndexBinary base classes (#3989) -- Update autoclose.yml (#4000) -- Migrate from RAFT to CUVS (#3549) -- Pin to numpy<2 (#4033) -- (1/n) - Preload datasets in manifold so that subsequent stages of training, indexing and search can use those instead of each trainer or indexer downloading data. (#4034) -- Constrain conda version for Windows build (#4040) -- Updates to faiss-gpu-cuvs nightly pkg (#4032) -- pin the dependecies version for x86_64 (#4046) -- pin arm64 dependency (#4060) -- Pin conda build (#4062) -- Improve naming due to codemod (#4063) -- Improve naming due to codemod (#4064) -- Improve naming due to codemod (#4065) -- separare the github build into two conditions (#4066) -- Improve naming due to codemod (#4070) -- improve naming due to codemod (#4067) -- improve naming due to codemod (#4071) -- improve naming due to codemod (#4072) -- fix nightily build (#4080) -- Change github action workflows name (#4083) -- Resolve Packaging Issues (#4044) -- Update __init__.py (#4086) -- Exhaustive IVF probing in scalar quantizer tests (#4075) -- Pin Nightlies with testing on PR (#4088) -- Update benchmarking library code to work for IdMap index as well (#4093) -- Update action.yml (#4100) -- Upgrade CUVS to 24.12 (#4021) -- Link cuVS Docs (#4084) -- Set KnnDescriptor.desc_name in the Benchmarking core framework in FAISS like other descriptors (#4109) -- enable quiet mode for conda install (#4112) -- Disable retry build (#4124) -- Add ngpu default argument to knn_ground_truth (#4123) -- Update code comment to reflect the range of IF from [1, k] (#4139) -- Reenable auto retry workflow (#4140) -- Migration off defaults to conda-forge channel (#4126) -- Benchmarking Scripts for cuVS Index, more docs updates (#4117) - -Fixed -- Fix total_rows (#3942) -- Fix INSTALL.md due to failure of conflict resolving (#3915) -- Back out "Add example of how to build, link, and test an external SWIG module" (#3954) -- Fix shadowed variable in faiss/IndexPQ.cpp (#3959) -- Fix shadowed variable in faiss/IndexIVFAdditiveQuantizer.cpp (#3958) -- Fix shadowed variable in faiss/impl/HNSW.cpp (#3961) -- Fix shadowed variable in faiss/impl/simd_result_handlers.h (#3960) -- Fix shadowed variable in faiss/utils/NeuralNet.cpp (#3952) -- Resolve "incorrect-portions-license" errors: add no license lint to top of GPU files with both licenses (#3965) -- Resolve "duplicate-license-header": Find and replace duplicate license headers (#3967) -- fix some more nvidia licenses that get erased (#3977) -- fix merge_flat_ondisk stress run failures (#3999) -- Fix reverse_index_factory formatting of ScalarQuantizers (#4003) -- Fix shadowed variable in faiss/IndexAdditiveQuantizer.cpp (#4011) -- facebook-unused-include-check in fbcode/faiss (#4029) -- fix linter (#4035) -- Some chore fixes (#4010) -- Fix unused variable compilation error (#4041) -- stop dealloc of coarse quantizer when it is deleted (#4045) -- Fix SCD Table test flakiness (#4069) -- Fix IndexIVFFastScan reconstruct_from_offset method (#4095) -- more fast-scan reconstruction (#4128) -- Fix nightly cuVS 11.8.0 failure (#4149) -- Correct capitalization of FAISS to Faiss (#4155) -- Fix cuVS 12.4.0 nightly failure (#4153) - -Deprecated -- Remove unused-variable in dumbo/backup/dumbo/service/tests/ChainReplicatorTests.cpp (#4024) -- remove inconsistent oom exception test (#4052) -- Remove unused(and wrong) io macro (#4122) - - -## [1.9.0] - 2024-10-04 -### Added -- Add AVX-512 implementation for the distance and scalar quantizer functions. (#3853) -- Allow k and M suffixes in IVF indexes (#3812) -- add reconstruct support to additive quantizers (#3752) -- introduce options for reducing the overhead for a clustering procedure (#3731) -- Add hnsw search params for bounded queue option (#3748) -- ROCm support (#3462) -- Add sve targets (#2886) -- add get_version() for c_api (#3688) -- QINCo implementation in CPU Faiss (#3608) -- Add search functionality to FlatCodes (#3611) -- add dispatcher for VectorDistance and ResultHandlers (#3627) -- Add SQ8bit signed quantization (#3501) -- Add ABS_INNER_PRODUCT metric (#3524) -- Interop between CAGRA and HNSW (#3252) -- add skip_storage flag to HNSW (#3487) -- QT_bf16 for scalar quantizer for bfloat16 (#3444) -- Implement METRIC.NaNEuclidean (#3414) -- TimeoutCallback C++ and Python (#3417) -- support big-endian machines (#3361) -- Support for Remove ids from IVFPQFastScan index (#3354) -- Implement reconstruct_n for GPU IVFFlat indexes (#3338) -- Support of skip_ids in merge_from_multiple function of OnDiskInvertedLists (#3327) -- Add the ability to clone and read binary indexes to the C API. (#3318) -- AVX512 for PQFastScan (#3276) - -### Changed -- faster hnsw CPU index training (#3822) -- Some small improvements. (#3692) -- First attempt at LSH matching with nbits (#3679) -- Set verbosoe before train (#3619) -- Remove duplicate NegativeDistanceComputer instances (#3450) -- interrupt for NNDescent (#3432) -- Get rid of redundant instructions in ScalarQuantizer (#3430) -- PowerPC, improve code generation for function fvec_L2sqr (#3416) -- Unroll loop in lookup_2_lanes (#3364) -- Improve filtering & search parameters propagation (#3304) -- Change index_cpu_to_gpu to throw for indices not implemented on GPU (#3336) -- Throw when attempting to move IndexPQ to GPU (#3328) -- Skip HNSWPQ sdc init with new io flag (#3250) - -### Fixed -- FIx a bug for a non-simdlib code of ResidualQuantizer (#3868) -- assign_index should default to null (#3855) -- Fix an incorrectly counted the number of computed distances for HNSW (#3840) -- Add error for overflowing nbits during PQ construction (#3833) -- Fix radius search with HSNW and IP (#3698) -- fix algorithm of spreading vectors over shards (#3374) -- Fix IndexBinary.assign Python method (#3384) -- Few fixes in bench_fw to enable IndexFromCodec (#3383) -- Fix the endianness issue in AIX while running the benchmark. (#3345) -- Fix faiss swig build with version > 4.2.x (#3315) -- Fix problems when using 64-bit integers. (#3322) -- Fix IVFPQFastScan decode function (#3312) -- Handling FaissException in few destructors of ResultHandler.h (#3311) -- Fix HNSW stats (#3309) -- AIX compilation fix for io classes (#3275) - - -## [1.8.0] - 2024-02-27 -### Added -- Added a new conda package faiss-gpu-raft alongside faiss-cpu and faiss-gpu -- Integrated IVF-Flat and IVF-PQ implementations in faiss-gpu-raft from RAFT by Nvidia [thanks Corey Nolet and Tarang Jain] -- Added a context parameter to InvertedLists and InvertedListsIterator -- Added Faiss on Rocksdb demo to showing how inverted lists can be persisted in a key-value store -- Introduced Offline IVF framework powered by Faiss big batch search -- Added SIMD NEON Optimization for QT_FP16 in Scalar Quantizer. [thanks Naveen Tatikonda] -- Generalized ResultHandler and supported range search for HNSW and FastScan -- Introduced avx512 optimization mode and FAISS_OPT_LEVEL env variable [thanks Alexandr Ghuzva] -- Added search parameters for IndexRefine::search() and IndexRefineFlat::search() -- Supported large two-level clustering -- Added support for Python 3.11 and 3.12 -- Added support for CUDA 12 - -### Changed -- Used the benchmark to find Pareto optimal indices. Intentionally limited to IVF(Flat|HNSW),PQ|SQ indices -- Splitted off RQ encoding steps to another file -- Supported better NaN handling -- HNSW speedup + Distance 4 points [thanks Alexandr Ghuzva] - -### Fixed -- Fixed DeviceVector reallocations in Faiss GPU -- Used efSearch from params if provided in HNSW search -- Fixed warp synchronous behavior in Faiss GPU CUDA 12 - - -## [1.7.4] - 2023-04-12 -### Added -- Added big batch IVF search for conducting efficient search with big batches of queries -- Checkpointing in big batch search support -- Precomputed centroids support -- Support for iterable inverted lists for eg. key value stores -- 64-bit indexing arithmetic support in FAISS GPU -- IndexIVFShards now handle IVF indexes with a common quantizer -- Jaccard distance support -- CodePacker for non-contiguous code layouts -- Approximate evaluation of top-k distances for ResidualQuantizer and IndexBinaryFlat -- Added support for 12-bit PQ / IVFPQ fine quantizer decoders for standalone vector codecs (faiss/cppcontrib) -- Conda packages for osx-arm64 (Apple M1) and linux-aarch64 (ARM64) architectures -- Support for Python 3.10 - -### Removed -- CUDA 10 is no longer supported in precompiled packages -- Removed Python 3.7 support for precompiled packages -- Removed constraint for using fine quantizer with no greater than 8 bits for IVFPQ, for example, now it is possible to use IVF256,PQ10x12 for a CPU index - -### Changed -- Various performance optimizations for PQ / IVFPQ for AVX2 and ARM for training (fused distance+nearest kernel), search (faster kernels for distance_to_code() and scan_list_*()) and vector encoding -- A magnitude faster CPU code for LSQ/PLSQ training and vector encoding (reworked code) -- Performance improvements for Hamming Code computations for AVX2 and ARM (reworked code) -- Improved auto-vectorization support for IP and L2 distance computations (better handling of pragmas) -- Improved ResidualQuantizer vector encoding (pooling memory allocations, avoid r/w to a temporary buffer) - -### Fixed -- HSNW bug fixed which improves the recall rate! Special thanks to zh Wang @hhy3 for this. -- Faiss GPU IVF large query batch fix -- Faiss + Torch fixes, re-enable k = 2048 -- Fix the number of distance computations to match max_codes parameter -- Fix decoding of large fast_scan blocks - - -## [1.7.3] - 2022-11-3 -### Added -- Added sparse k-means routines and moved the generic kmeans to contrib -- Added FlatDistanceComputer for all FlatCodes indexes -- Support for fast accumulation of 4-bit LSQ and RQ -- Added product additive quantization -- Support per-query search parameters for many indexes + filtering by ids -- write_VectorTransform and read_vectorTransform were added to the public API (by @AbdelrahmanElmeniawy) -- Support for IDMap2 in index_factory by adding "IDMap2" to prefix or suffix of the input String (by @AbdelrahmanElmeniawy) -- Support for merging all IndexFlatCodes descendants (by @AbdelrahmanElmeniawy) -- Remove and merge features for IndexFastScan (by @AbdelrahmanElmeniawy) -- Performance improvements: 1) specialized the AVX2 pieces of code speeding up certain hotspots, 2) specialized kernels for vector codecs (this can be found in faiss/cppcontrib) - - -### Fixed -- Fixed memory leak in OnDiskInvertedLists::do_mmap when the file is not closed (by @AbdelrahmanElmeniawy) -- LSH correctly throws error for metric types other than METRIC_L2 (by @AbdelrahmanElmeniawy) - -## [1.7.2] - 2021-12-15 -### Added -- Support LSQ on GPU (by @KinglittleQ) -- Support for exact 1D kmeans (by @KinglittleQ) - -## [1.7.1] - 2021-05-27 -### Added -- Support for building C bindings through the `FAISS_ENABLE_C_API` CMake option. -- Serializing the indexes with the python pickle module -- Support for the NNDescent k-NN graph building method (by @KinglittleQ) -- Support for the NSG graph indexing method (by @KinglittleQ) -- Residual quantizers: support as codec and unoptimized search -- Support for 4-bit PQ implementation for ARM (by @vorj, @n-miyamoto-fixstars, @LWisteria, and @matsui528) -- Implementation of Local Search Quantization (by @KinglittleQ) - -### Changed -- The order of xb an xq was different between `faiss.knn` and `faiss.knn_gpu`. -Also the metric argument was called distance_type. -- The typed vectors (LongVector, LongLongVector, etc.) of the SWIG interface have -been deprecated. They have been replaced with Int32Vector, Int64Vector, etc. (by h-vetinari) - -### Fixed -- Fixed a bug causing kNN search functions for IndexBinaryHash and -IndexBinaryMultiHash to return results in a random order. -- Copy constructor of AlignedTable had a bug leading to crashes when cloning -IVFPQ indices. - -## [1.7.0] - 2021-01-27 - -## [1.6.5] - 2020-11-22 - -## [1.6.4] - 2020-10-12 -### Added -- Arbitrary dimensions per sub-quantizer now allowed for `GpuIndexIVFPQ`. -- Brute-force kNN on GPU (`bfKnn`) now accepts `int32` indices. -- Nightly conda builds now available (for CPU). -- Faiss is now supported on Windows. - -## [1.6.3] - 2020-03-24 -### Added -- Support alternative distances on GPU for GpuIndexFlat, including L1, Linf and -Lp metrics. -- Support METRIC_INNER_PRODUCT for GpuIndexIVFPQ. -- Support float16 coarse quantizer for GpuIndexIVFFlat and GpuIndexIVFPQ. GPU -Tensor Core operations (mixed-precision arithmetic) are enabled on supported -hardware when operating with float16 data. -- Support k-means clustering with encoded vectors. This makes it possible to -train on larger datasets without decompressing them in RAM, and is especially -useful for binary datasets (see https://github.com/facebookresearch/faiss/blob/main/tests/test_build_blocks.py#L92). -- Support weighted k-means. Weights can be associated to each training point -(see https://github.com/facebookresearch/faiss/blob/main/tests/test_build_blocks.py). -- Serialize callback in python, to write to pipes or sockets (see -https://github.com/facebookresearch/faiss/wiki/Index-IO,-cloning-and-hyper-parameter-tuning). -- Reconstruct arbitrary ids from IndexIVF + efficient remove of a small number -of ids. This avoids 2 inefficiencies: O(ntotal) removal of vectors and -IndexIDMap2 on top of indexIVF. Documentation here: -https://github.com/facebookresearch/faiss/wiki/Special-operations-on-indexes. -- Support inner product as a metric in IndexHNSW (see -https://github.com/facebookresearch/faiss/blob/main/tests/test_index.py#L490). -- Support PQ of sizes other than 8 bit in IndexIVFPQ. -- Demo on how to perform searches sequentially on an IVF index. This is useful -for an OnDisk index with a very large batch of queries. In that case, it is -worthwhile to scan the index sequentially (see -https://github.com/facebookresearch/faiss/blob/main/tests/test_ivflib.py#L62). -- Range search support for most binary indexes. -- Support for hashing-based binary indexes (see -https://github.com/facebookresearch/faiss/wiki/Binary-indexes). - -### Changed -- Replaced obj table in Clustering object: now it is a ClusteringIterationStats -structure that contains additional statistics. - -### Removed -- Removed support for useFloat16Accumulator for accumulators on GPU (all -accumulations are now done in float32, regardless of whether float16 or float32 -input data is used). - -### Fixed -- Some python3 fixes in benchmarks. -- Fixed GpuCloner (some fields were not copied, default to no precomputed tables -with IndexIVFPQ). -- Fixed support for new pytorch versions. -- Serialization bug with alternative distances. -- Removed test on multiple-of-4 dimensions when switching between blas and AVX -implementations. - -## [1.6.2] - 2020-03-10 - -## [1.6.1] - 2019-12-04 - -## [1.6.0] - 2019-09-24 -### Added -- Faiss as a codec: We introduce a new API within Faiss to encode fixed-size -vectors into fixed-size codes. The encoding is lossy and the tradeoff between -compression and reconstruction accuracy can be adjusted. -- ScalarQuantizer support for GPU, see gpu/GpuIndexIVFScalarQuantizer.h. This is -particularly useful as GPU memory is often less abundant than CPU. -- Added easy-to-use serialization functions for indexes to byte arrays in Python -(faiss.serialize_index, faiss.deserialize_index). -- The Python KMeans object can be used to use the GPU directly, just add -gpu=True to the constuctor see gpu/test/test_gpu_index.py test TestGPUKmeans. - -### Changed -- Change in the code layout: many C++ sources are now in subdirectories impl/ -and utils/. - -## [1.5.3] - 2019-06-24 -### Added -- Basic support for 6 new metrics in CPU IndexFlat and IndexHNSW (https://github.com/facebookresearch/faiss/issues/848). -- Support for IndexIDMap/IndexIDMap2 with binary indexes (https://github.com/facebookresearch/faiss/issues/780). - -### Changed -- Throw python exception for OOM (https://github.com/facebookresearch/faiss/issues/758). -- Make DistanceComputer available for all random access indexes. -- Gradually moving from long to uint64_t for portability. - -### Fixed -- Slow scanning of inverted lists (https://github.com/facebookresearch/faiss/issues/836). - -## [1.5.2] - 2019-05-28 -### Added -- Support for searching several inverted lists in parallel (parallel_mode != 0). -- Better support for PQ codes where nbit != 8 or 16. -- IVFSpectralHash implementation: spectral hash codes inside an IVF. -- 6-bit per component scalar quantizer (4 and 8 bit were already supported). -- Combinations of inverted lists: HStackInvertedLists and VStackInvertedLists. -- Configurable number of threads for OnDiskInvertedLists prefetching (including -0=no prefetch). -- More test and demo code compatible with Python 3 (print with parentheses). - -### Changed -- License was changed from BSD+Patents to MIT. -- Exceptions raised in sub-indexes of IndexShards and IndexReplicas are now -propagated. -- Refactored benchmark code: data loading is now in a single file. - -## [1.5.1] - 2019-04-05 -### Added -- MatrixStats object, which reports useful statistics about a dataset. -- Option to round coordinates during k-means optimization. -- An alternative option for search in HNSW. -- Support for range search in IVFScalarQuantizer. -- Support for direct uint_8 codec in ScalarQuantizer. -- Better support for PQ code assignment with external index. -- Support for IMI2x16 (4B virtual centroids). -- Support for k = 2048 search on GPU (instead of 1024). -- Support for renaming an ondisk invertedlists. -- Support for nterrupting computations with interrupt signal (ctrl-C) in python. -- Simplified build system (with --with-cuda/--with-cuda-arch options). - -### Changed -- Moved stats() and imbalance_factor() from IndexIVF to InvertedLists object. -- Renamed IndexProxy to IndexReplicas. -- Most CUDA mem alloc failures now throw exceptions instead of terminating on an -assertion. -- Updated example Dockerfile. -- Conda packages now depend on the cudatoolkit packages, which fixes some -interferences with pytorch. Consequentially, faiss-gpu should now be installed -by conda install -c pytorch faiss-gpu cudatoolkit=10.0. - -## [1.5.0] - 2018-12-19 -### Added -- New GpuIndexBinaryFlat index. -- New IndexBinaryHNSW index. - -## [1.4.0] - 2018-08-30 -### Added -- Automatic tracking of C++ references in Python. -- Support for non-intel platforms, some functions optimized for ARM. -- Support for overriding nprobe for concurrent searches. -- Support for floating-point quantizers in binary indices. - -### Fixed -- No more segfaults due to Python's GC. -- GpuIndexIVFFlat issues for float32 with 64 / 128 dims. -- Sharding of flat indexes on GPU with index_cpu_to_gpu_multiple. - -## [1.3.0] - 2018-07-10 -### Added -- Support for binary indexes (IndexBinaryFlat, IndexBinaryIVF). -- Support fp16 encoding in scalar quantizer. -- Support for deduplication in IndexIVFFlat. -- Support for index serialization. - -### Fixed -- MMAP bug for normal indices. -- Propagation of io_flags in read func. -- k-selection for CUDA 9. -- Race condition in OnDiskInvertedLists. - -## [1.2.1] - 2018-02-28 -### Added -- Support for on-disk storage of IndexIVF data. -- C bindings. -- Extended tutorial to GPU indices. - -[Unreleased]: https://github.com/facebookresearch/faiss/compare/v1.9.0...HEAD -[1.9.0]: https://github.com/facebookresearch/faiss/compare/v1.8.0...v1.9.0 -[1.8.0]: https://github.com/facebookresearch/faiss/compare/v1.7.4...v1.8.0 -[1.7.4]: https://github.com/facebookresearch/faiss/compare/v1.7.3...v1.7.4 -[1.7.3]: https://github.com/facebookresearch/faiss/compare/v1.7.2...v1.7.3 -[1.7.2]: https://github.com/facebookresearch/faiss/compare/v1.7.1...v1.7.2 -[1.7.1]: https://github.com/facebookresearch/faiss/compare/v1.7.0...v1.7.1 -[1.7.0]: https://github.com/facebookresearch/faiss/compare/v1.6.5...v1.7.0 -[1.6.5]: https://github.com/facebookresearch/faiss/compare/v1.6.4...v1.6.5 -[1.6.4]: https://github.com/facebookresearch/faiss/compare/v1.6.3...v1.6.4 -[1.6.3]: https://github.com/facebookresearch/faiss/compare/v1.6.2...v1.6.3 -[1.6.2]: https://github.com/facebookresearch/faiss/compare/v1.6.1...v1.6.2 -[1.6.1]: https://github.com/facebookresearch/faiss/compare/v1.6.0...v1.6.1 -[1.6.0]: https://github.com/facebookresearch/faiss/compare/v1.5.3...v1.6.0 -[1.5.3]: https://github.com/facebookresearch/faiss/compare/v1.5.2...v1.5.3 -[1.5.2]: https://github.com/facebookresearch/faiss/compare/v1.5.1...v1.5.2 -[1.5.1]: https://github.com/facebookresearch/faiss/compare/v1.5.0...v1.5.1 -[1.5.0]: https://github.com/facebookresearch/faiss/compare/v1.4.0...v1.5.0 -[1.4.0]: https://github.com/facebookresearch/faiss/compare/v1.3.0...v1.4.0 -[1.3.0]: https://github.com/facebookresearch/faiss/compare/v1.2.1...v1.3.0 -[1.2.1]: https://github.com/facebookresearch/faiss/releases/tag/v1.2.1 diff --git a/packages/leann-backend-hnsw/third_party/faiss/CMakeLists.txt b/packages/leann-backend-hnsw/third_party/faiss/CMakeLists.txt deleted file mode 100644 index 4a70aaf..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/CMakeLists.txt +++ /dev/null @@ -1,126 +0,0 @@ -# @lint-ignore-every LICENSELINT -# Copyright (c) Facebook, Inc. and its affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# ============================================================================= -# Copyright (c) 2023-2024, NVIDIA CORPORATION. -# -# 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. -# ============================================================================= - -cmake_minimum_required(VERSION 3.24.0 FATAL_ERROR) - -set(FAISS_LANGUAGES CXX) - -if(FAISS_ENABLE_GPU) - if (FAISS_ENABLE_ROCM) - list(APPEND FAISS_LANGUAGES HIP) - list(PREPEND CMAKE_MODULE_PATH "/opt/rocm/lib/cmake") - list(PREPEND CMAKE_PREFIX_PATH "/opt/rocm") - else() - list(APPEND FAISS_LANGUAGES CUDA) - endif() -endif() - -if(FAISS_ENABLE_CUVS) -include(cmake/thirdparty/fetch_rapids.cmake) -include(rapids-cmake) -include(rapids-cpm) -include(rapids-cuda) -include(rapids-export) -include(rapids-find) - -rapids_cuda_init_architectures(faiss) -rapids_cuda_init_architectures(pyfaiss) -rapids_cuda_init_architectures(faiss_c_library) -endif() - -project(faiss - VERSION 1.10.0 - DESCRIPTION "A library for efficient similarity search and clustering of dense vectors." - HOMEPAGE_URL "https://github.com/facebookresearch/faiss" - LANGUAGES ${FAISS_LANGUAGES}) -include(GNUInstallDirs) - -set(CMAKE_CXX_STANDARD 17) - -list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake") - -# Valid values are "generic", "avx2", "avx512", "avx512_spr", "sve". -option(FAISS_OPT_LEVEL "" "generic") -option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON) -option(FAISS_ENABLE_CUVS "Enable cuVS for GPU indexes." OFF) -option(FAISS_ENABLE_ROCM "Enable ROCm for GPU indexes." OFF) -option(FAISS_ENABLE_MKL "Enable MKL." ON) -option(FAISS_ENABLE_PYTHON "Build Python extension." ON) -option(FAISS_ENABLE_C_API "Build C API." OFF) -option(FAISS_ENABLE_EXTRAS "Build extras like benchmarks and demos" ON) -option(FAISS_USE_LTO "Enable Link-Time optimization" OFF) - -if(FAISS_ENABLE_GPU) - if(FAISS_ENABLE_ROCM) - enable_language(HIP) - add_definitions(-DUSE_AMD_ROCM) - find_package(HIP REQUIRED) - find_package(hipBLAS REQUIRED) - set(GPU_EXT_PREFIX "hip") - execute_process(COMMAND ${PROJECT_SOURCE_DIR}/faiss/gpu/hipify.sh) - else () - set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER}) - enable_language(CUDA) - set(GPU_EXT_PREFIX "cu") - endif() -endif() - -if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs) - find_package(cuvs) - endif() - -add_subdirectory(faiss) - -if(FAISS_ENABLE_GPU) - if(FAISS_ENABLE_ROCM) - add_subdirectory(faiss/gpu-rocm) - else() - add_subdirectory(faiss/gpu) - endif() -endif() - -if(FAISS_ENABLE_PYTHON) - add_subdirectory(faiss/python) -endif() - -if(FAISS_ENABLE_C_API) - add_subdirectory(c_api) -endif() - -if(FAISS_ENABLE_EXTRAS) - add_subdirectory(demos) - add_subdirectory(benchs) - add_subdirectory(tutorial/cpp) -endif() - -# CTest must be included in the top level to enable `make test` target. -include(CTest) -if(BUILD_TESTING) - add_subdirectory(tests) - add_subdirectory(perf_tests) - if(FAISS_ENABLE_GPU) - if(FAISS_ENABLE_ROCM) - add_subdirectory(faiss/gpu-rocm/test) - else() - add_subdirectory(faiss/gpu/test) - endif() - endif() -endif() diff --git a/packages/leann-backend-hnsw/third_party/faiss/CODE_OF_CONDUCT.md b/packages/leann-backend-hnsw/third_party/faiss/CODE_OF_CONDUCT.md deleted file mode 100644 index ac27d8a..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/CODE_OF_CONDUCT.md +++ /dev/null @@ -1,2 +0,0 @@ -# Code of Conduct -Facebook has adopted a Code of Conduct that we expect project participants to adhere to. Please [read the full text](https://code.fb.com/codeofconduct) so that you can understand what actions will and will not be tolerated. \ No newline at end of file diff --git a/packages/leann-backend-hnsw/third_party/faiss/CONTRIBUTING.md b/packages/leann-backend-hnsw/third_party/faiss/CONTRIBUTING.md deleted file mode 100644 index 10fc815..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/CONTRIBUTING.md +++ /dev/null @@ -1,52 +0,0 @@ -# Contributing to Faiss - -We want to make contributing to this project as easy and transparent as -possible. - -## Our Development Process - -We mainly develop Faiss within Facebook. Sometimes, we will sync the -github version of Faiss with the internal state. - -## Pull Requests - -We welcome pull requests that add significant value to Faiss. If you plan to do -a major development and contribute it back to Faiss, please contact us first before -putting too much effort into it. - -1. Fork the repo and create your branch from `main`. -2. If you've added code that should be tested, add tests. -3. If you've changed APIs, update the documentation. -4. Ensure the test suite passes. -5. Make sure your code lints. -6. If you haven't already, complete the Contributor License Agreement ("CLA"). - -There is a Facebook internal test suite for Faiss, and we need to run -all changes to Faiss through it. - -## Contributor License Agreement ("CLA") - -In order to accept your pull request, we need you to submit a CLA. You only need -to do this once to work on any of Facebook's open source projects. - -Complete your CLA here: - -## Issues - -We use GitHub issues to track public bugs. Please ensure your description is -clear and has sufficient instructions to be able to reproduce the issue. - -Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe -disclosure of security bugs. In those cases, please go through the process -outlined on that page and do not file a public issue. - -## Coding Style - -* 4 spaces for indentation in C++ (no tabs) -* 80 character line length (both for C++ and Python) -* C++ language level: C++17 - -## License - -By contributing to Faiss, you agree that your contributions will be licensed -under the LICENSE file in the root directory of this source tree. diff --git a/packages/leann-backend-hnsw/third_party/faiss/Doxyfile b/packages/leann-backend-hnsw/third_party/faiss/Doxyfile deleted file mode 100644 index 3a112d0..0000000 --- a/packages/leann-backend-hnsw/third_party/faiss/Doxyfile +++ /dev/null @@ -1,2282 +0,0 @@ - - -# Doxyfile 1.8.5 - -# This file describes the settings to be used by the documentation system -# doxygen (www.doxygen.org) for a project. -# -# All text after a double hash (##) is considered a comment and is placed in -# front of the TAG it is preceding. -# -# All text after a single hash (#) is considered a comment and will be ignored. -# The format is: -# TAG = value [value, ...] -# For lists, items can also be appended using: -# TAG += value [value, ...] -# Values that contain spaces should be placed between quotes (\" \"). - -#--------------------------------------------------------------------------- -# Project related configuration options -#--------------------------------------------------------------------------- - -# This tag specifies the encoding used for all characters in the config file -# that follow. The default is UTF-8 which is also the encoding used for all text -# before the first occurrence of this tag. Doxygen uses libiconv (or the iconv -# built into libc) for the transcoding. See http://www.gnu.org/software/libiconv -# for the list of possible encodings. -# The default value is: UTF-8. - -DOXYFILE_ENCODING = UTF-8 - -# The PROJECT_NAME tag is a single word (or a sequence of words surrounded by -# double-quotes, unless you are using Doxywizard) that should identify the -# project for which the documentation is generated. This name is used in the -# title of most generated pages and in a few other places. -# The default value is: My Project. - -PROJECT_NAME = "Faiss" - -# The PROJECT_NUMBER tag can be used to enter a project or revision number. This -# could be handy for archiving the generated documentation or if some version -# control system is used. - -PROJECT_NUMBER = - -# Using the PROJECT_BRIEF tag one can provide an optional one line description -# for a project that appears at the top of each page and should give viewer a -# quick idea about the purpose of the project. Keep the description short. - -PROJECT_BRIEF = - -# With the PROJECT_LOGO tag one can specify an logo or icon that is included in -# the documentation. The maximum height of the logo should not exceed 55 pixels -# and the maximum width should not exceed 200 pixels. Doxygen will copy the logo -# to the output directory. - -PROJECT_LOGO = - -# The OUTPUT_DIRECTORY tag is used to specify the (relative or absolute) path -# into which the generated documentation will be written. If a relative path is -# entered, it will be relative to the location where doxygen was started. If -# left blank the current directory will be used. - -OUTPUT_DIRECTORY = - -# If the CREATE_SUBDIRS tag is set to YES, then doxygen will create 4096 sub- -# directories (in 2 levels) under the output directory of each output format and -# will distribute the generated files over these directories. Enabling this -# option can be useful when feeding doxygen a huge amount of source files, where -# putting all generated files in the same directory would otherwise causes -# performance problems for the file system. -# The default value is: NO. - -CREATE_SUBDIRS = NO - -# The OUTPUT_LANGUAGE tag is used to specify the language in which all -# documentation generated by doxygen is written. Doxygen will use this -# information to generate all constant output in the proper language. -# Possible values are: Afrikaans, Arabic, Brazilian, Catalan, Chinese, Chinese- -# Traditional, Croatian, Czech, Danish, Dutch, English, Esperanto, Farsi, -# Finnish, French, German, Greek, Hungarian, Italian, Japanese, Japanese-en, -# Korean, Korean-en, Latvian, Norwegian, Macedonian, Persian, Polish, -# Portuguese, Romanian, Russian, Serbian, Slovak, Slovene, Spanish, Swedish, -# Turkish, Ukrainian and Vietnamese. -# The default value is: English. - -OUTPUT_LANGUAGE = English - -# If the BRIEF_MEMBER_DESC tag is set to YES doxygen will include brief member -# descriptions after the members that are listed in the file and class -# documentation (similar to Javadoc). Set to NO to disable this. -# The default value is: YES. - -BRIEF_MEMBER_DESC = YES - -# If the REPEAT_BRIEF tag is set to YES doxygen will prepend the brief -# description of a member or function before the detailed description -# -# Note: If both HIDE_UNDOC_MEMBERS and BRIEF_MEMBER_DESC are set to NO, the -# brief descriptions will be completely suppressed. -# The default value is: YES. - -REPEAT_BRIEF = YES - -# This tag implements a quasi-intelligent brief description abbreviator that is -# used to form the text in various listings. Each string in this list, if found -# as the leading text of the brief description, will be stripped from the text -# and the result, after processing the whole list, is used as the annotated -# text. Otherwise, the brief description is used as-is. If left blank, the -# following values are used ($name is automatically replaced with the name of -# the entity):The $name class, The $name widget, The $name file, is, provides, -# specifies, contains, represents, a, an and the. - -ABBREVIATE_BRIEF = - -# If the ALWAYS_DETAILED_SEC and REPEAT_BRIEF tags are both set to YES then -# doxygen will generate a detailed section even if there is only a brief -# description. -# The default value is: NO. - -ALWAYS_DETAILED_SEC = NO - -# If the INLINE_INHERITED_MEMB tag is set to YES, doxygen will show all -# inherited members of a class in the documentation of that class as if those -# members were ordinary class members. Constructors, destructors and assignment -# operators of the base classes will not be shown. -# The default value is: NO. - -INLINE_INHERITED_MEMB = YES - -# If the FULL_PATH_NAMES tag is set to YES doxygen will prepend the full path -# before files name in the file list and in the header files. If set to NO the -# shortest path that makes the file name unique will be used -# The default value is: YES. - -FULL_PATH_NAMES = YES - -# The STRIP_FROM_PATH tag can be used to strip a user-defined part of the path. -# Stripping is only done if one of the specified strings matches the left-hand -# part of the path. The tag can be used to show relative paths in the file list. -# If left blank the directory from which doxygen is run is used as the path to -# strip. -# -# Note that you can specify absolute paths here, but also relative paths, which -# will be relative from the directory where doxygen is started. -# This tag requires that the tag FULL_PATH_NAMES is set to YES. - -STRIP_FROM_PATH = - -# The STRIP_FROM_INC_PATH tag can be used to strip a user-defined part of the -# path mentioned in the documentation of a class, which tells the reader which -# header file to include in order to use a class. If left blank only the name of -# the header file containing the class definition is used. Otherwise one should -# specify the list of include paths that are normally passed to the compiler -# using the -I flag. - -STRIP_FROM_INC_PATH = - -# If the SHORT_NAMES tag is set to YES, doxygen will generate much shorter (but -# less readable) file names. This can be useful is your file systems doesn't -# support long names like on DOS, Mac, or CD-ROM. -# The default value is: NO. - -SHORT_NAMES = NO - -# If the JAVADOC_AUTOBRIEF tag is set to YES then doxygen will interpret the -# first line (until the first dot) of a Javadoc-style comment as the brief -# description. If set to NO, the Javadoc-style will behave just like regular Qt- -# style comments (thus requiring an explicit @brief command for a brief -# description.) -# The default value is: NO. - -JAVADOC_AUTOBRIEF = NO - -# If the QT_AUTOBRIEF tag is set to YES then doxygen will interpret the first -# line (until the first dot) of a Qt-style comment as the brief description. If -# set to NO, the Qt-style will behave just like regular Qt-style comments (thus -# requiring an explicit \brief command for a brief description.) -# The default value is: NO. - -QT_AUTOBRIEF = NO - -# The MULTILINE_CPP_IS_BRIEF tag can be set to YES to make doxygen treat a -# multi-line C++ special comment block (i.e. a block of //! or /// comments) as -# a brief description. This used to be the default behavior. The new default is -# to treat a multi-line C++ comment block as a detailed description. Set this -# tag to YES if you prefer the old behavior instead. -# -# Note that setting this tag to YES also means that rational rose comments are -# not recognized any more. -# The default value is: NO. - -MULTILINE_CPP_IS_BRIEF = NO - -# If the INHERIT_DOCS tag is set to YES then an undocumented member inherits the -# documentation from any documented member that it re-implements. -# The default value is: YES. - -INHERIT_DOCS = YES - -# If the SEPARATE_MEMBER_PAGES tag is set to YES, then doxygen will produce a -# new page for each member. If set to NO, the documentation of a member will be -# part of the file/class/namespace that contains it. -# The default value is: NO. - -SEPARATE_MEMBER_PAGES = NO - -# The TAB_SIZE tag can be used to set the number of spaces in a tab. Doxygen -# uses this value to replace tabs by spaces in code fragments. -# Minimum value: 1, maximum value: 16, default value: 4. - -TAB_SIZE = 4 - -# This tag can be used to specify a number of aliases that act as commands in -# the documentation. An alias has the form: -# name=value -# For example adding -# "sideeffect=@par Side Effects:\n" -# will allow you to put the command \sideeffect (or @sideeffect) in the -# documentation, which will result in a user-defined paragraph with heading -# "Side Effects:". You can put \n's in the value part of an alias to insert -# newlines. - -ALIASES = - -# This tag can be used to specify a number of word-keyword mappings (TCL only). -# A mapping has the form "name=value". For example adding "class=itcl::class" -# will allow you to use the command class in the itcl::class meaning. - -TCL_SUBST = - -# Set the OPTIMIZE_OUTPUT_FOR_C tag to YES if your project consists of C sources -# only. Doxygen will then generate output that is more tailored for C. For -# instance, some of the names that are used will be different. The list of all -# members will be omitted, etc. -# The default value is: NO. - -OPTIMIZE_OUTPUT_FOR_C = NO - -# Set the OPTIMIZE_OUTPUT_JAVA tag to YES if your project consists of Java or -# Python sources only. Doxygen will then generate output that is more tailored -# for that language. For instance, namespaces will be presented as packages, -# qualified scopes will look different, etc. -# The default value is: NO. - -OPTIMIZE_OUTPUT_JAVA = NO - -# Set the OPTIMIZE_FOR_FORTRAN tag to YES if your project consists of Fortran -# sources. Doxygen will then generate output that is tailored for Fortran. -# The default value is: NO. - -OPTIMIZE_FOR_FORTRAN = NO - -# Set the OPTIMIZE_OUTPUT_VHDL tag to YES if your project consists of VHDL -# sources. Doxygen will then generate output that is tailored for VHDL. -# The default value is: NO. - -OPTIMIZE_OUTPUT_VHDL = NO - -# Doxygen selects the parser to use depending on the extension of the files it -# parses. With this tag you can assign which parser to use for a given -# extension. Doxygen has a built-in mapping, but you can override or extend it -# using this tag. The format is ext=language, where ext is a file extension, and -# language is one of the parsers supported by doxygen: IDL, Java, Javascript, -# C#, C, C++, D, PHP, Objective-C, Python, Fortran, VHDL. For instance to make -# doxygen treat .inc files as Fortran files (default is PHP), and .f files as C -# (default is Fortran), use: inc=Fortran f=C. -# -# Note For files without extension you can use no_extension as a placeholder. -# -# Note that for custom extensions you also need to set FILE_PATTERNS otherwise -# the files are not read by doxygen. - -EXTENSION_MAPPING = - -# If the MARKDOWN_SUPPORT tag is enabled then doxygen pre-processes all comments -# according to the Markdown format, which allows for more readable -# documentation. See http://daringfireball.net/projects/markdown/ for details. -# The output of markdown processing is further processed by doxygen, so you can -# mix doxygen, HTML, and XML commands with Markdown formatting. Disable only in -# case of backward compatibilities issues. -# The default value is: YES. - -MARKDOWN_SUPPORT = YES - -# When enabled doxygen tries to link words that correspond to documented -# classes, or namespaces to their corresponding documentation. Such a link can -# be prevented in individual cases by by putting a % sign in front of the word -# or globally by setting AUTOLINK_SUPPORT to NO. -# The default value is: YES. - -AUTOLINK_SUPPORT = YES - -# If you use STL classes (i.e. std::string, std::vector, etc.) but do not want -# to include (a tag file for) the STL sources as input, then you should set this -# tag to YES in order to let doxygen match functions declarations and -# definitions whose arguments contain STL classes (e.g. func(std::string); -# versus func(std::string) {}). This also make the inheritance and collaboration -# diagrams that involve STL classes more complete and accurate. -# The default value is: NO. - -BUILTIN_STL_SUPPORT = NO - -# If you use Microsoft's C++/CLI language, you should set this option to YES to -# enable parsing support. -# The default value is: NO. - -CPP_CLI_SUPPORT = NO - -# Set the SIP_SUPPORT tag to YES if your project consists of sip (see: -# http://www.riverbankcomputing.co.uk/software/sip/intro) sources only. Doxygen -# will parse them like normal C++ but will assume all classes use public instead -# of private inheritance when no explicit protection keyword is present. -# The default value is: NO. - -SIP_SUPPORT = NO - -# For Microsoft's IDL there are propget and propput attributes to indicate -# getter and setter methods for a property. Setting this option to YES will make -# doxygen to replace the get and set methods by a property in the documentation. -# This will only work if the methods are indeed getting or setting a simple -# type. If this is not the case, or you want to show the methods anyway, you -# should set this option to NO. -# The default value is: YES. - -IDL_PROPERTY_SUPPORT = YES - -# If member grouping is used in the documentation and the DISTRIBUTE_GROUP_DOC -# tag is set to YES, then doxygen will reuse the documentation of the first -# member in the group (if any) for the other members of the group. By default -# all members of a group must be documented explicitly. -# The default value is: NO. - -DISTRIBUTE_GROUP_DOC = NO - -# Set the SUBGROUPING tag to YES to allow class member groups of the same type -# (for instance a group of public functions) to be put as a subgroup of that -# type (e.g. under the Public Functions section). Set it to NO to prevent -# subgrouping. Alternatively, this can be done per class using the -# \nosubgrouping command. -# The default value is: YES. - -SUBGROUPING = YES - -# When the INLINE_GROUPED_CLASSES tag is set to YES, classes, structs and unions -# are shown inside the group in which they are included (e.g. using \ingroup) -# instead of on a separate page (for HTML and Man pages) or section (for LaTeX -# and RTF). -# -# Note that this feature does not work in combination with -# SEPARATE_MEMBER_PAGES. -# The default value is: NO. - -INLINE_GROUPED_CLASSES = NO - -# When the INLINE_SIMPLE_STRUCTS tag is set to YES, structs, classes, and unions -# with only public data fields or simple typedef fields will be shown inline in -# the documentation of the scope in which they are defined (i.e. file, -# namespace, or group documentation), provided this scope is documented. If set -# to NO, structs, classes, and unions are shown on a separate page (for HTML and -# Man pages) or section (for LaTeX and RTF). -# The default value is: NO. - -INLINE_SIMPLE_STRUCTS = NO - -# When TYPEDEF_HIDES_STRUCT tag is enabled, a typedef of a struct, union, or -# enum is documented as struct, union, or enum with the name of the typedef. So -# typedef struct TypeS {} TypeT, will appear in the documentation as a struct -# with name TypeT. When disabled the typedef will appear as a member of a file, -# namespace, or class. And the struct will be named TypeS. This can typically be -# useful for C code in case the coding convention dictates that all compound -# types are typedef'ed and only the typedef is referenced, never the tag name. -# The default value is: NO. - -TYPEDEF_HIDES_STRUCT = NO - -# The size of the symbol lookup cache can be set using LOOKUP_CACHE_SIZE. This -# cache is used to resolve symbols given their name and scope. Since this can be -# an expensive process and often the same symbol appears multiple times in the -# code, doxygen keeps a cache of pre-resolved symbols. If the cache is too small -# doxygen will become slower. If the cache is too large, memory is wasted. The -# cache size is given by this formula: 2^(16+LOOKUP_CACHE_SIZE). The valid range -# is 0..9, the default is 0, corresponding to a cache size of 2^16=65536 -# symbols. At the end of a run doxygen will report the cache usage and suggest -# the optimal cache size from a speed point of view. -# Minimum value: 0, maximum value: 9, default value: 0. - -LOOKUP_CACHE_SIZE = 0 - -#--------------------------------------------------------------------------- -# Build related configuration options -#--------------------------------------------------------------------------- - -# If the EXTRACT_ALL tag is set to YES doxygen will assume all entities in -# documentation are documented, even if no documentation was available. Private -# class members and static file members will be hidden unless the -# EXTRACT_PRIVATE respectively EXTRACT_STATIC tags are set to YES. -# Note: This will also disable the warnings about undocumented members that are -# normally produced when WARNINGS is set to YES. -# The default value is: NO. - -EXTRACT_ALL = NO - -# If the EXTRACT_PRIVATE tag is set to YES all private members of a class will -# be included in the documentation. -# The default value is: NO. - -EXTRACT_PRIVATE = NO - -# If the EXTRACT_PACKAGE tag is set to YES all members with package or internal -# scope will be included in the documentation. -# The default value is: NO. - -EXTRACT_PACKAGE = NO - -# If the EXTRACT_STATIC tag is set to YES all static members of a file will be -# included in the documentation. -# The default value is: NO. - -EXTRACT_STATIC = NO - -# If the EXTRACT_LOCAL_CLASSES tag is set to YES classes (and structs) defined -# locally in source files will be included in the documentation. If set to NO -# only classes defined in header files are included. Does not have any effect -# for Java sources. -# The default value is: YES. - -EXTRACT_LOCAL_CLASSES = NO - -# This flag is only useful for Objective-C code. When set to YES local methods, -# which are defined in the implementation section but not in the interface are -# included in the documentation. If set to NO only methods in the interface are -# included. -# The default value is: NO. - -EXTRACT_LOCAL_METHODS = NO - -# If this flag is set to YES, the members of anonymous namespaces will be -# extracted and appear in the documentation as a namespace called -# 'anonymous_namespace{file}', where file will be replaced with the base name of -# the file that contains the anonymous namespace. By default anonymous namespace -# are hidden. -# The default value is: NO. - -EXTRACT_ANON_NSPACES = NO - -# If the HIDE_UNDOC_MEMBERS tag is set to YES, doxygen will hide all -# undocumented members inside documented classes or files. If set to NO these -# members will be included in the various overviews, but no documentation -# section is generated. This option has no effect if EXTRACT_ALL is enabled. -# The default value is: NO. - -HIDE_UNDOC_MEMBERS = NO - -# If the HIDE_UNDOC_CLASSES tag is set to YES, doxygen will hide all -# undocumented classes that are normally visible in the class hierarchy. If set -# to NO these classes will be included in the various overviews. This option has -# no effect if EXTRACT_ALL is enabled. -# The default value is: NO. - -HIDE_UNDOC_CLASSES = NO - -# If the HIDE_FRIEND_COMPOUNDS tag is set to YES, doxygen will hide all friend -# (class|struct|union) declarations. If set to NO these declarations will be -# included in the documentation. -# The default value is: NO. - -HIDE_FRIEND_COMPOUNDS = NO - -# If the HIDE_IN_BODY_DOCS tag is set to YES, doxygen will hide any -# documentation blocks found inside the body of a function. If set to NO these -# blocks will be appended to the function's detailed documentation block. -# The default value is: NO. - -HIDE_IN_BODY_DOCS = NO - -# The INTERNAL_DOCS tag determines if documentation that is typed after a -# \internal command is included. If the tag is set to NO then the documentation -# will be excluded. Set it to YES to include the internal documentation. -# The default value is: NO. - -INTERNAL_DOCS = NO - -# If the CASE_SENSE_NAMES tag is set to NO then doxygen will only generate file -# names in lower-case letters. If set to YES upper-case letters are also -# allowed. This is useful if you have classes or files whose names only differ -# in case and if your file system supports case sensitive file names. Windows -# and Mac users are advised to set this option to NO. -# The default value is: system dependent. - -CASE_SENSE_NAMES = YES - -# If the HIDE_SCOPE_NAMES tag is set to NO then doxygen will show members with -# their full class and namespace scopes in the documentation. If set to YES the -# scope will be hidden. -# The default value is: NO. - -HIDE_SCOPE_NAMES = NO - -# If the SHOW_INCLUDE_FILES tag is set to YES then doxygen will put a list of -# the files that are included by a file in the documentation of that file. -# The default value is: YES. - -SHOW_INCLUDE_FILES = YES - -# If the FORCE_LOCAL_INCLUDES tag is set to YES then doxygen will list include -# files with double quotes in the documentation rather than with sharp brackets. -# The default value is: NO. - -FORCE_LOCAL_INCLUDES = NO - -# If the INLINE_INFO tag is set to YES then a tag [inline] is inserted in the -# documentation for inline members. -# The default value is: YES. - -INLINE_INFO = YES - -# If the SORT_MEMBER_DOCS tag is set to YES then doxygen will sort the -# (detailed) documentation of file and class members alphabetically by member -# name. If set to NO the members will appear in declaration order. -# The default value is: YES. - -SORT_MEMBER_DOCS = YES - -# If the SORT_BRIEF_DOCS tag is set to YES then doxygen will sort the brief -# descriptions of file, namespace and class members alphabetically by member -# name. If set to NO the members will appear in declaration order. -# The default value is: NO. - -SORT_BRIEF_DOCS = NO - -# If the SORT_MEMBERS_CTORS_1ST tag is set to YES then doxygen will sort the -# (brief and detailed) documentation of class members so that constructors and -# destructors are listed first. If set to NO the constructors will appear in the -# respective orders defined by SORT_BRIEF_DOCS and SORT_MEMBER_DOCS. -# Note: If SORT_BRIEF_DOCS is set to NO this option is ignored for sorting brief -# member documentation. -# Note: If SORT_MEMBER_DOCS is set to NO this option is ignored for sorting -# detailed member documentation. -# The default value is: NO. - -SORT_MEMBERS_CTORS_1ST = NO - -# If the SORT_GROUP_NAMES tag is set to YES then doxygen will sort the hierarchy -# of group names into alphabetical order. If set to NO the group names will -# appear in their defined order. -# The default value is: NO. - -SORT_GROUP_NAMES = NO - -# If the SORT_BY_SCOPE_NAME tag is set to YES, the class list will be sorted by -# fully-qualified names, including namespaces. If set to NO, the class list will -# be sorted only by class name, not including the namespace part. -# Note: This option is not very useful if HIDE_SCOPE_NAMES is set to YES. -# Note: This option applies only to the class list, not to the alphabetical -# list. -# The default value is: NO. - -SORT_BY_SCOPE_NAME = NO - -# If the STRICT_PROTO_MATCHING option is enabled and doxygen fails to do proper -# type resolution of all parameters of a function it will reject a match between -# the prototype and the implementation of a member function even if there is -# only one candidate or it is obvious which candidate to choose by doing a -# simple string match. By disabling STRICT_PROTO_MATCHING doxygen will still -# accept a match between prototype and implementation in such cases. -# The default value is: NO. - -STRICT_PROTO_MATCHING = NO - -# The GENERATE_TODOLIST tag can be used to enable ( YES) or disable ( NO) the -# todo list. This list is created by putting \todo commands in the -# documentation. -# The default value is: YES. - -GENERATE_TODOLIST = YES - -# The GENERATE_TESTLIST tag can be used to enable ( YES) or disable ( NO) the -# test list. This list is created by putting \test commands in the -# documentation. -# The default value is: YES. - -GENERATE_TESTLIST = YES - -# The GENERATE_BUGLIST tag can be used to enable ( YES) or disable ( NO) the bug -# list. This list is created by putting \bug commands in the documentation. -# The default value is: YES. - -GENERATE_BUGLIST = YES - -# The GENERATE_DEPRECATEDLIST tag can be used to enable ( YES) or disable ( NO) -# the deprecated list. This list is created by putting \deprecated commands in -# the documentation. -# The default value is: YES. - -GENERATE_DEPRECATEDLIST= YES - -# The ENABLED_SECTIONS tag can be used to enable conditional documentation -# sections, marked by \if ... \endif and \cond -# ... \endcond blocks. - -ENABLED_SECTIONS = - -# The MAX_INITIALIZER_LINES tag determines the maximum number of lines that the -# initial value of a variable or macro / define can have for it to appear in the -# documentation. If the initializer consists of more lines than specified here -# it will be hidden. Use a value of 0 to hide initializers completely. The -# appearance of the value of individual variables and macros / defines can be -# controlled using \showinitializer or \hideinitializer command in the -# documentation regardless of this setting. -# Minimum value: 0, maximum value: 10000, default value: 30. - -MAX_INITIALIZER_LINES = 30 - -# Set the SHOW_USED_FILES tag to NO to disable the list of files generated at -# the bottom of the documentation of classes and structs. If set to YES the list -# will mention the files that were used to generate the documentation. -# The default value is: YES. - -SHOW_USED_FILES = YES - -# Set the SHOW_FILES tag to NO to disable the generation of the Files page. This -# will remove the Files entry from the Quick Index and from the Folder Tree View -# (if specified). -# The default value is: YES. - -SHOW_FILES = YES - -# Set the SHOW_NAMESPACES tag to NO to disable the generation of the Namespaces -# page. This will remove the Namespaces entry from the Quick Index and from the -# Folder Tree View (if specified). -# The default value is: YES. - -SHOW_NAMESPACES = YES - -# The FILE_VERSION_FILTER tag can be used to specify a program or script that -# doxygen should invoke to get the current version for each file (typically from -# the version control system). Doxygen will invoke the program by executing (via -# popen()) the command command input-file, where command is the value of the -# FILE_VERSION_FILTER tag, and input-file is the name of an input file provided -# by doxygen. Whatever the program writes to standard output is used as the file -# version. For an example see the documentation. - -FILE_VERSION_FILTER = - -# The LAYOUT_FILE tag can be used to specify a layout file which will be parsed -# by doxygen. The layout file controls the global structure of the generated -# output files in an output format independent way. To create the layout file -# that represents doxygen's defaults, run doxygen with the -l option. You can -# optionally specify a file name after the option, if omitted DoxygenLayout.xml -# will be used as the name of the layout file. -# -# Note that if you run doxygen from a directory containing a file called -# DoxygenLayout.xml, doxygen will parse it automatically even if the LAYOUT_FILE -# tag is left empty. - -LAYOUT_FILE = - -# The CITE_BIB_FILES tag can be used to specify one or more bib files containing -# the reference definitions. This must be a list of .bib files. The .bib -# extension is automatically appended if omitted. This requires the bibtex tool -# to be installed. See also http://en.wikipedia.org/wiki/BibTeX for more info. -# For LaTeX the style of the bibliography can be controlled using -# LATEX_BIB_STYLE. To use this feature you need bibtex and perl available in the -# search path. Do not use file names with spaces, bibtex cannot handle them. See -# also \cite for info how to create references. - -CITE_BIB_FILES = - -#--------------------------------------------------------------------------- -# Configuration options related to warning and progress messages -#--------------------------------------------------------------------------- - -# The QUIET tag can be used to turn on/off the messages that are generated to -# standard output by doxygen. If QUIET is set to YES this implies that the -# messages are off. -# The default value is: NO. - -QUIET = NO - -# The WARNINGS tag can be used to turn on/off the warning messages that are -# generated to standard error ( stderr) by doxygen. If WARNINGS is set to YES -# this implies that the warnings are on. -# -# Tip: Turn warnings on while writing the documentation. -# The default value is: YES. - -WARNINGS = YES - -# If the WARN_IF_UNDOCUMENTED tag is set to YES, then doxygen will generate -# warnings for undocumented members. If EXTRACT_ALL is set to YES then this flag -# will automatically be disabled. -# The default value is: YES. - -WARN_IF_UNDOCUMENTED = YES - -# If the WARN_IF_DOC_ERROR tag is set to YES, doxygen will generate warnings for -# potential errors in the documentation, such as not documenting some parameters -# in a documented function, or documenting parameters that don't exist or using -# markup commands wrongly. -# The default value is: YES. - -WARN_IF_DOC_ERROR = YES - -# This WARN_NO_PARAMDOC option can be enabled to get warnings for functions that -# are documented, but have no documentation for their parameters or return -# value. If set to NO doxygen will only warn about wrong or incomplete parameter -# documentation, but not about the absence of documentation. -# The default value is: NO. - -WARN_NO_PARAMDOC = NO - -# The WARN_FORMAT tag determines the format of the warning messages that doxygen -# can produce. The string should contain the $file, $line, and $text tags, which -# will be replaced by the file and line number from which the warning originated -# and the warning text. Optionally the format may contain $version, which will -# be replaced by the version of the file (if it could be obtained via -# FILE_VERSION_FILTER) -# The default value is: $file:$line: $text. - -WARN_FORMAT = "$file:$line: $text" - -# The WARN_LOGFILE tag can be used to specify a file to which warning and error -# messages should be written. If left blank the output is written to standard -# error (stderr). - -WARN_LOGFILE = - -#--------------------------------------------------------------------------- -# Configuration options related to the input files -#--------------------------------------------------------------------------- - -# The INPUT tag is used to specify the files and/or directories that contain -# documented source files. You may enter file names like myfile.cpp or -# directories like /usr/src/myproject. Separate the files or directories with -# spaces. -# Note: If this tag is empty the current directory is searched. - -INPUT = ./faiss - -# This tag can be used to specify the character encoding of the source files -# that doxygen parses. Internally doxygen uses the UTF-8 encoding. Doxygen uses -# libiconv (or the iconv built into libc) for the transcoding. See the libiconv -# documentation (see: http://www.gnu.org/software/libiconv) for the list of -# possible encodings. -# The default value is: UTF-8. - -INPUT_ENCODING = UTF-8 - -# If the value of the INPUT tag contains directories, you can use the -# FILE_PATTERNS tag to specify one or more wildcard patterns (like *.cpp and -# *.h) to filter out the source-files in the directories. If left blank the -# following patterns are tested:*.c, *.cc, *.cxx, *.cpp, *.c++, *.java, *.ii, -# *.ixx, *.ipp, *.i++, *.inl, *.idl, *.ddl, *.odl, *.h, *.hh, *.hxx, *.hpp, -# *.h++, *.cs, *.d, *.php, *.php4, *.php5, *.phtml, *.inc, *.m, *.markdown, -# *.md, *.mm, *.dox, *.py, *.f90, *.f, *.for, *.tcl, *.vhd, *.vhdl, *.ucf, -# *.qsf, *.as and *.js. - -FILE_PATTERNS = *.h *.cuh - -# The RECURSIVE tag can be used to specify whether or not subdirectories should -# be searched for input files as well. -# The default value is: NO. - -RECURSIVE = YES - -# The EXCLUDE tag can be used to specify files and/or directories that should be -# excluded from the INPUT source files. This way you can easily exclude a -# subdirectory from a directory tree whose root is specified with the INPUT tag. -# -# Note that relative paths are relative to the directory from which doxygen is -# run. - -EXCLUDE = gpu/test - -# The EXCLUDE_SYMLINKS tag can be used to select whether or not files or -# directories that are symbolic links (a Unix file system feature) are excluded -# from the input. -# The default value is: NO. - -EXCLUDE_SYMLINKS = NO - -# If the value of the INPUT tag contains directories, you can use the -# EXCLUDE_PATTERNS tag to specify one or more wildcard patterns to exclude -# certain files from those directories. -# -# Note that the wildcards are matched against the file with absolute path, so to -# exclude all test directories for example use the pattern */test/* - -EXCLUDE_PATTERNS = - -# The EXCLUDE_SYMBOLS tag can be used to specify one or more symbol names -# (namespaces, classes, functions, etc.) that should be excluded from the -# output. The symbol name can be a fully qualified name, a word, or if the -# wildcard * is used, a substring. Examples: ANamespace, AClass, -# AClass::ANamespace, ANamespace::*Test -# -# Note that the wildcards are matched against the file with absolute path, so to -# exclude all test directories use the pattern */test/* - -EXCLUDE_SYMBOLS = - -# The EXAMPLE_PATH tag can be used to specify one or more files or directories -# that contain example code fragments that are included (see the \include -# command). - -EXAMPLE_PATH = - -# If the value of the EXAMPLE_PATH tag contains directories, you can use the -# EXAMPLE_PATTERNS tag to specify one or more wildcard pattern (like *.cpp and -# *.h) to filter out the source-files in the directories. If left blank all -# files are included. - -EXAMPLE_PATTERNS = - -# If the EXAMPLE_RECURSIVE tag is set to YES then subdirectories will be -# searched for input files to be used with the \include or \dontinclude commands -# irrespective of the value of the RECURSIVE tag. -# The default value is: NO. - -EXAMPLE_RECURSIVE = NO - -# The IMAGE_PATH tag can be used to specify one or more files or directories -# that contain images that are to be included in the documentation (see the -# \image command). - -IMAGE_PATH = - -# The INPUT_FILTER tag can be used to specify a program that doxygen should -# invoke to filter for each input file. Doxygen will invoke the filter program -# by executing (via popen()) the command: -# -# -# -# where is the value of the INPUT_FILTER tag, and is the -# name of an input file. Doxygen will then use the output that the filter -# program writes to standard output. If FILTER_PATTERNS is specified, this tag -# will be ignored. -# -# Note that the filter must not add or remove lines; it is applied before the -# code is scanned, but not when the output code is generated. If lines are added -# or removed, the anchors will not be placed correctly. - -INPUT_FILTER = - -# The FILTER_PATTERNS tag can be used to specify filters on a per file pattern -# basis. Doxygen will compare the file name with each pattern and apply the -# filter if there is a match. The filters are a list of the form: pattern=filter -# (like *.cpp=my_cpp_filter). See INPUT_FILTER for further information on how -# filters are used. If the FILTER_PATTERNS tag is empty or if none of the -# patterns match the file name, INPUT_FILTER is applied. - -FILTER_PATTERNS = - -# If the FILTER_SOURCE_FILES tag is set to YES, the input filter (if set using -# INPUT_FILTER ) will also be used to filter the input files that are used for -# producing the source files to browse (i.e. when SOURCE_BROWSER is set to YES). -# The default value is: NO. - -FILTER_SOURCE_FILES = NO - -# The FILTER_SOURCE_PATTERNS tag can be used to specify source filters per file -# pattern. A pattern will override the setting for FILTER_PATTERN (if any) and -# it is also possible to disable source filtering for a specific pattern using -# *.ext= (so without naming a filter). -# This tag requires that the tag FILTER_SOURCE_FILES is set to YES. - -FILTER_SOURCE_PATTERNS = - -# If the USE_MDFILE_AS_MAINPAGE tag refers to the name of a markdown file that -# is part of the input, its contents will be placed on the main page -# (index.html). This can be useful if you have a project on for instance GitHub -# and want to reuse the introduction page also for the doxygen output. - -USE_MDFILE_AS_MAINPAGE = - -#--------------------------------------------------------------------------- -# Configuration options related to source browsing -#--------------------------------------------------------------------------- - -# If the SOURCE_BROWSER tag is set to YES then a list of source files will be -# generated. Documented entities will be cross-referenced with these sources. -# -# Note: To get rid of all source code in the generated output, make sure that -# also VERBATIM_HEADERS is set to NO. -# The default value is: NO. - -SOURCE_BROWSER = YES - -# Setting the INLINE_SOURCES tag to YES will include the body of functions, -# classes and enums directly into the documentation. -# The default value is: NO. - -INLINE_SOURCES = NO - -# Setting the STRIP_CODE_COMMENTS tag to YES will instruct doxygen to hide any -# special comment blocks from generated source code fragments. Normal C, C++ and -# Fortran comments will always remain visible. -# The default value is: YES. - -STRIP_CODE_COMMENTS = NO - -# If the REFERENCED_BY_RELATION tag is set to YES then for each documented -# function all documented functions referencing it will be listed. -# The default value is: NO. - -REFERENCED_BY_RELATION = NO - -# If the REFERENCES_RELATION tag is set to YES then for each documented function -# all documented entities called/used by that function will be listed. -# The default value is: NO. - -REFERENCES_RELATION = NO - -# If the REFERENCES_LINK_SOURCE tag is set to YES and SOURCE_BROWSER tag is set -# to YES, then the hyperlinks from functions in REFERENCES_RELATION and -# REFERENCED_BY_RELATION lists will link to the source code. Otherwise they will -# link to the documentation. -# The default value is: YES. - -REFERENCES_LINK_SOURCE = YES - -# If SOURCE_TOOLTIPS is enabled (the default) then hovering a hyperlink in the -# source code will show a tooltip with additional information such as prototype, -# brief description and links to the definition and documentation. Since this -# will make the HTML file larger and loading of large files a bit slower, you -# can opt to disable this feature. -# The default value is: YES. -# This tag requires that the tag SOURCE_BROWSER is set to YES. - -SOURCE_TOOLTIPS = YES - -# If the USE_HTAGS tag is set to YES then the references to source code will -# point to the HTML generated by the htags(1) tool instead of doxygen built-in -# source browser. The htags tool is part of GNU's global source tagging system -# (see http://www.gnu.org/software/global/global.html). You will need version -# 4.8.6 or higher. -# -# To use it do the following: -# - Install the latest version of global -# - Enable SOURCE_BROWSER and USE_HTAGS in the config file -# - Make sure the INPUT points to the root of the source tree -# - Run doxygen as normal -# -# Doxygen will invoke htags (and that will in turn invoke gtags), so these -# tools must be available from the command line (i.e. in the search path). -# -# The result: instead of the source browser generated by doxygen, the links to -# source code will now point to the output of htags. -# The default value is: NO. -# This tag requires that the tag SOURCE_BROWSER is set to YES. - -USE_HTAGS = NO - -# If the VERBATIM_HEADERS tag is set the YES then doxygen will generate a -# verbatim copy of the header file for each class for which an include is -# specified. Set to NO to disable this. -# See also: Section \class. -# The default value is: YES. - -VERBATIM_HEADERS = YES - -#--------------------------------------------------------------------------- -# Configuration options related to the alphabetical class index -#--------------------------------------------------------------------------- - -# If the ALPHABETICAL_INDEX tag is set to YES, an alphabetical index of all -# compounds will be generated. Enable this if the project contains a lot of -# classes, structs, unions or interfaces. -# The default value is: YES. - -ALPHABETICAL_INDEX = YES - -# The COLS_IN_ALPHA_INDEX tag can be used to specify the number of columns in -# which the alphabetical index list will be split. -# Minimum value: 1, maximum value: 20, default value: 5. -# This tag requires that the tag ALPHABETICAL_INDEX is set to YES. - -COLS_IN_ALPHA_INDEX = 5 - -# In case all classes in a project start with a common prefix, all classes will -# be put under the same header in the alphabetical index. The IGNORE_PREFIX tag -# can be used to specify a prefix (or a list of prefixes) that should be ignored -# while generating the index headers. -# This tag requires that the tag ALPHABETICAL_INDEX is set to YES. - -IGNORE_PREFIX = - -#--------------------------------------------------------------------------- -# Configuration options related to the HTML output -#--------------------------------------------------------------------------- - -# If the GENERATE_HTML tag is set to YES doxygen will generate HTML output -# The default value is: YES. - -GENERATE_HTML = NO - -# The HTML_OUTPUT tag is used to specify where the HTML docs will be put. If a -# relative path is entered the value of OUTPUT_DIRECTORY will be put in front of -# it. -# The default directory is: html. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_OUTPUT = html - -# The HTML_FILE_EXTENSION tag can be used to specify the file extension for each -# generated HTML page (for example: .htm, .php, .asp). -# The default value is: .html. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_FILE_EXTENSION = .html - -# The HTML_HEADER tag can be used to specify a user-defined HTML header file for -# each generated HTML page. If the tag is left blank doxygen will generate a -# standard header. -# -# To get valid HTML the header file that includes any scripts and style sheets -# that doxygen needs, which is dependent on the configuration options used (e.g. -# the setting GENERATE_TREEVIEW). It is highly recommended to start with a -# default header using -# doxygen -w html new_header.html new_footer.html new_stylesheet.css -# YourConfigFile -# and then modify the file new_header.html. See also section "Doxygen usage" -# for information on how to generate the default header that doxygen normally -# uses. -# Note: The header is subject to change so you typically have to regenerate the -# default header when upgrading to a newer version of doxygen. For a description -# of the possible markers and block names see the documentation. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_HEADER = - -# The HTML_FOOTER tag can be used to specify a user-defined HTML footer for each -# generated HTML page. If the tag is left blank doxygen will generate a standard -# footer. See HTML_HEADER for more information on how to generate a default -# footer and what special commands can be used inside the footer. See also -# section "Doxygen usage" for information on how to generate the default footer -# that doxygen normally uses. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_FOOTER = - -# The HTML_STYLESHEET tag can be used to specify a user-defined cascading style -# sheet that is used by each HTML page. It can be used to fine-tune the look of -# the HTML output. If left blank doxygen will generate a default style sheet. -# See also section "Doxygen usage" for information on how to generate the style -# sheet that doxygen normally uses. -# Note: It is recommended to use HTML_EXTRA_STYLESHEET instead of this tag, as -# it is more robust and this tag (HTML_STYLESHEET) will in the future become -# obsolete. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_STYLESHEET = - -# The HTML_EXTRA_STYLESHEET tag can be used to specify an additional user- -# defined cascading style sheet that is included after the standard style sheets -# created by doxygen. Using this option one can overrule certain style aspects. -# This is preferred over using HTML_STYLESHEET since it does not replace the -# standard style sheet and is therefor more robust against future updates. -# Doxygen will copy the style sheet file to the output directory. For an example -# see the documentation. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_EXTRA_STYLESHEET = - -# The HTML_EXTRA_FILES tag can be used to specify one or more extra images or -# other source files which should be copied to the HTML output directory. Note -# that these files will be copied to the base HTML output directory. Use the -# $relpath^ marker in the HTML_HEADER and/or HTML_FOOTER files to load these -# files. In the HTML_STYLESHEET file, use the file name only. Also note that the -# files will be copied as-is; there are no commands or markers available. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_EXTRA_FILES = - -# The HTML_COLORSTYLE_HUE tag controls the color of the HTML output. Doxygen -# will adjust the colors in the stylesheet and background images according to -# this color. Hue is specified as an angle on a colorwheel, see -# http://en.wikipedia.org/wiki/Hue for more information. For instance the value -# 0 represents red, 60 is yellow, 120 is green, 180 is cyan, 240 is blue, 300 -# purple, and 360 is red again. -# Minimum value: 0, maximum value: 359, default value: 220. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_COLORSTYLE_HUE = 220 - -# The HTML_COLORSTYLE_SAT tag controls the purity (or saturation) of the colors -# in the HTML output. For a value of 0 the output will use grayscales only. A -# value of 255 will produce the most vivid colors. -# Minimum value: 0, maximum value: 255, default value: 100. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_COLORSTYLE_SAT = 100 - -# The HTML_COLORSTYLE_GAMMA tag controls the gamma correction applied to the -# luminance component of the colors in the HTML output. Values below 100 -# gradually make the output lighter, whereas values above 100 make the output -# darker. The value divided by 100 is the actual gamma applied, so 80 represents -# a gamma of 0.8, The value 220 represents a gamma of 2.2, and 100 does not -# change the gamma. -# Minimum value: 40, maximum value: 240, default value: 80. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_COLORSTYLE_GAMMA = 80 - -# If the HTML_TIMESTAMP tag is set to YES then the footer of each generated HTML -# page will contain the date and time when the page was generated. Setting this -# to NO can help when comparing the output of multiple runs. -# The default value is: YES. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_TIMESTAMP = NO - -# If the HTML_DYNAMIC_SECTIONS tag is set to YES then the generated HTML -# documentation will contain sections that can be hidden and shown after the -# page has loaded. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_DYNAMIC_SECTIONS = NO - -# With HTML_INDEX_NUM_ENTRIES one can control the preferred number of entries -# shown in the various tree structured indices initially; the user can expand -# and collapse entries dynamically later on. Doxygen will expand the tree to -# such a level that at most the specified number of entries are visible (unless -# a fully collapsed tree already exceeds this amount). So setting the number of -# entries 1 will produce a full collapsed tree by default. 0 is a special value -# representing an infinite number of entries and will result in a full expanded -# tree by default. -# Minimum value: 0, maximum value: 9999, default value: 100. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_INDEX_NUM_ENTRIES = 100 - -# If the GENERATE_DOCSET tag is set to YES, additional index files will be -# generated that can be used as input for Apple's Xcode 3 integrated development -# environment (see: http://developer.apple.com/tools/xcode/), introduced with -# OSX 10.5 (Leopard). To create a documentation set, doxygen will generate a -# Makefile in the HTML output directory. Running make will produce the docset in -# that directory and running make install will install the docset in -# ~/Library/Developer/Shared/Documentation/DocSets so that Xcode will find it at -# startup. See http://developer.apple.com/tools/creatingdocsetswithdoxygen.html -# for more information. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -GENERATE_DOCSET = NO - -# This tag determines the name of the docset feed. A documentation feed provides -# an umbrella under which multiple documentation sets from a single provider -# (such as a company or product suite) can be grouped. -# The default value is: Doxygen generated docs. -# This tag requires that the tag GENERATE_DOCSET is set to YES. - -DOCSET_FEEDNAME = "Doxygen generated docs" - -# This tag specifies a string that should uniquely identify the documentation -# set bundle. This should be a reverse domain-name style string, e.g. -# com.mycompany.MyDocSet. Doxygen will append .docset to the name. -# The default value is: org.doxygen.Project. -# This tag requires that the tag GENERATE_DOCSET is set to YES. - -DOCSET_BUNDLE_ID = org.doxygen.Project - -# The DOCSET_PUBLISHER_ID tag specifies a string that should uniquely identify -# the documentation publisher. This should be a reverse domain-name style -# string, e.g. com.mycompany.MyDocSet.documentation. -# The default value is: org.doxygen.Publisher. -# This tag requires that the tag GENERATE_DOCSET is set to YES. - -DOCSET_PUBLISHER_ID = org.doxygen.Publisher - -# The DOCSET_PUBLISHER_NAME tag identifies the documentation publisher. -# The default value is: Publisher. -# This tag requires that the tag GENERATE_DOCSET is set to YES. - -DOCSET_PUBLISHER_NAME = Publisher - -# If the GENERATE_HTMLHELP tag is set to YES then doxygen generates three -# additional HTML index files: index.hhp, index.hhc, and index.hhk. The -# index.hhp is a project file that can be read by Microsoft's HTML Help Workshop -# (see: http://www.microsoft.com/en-us/download/details.aspx?id=21138) on -# Windows. -# -# The HTML Help Workshop contains a compiler that can convert all HTML output -# generated by doxygen into a single compiled HTML file (.chm). Compiled HTML -# files are now used as the Windows 98 help format, and will replace the old -# Windows help format (.hlp) on all Windows platforms in the future. Compressed -# HTML files also contain an index, a table of contents, and you can search for -# words in the documentation. The HTML workshop also contains a viewer for -# compressed HTML files. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -GENERATE_HTMLHELP = NO - -# The CHM_FILE tag can be used to specify the file name of the resulting .chm -# file. You can add a path in front of the file if the result should not be -# written to the html output directory. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -CHM_FILE = - -# The HHC_LOCATION tag can be used to specify the location (absolute path -# including file name) of the HTML help compiler ( hhc.exe). If non-empty -# doxygen will try to run the HTML help compiler on the generated index.hhp. -# The file has to be specified with full path. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -HHC_LOCATION = - -# The GENERATE_CHI flag controls if a separate .chi index file is generated ( -# YES) or that it should be included in the master .chm file ( NO). -# The default value is: NO. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -GENERATE_CHI = NO - -# The CHM_INDEX_ENCODING is used to encode HtmlHelp index ( hhk), content ( hhc) -# and project file content. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -CHM_INDEX_ENCODING = - -# The BINARY_TOC flag controls whether a binary table of contents is generated ( -# YES) or a normal table of contents ( NO) in the .chm file. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -BINARY_TOC = NO - -# The TOC_EXPAND flag can be set to YES to add extra items for group members to -# the table of contents of the HTML help documentation and to the tree view. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTMLHELP is set to YES. - -TOC_EXPAND = NO - -# If the GENERATE_QHP tag is set to YES and both QHP_NAMESPACE and -# QHP_VIRTUAL_FOLDER are set, an additional index file will be generated that -# can be used as input for Qt's qhelpgenerator to generate a Qt Compressed Help -# (.qch) of the generated HTML documentation. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -GENERATE_QHP = NO - -# If the QHG_LOCATION tag is specified, the QCH_FILE tag can be used to specify -# the file name of the resulting .qch file. The path specified is relative to -# the HTML output folder. -# This tag requires that the tag GENERATE_QHP is set to YES. - -QCH_FILE = - -# The QHP_NAMESPACE tag specifies the namespace to use when generating Qt Help -# Project output. For more information please see Qt Help Project / Namespace -# (see: http://qt-project.org/doc/qt-4.8/qthelpproject.html#namespace). -# The default value is: org.doxygen.Project. -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHP_NAMESPACE = org.doxygen.Project - -# The QHP_VIRTUAL_FOLDER tag specifies the namespace to use when generating Qt -# Help Project output. For more information please see Qt Help Project / Virtual -# Folders (see: http://qt-project.org/doc/qt-4.8/qthelpproject.html#virtual- -# folders). -# The default value is: doc. -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHP_VIRTUAL_FOLDER = doc - -# If the QHP_CUST_FILTER_NAME tag is set, it specifies the name of a custom -# filter to add. For more information please see Qt Help Project / Custom -# Filters (see: http://qt-project.org/doc/qt-4.8/qthelpproject.html#custom- -# filters). -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHP_CUST_FILTER_NAME = - -# The QHP_CUST_FILTER_ATTRS tag specifies the list of the attributes of the -# custom filter to add. For more information please see Qt Help Project / Custom -# Filters (see: http://qt-project.org/doc/qt-4.8/qthelpproject.html#custom- -# filters). -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHP_CUST_FILTER_ATTRS = - -# The QHP_SECT_FILTER_ATTRS tag specifies the list of the attributes this -# project's filter section matches. Qt Help Project / Filter Attributes (see: -# http://qt-project.org/doc/qt-4.8/qthelpproject.html#filter-attributes). -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHP_SECT_FILTER_ATTRS = - -# The QHG_LOCATION tag can be used to specify the location of Qt's -# qhelpgenerator. If non-empty doxygen will try to run qhelpgenerator on the -# generated .qhp file. -# This tag requires that the tag GENERATE_QHP is set to YES. - -QHG_LOCATION = - -# If the GENERATE_ECLIPSEHELP tag is set to YES, additional index files will be -# generated, together with the HTML files, they form an Eclipse help plugin. To -# install this plugin and make it available under the help contents menu in -# Eclipse, the contents of the directory containing the HTML and XML files needs -# to be copied into the plugins directory of eclipse. The name of the directory -# within the plugins directory should be the same as the ECLIPSE_DOC_ID value. -# After copying Eclipse needs to be restarted before the help appears. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -GENERATE_ECLIPSEHELP = NO - -# A unique identifier for the Eclipse help plugin. When installing the plugin -# the directory name containing the HTML and XML files should also have this -# name. Each documentation set should have its own identifier. -# The default value is: org.doxygen.Project. -# This tag requires that the tag GENERATE_ECLIPSEHELP is set to YES. - -ECLIPSE_DOC_ID = org.doxygen.Project - -# If you want full control over the layout of the generated HTML pages it might -# be necessary to disable the index and replace it with your own. The -# DISABLE_INDEX tag can be used to turn on/off the condensed index (tabs) at top -# of each HTML page. A value of NO enables the index and the value YES disables -# it. Since the tabs in the index contain the same information as the navigation -# tree, you can set this option to YES if you also set GENERATE_TREEVIEW to YES. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -DISABLE_INDEX = NO - -# The GENERATE_TREEVIEW tag is used to specify whether a tree-like index -# structure should be generated to display hierarchical information. If the tag -# value is set to YES, a side panel will be generated containing a tree-like -# index structure (just like the one that is generated for HTML Help). For this -# to work a browser that supports JavaScript, DHTML, CSS and frames is required -# (i.e. any modern browser). Windows users are probably better off using the -# HTML help feature. Via custom stylesheets (see HTML_EXTRA_STYLESHEET) one can -# further fine-tune the look of the index. As an example, the default style -# sheet generated by doxygen has an example that shows how to put an image at -# the root of the tree instead of the PROJECT_NAME. Since the tree basically has -# the same information as the tab index, you could consider setting -# DISABLE_INDEX to YES when enabling this option. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -GENERATE_TREEVIEW = NO - -# The ENUM_VALUES_PER_LINE tag can be used to set the number of enum values that -# doxygen will group on one line in the generated HTML documentation. -# -# Note that a value of 0 will completely suppress the enum values from appearing -# in the overview section. -# Minimum value: 0, maximum value: 20, default value: 4. -# This tag requires that the tag GENERATE_HTML is set to YES. - -ENUM_VALUES_PER_LINE = 4 - -# If the treeview is enabled (see GENERATE_TREEVIEW) then this tag can be used -# to set the initial width (in pixels) of the frame in which the tree is shown. -# Minimum value: 0, maximum value: 1500, default value: 250. -# This tag requires that the tag GENERATE_HTML is set to YES. - -TREEVIEW_WIDTH = 250 - -# When the EXT_LINKS_IN_WINDOW option is set to YES doxygen will open links to -# external symbols imported via tag files in a separate window. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -EXT_LINKS_IN_WINDOW = NO - -# Use this tag to change the font size of LaTeX formulas included as images in -# the HTML documentation. When you change the font size after a successful -# doxygen run you need to manually remove any form_*.png images from the HTML -# output directory to force them to be regenerated. -# Minimum value: 8, maximum value: 50, default value: 10. -# This tag requires that the tag GENERATE_HTML is set to YES. - -FORMULA_FONTSIZE = 10 - -# Use the FORMULA_TRANPARENT tag to determine whether or not the images -# generated for formulas are transparent PNGs. Transparent PNGs are not -# supported properly for IE 6.0, but are supported on all modern browsers. -# -# Note that when changing this option you need to delete any form_*.png files in -# the HTML output directory before the changes have effect. -# The default value is: YES. -# This tag requires that the tag GENERATE_HTML is set to YES. - -FORMULA_TRANSPARENT = YES - -# Enable the USE_MATHJAX option to render LaTeX formulas using MathJax (see -# http://www.mathjax.org) which uses client side Javascript for the rendering -# instead of using prerendered bitmaps. Use this if you do not have LaTeX -# installed or if you want to formulas look prettier in the HTML output. When -# enabled you may also need to install MathJax separately and configure the path -# to it using the MATHJAX_RELPATH option. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -USE_MATHJAX = NO - -# When MathJax is enabled you can set the default output format to be used for -# the MathJax output. See the MathJax site (see: -# http://docs.mathjax.org/en/latest/output.html) for more details. -# Possible values are: HTML-CSS (which is slower, but has the best -# compatibility), NativeMML (i.e. MathML) and SVG. -# The default value is: HTML-CSS. -# This tag requires that the tag USE_MATHJAX is set to YES. - -MATHJAX_FORMAT = HTML-CSS - -# When MathJax is enabled you need to specify the location relative to the HTML -# output directory using the MATHJAX_RELPATH option. The destination directory -# should contain the MathJax.js script. For instance, if the mathjax directory -# is located at the same level as the HTML output directory, then -# MATHJAX_RELPATH should be ../mathjax. The default value points to the MathJax -# Content Delivery Network so you can quickly see the result without installing -# MathJax. However, it is strongly recommended to install a local copy of -# MathJax from http://www.mathjax.org before deployment. -# The default value is: http://cdn.mathjax.org/mathjax/latest. -# This tag requires that the tag USE_MATHJAX is set to YES. - -MATHJAX_RELPATH = http://cdn.mathjax.org/mathjax/latest - -# The MATHJAX_EXTENSIONS tag can be used to specify one or more MathJax -# extension names that should be enabled during MathJax rendering. For example -# MATHJAX_EXTENSIONS = TeX/AMSmath TeX/AMSsymbols -# This tag requires that the tag USE_MATHJAX is set to YES. - -MATHJAX_EXTENSIONS = - -# The MATHJAX_CODEFILE tag can be used to specify a file with javascript pieces -# of code that will be used on startup of the MathJax code. See the MathJax site -# (see: http://docs.mathjax.org/en/latest/output.html) for more details. For an -# example see the documentation. -# This tag requires that the tag USE_MATHJAX is set to YES. - -MATHJAX_CODEFILE = - -# When the SEARCHENGINE tag is enabled doxygen will generate a search box for -# the HTML output. The underlying search engine uses javascript and DHTML and -# should work on any modern browser. Note that when using HTML help -# (GENERATE_HTMLHELP), Qt help (GENERATE_QHP), or docsets (GENERATE_DOCSET) -# there is already a search function so this one should typically be disabled. -# For large projects the javascript based search engine can be slow, then -# enabling SERVER_BASED_SEARCH may provide a better solution. It is possible to -# search using the keyboard; to jump to the search box use + S -# (what the is depends on the OS and browser, but it is typically -# , /