[
{
"path": ".gitattributes",
"content": "notebooks/* linguist-vendored\ninclude/* linguist-vendored\nbindings/* linguist-vendored\n*.h linguist-detectable=false\n*.cpp linguist-detectable=false\n"
},
{
"path": ".github/release-template.ejs",
"content": "<% var groupCommits = [\n{\n name: 'breaking',\n show: true,\n list: []\n}, {\n name: 'feat',\n show: true,\n list: []\n}, {\n name: 'perf',\n show: true,\n list: []\n}, {\n name: 'fix',\n show: true,\n list: []\n}, {\n name: 'refactor',\n show: true,\n list: []\n}, {\n name: 'docs',\n show: true,\n list: []\n}, {\n name: 'test',\n show: true,\n list: []\n}, {\n name: 'other',\n show: true,\n list: []\n}\n]\n\nvar all_titles = {};\nvar all_commiters = {};\nvar commitHref = \"https://github.com/jina-ai/docarray/commit/\"\n\ncommits.forEach(function (commit) {\n\n var result = (commit.title).match(/^(\\w*)(\\((.*)\\))?\\: (.*)$/);\n\n var type = result && result[1];\n var scope = result && result[3];\n var title = result && result[4];\n var committer = commit.authorName\n\n if (!(committer in all_commiters)) {\n all_commiters[committer] = 1\n }\n\n if (!(title in all_titles)) {\n all_titles[title] = 1\n if( title != null && (title.indexOf('๐ฅ')>-1 || title.indexOf(':boom:')>-1) ){\n groupCommits.find(item => item.name === 'breaking').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'fix' || type == 'fixed'){\n groupCommits.find(item => item.name === 'fix').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'perf' || type == 'performance'){\n groupCommits.find(item => item.name === 'perf').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'feat' || type == 'feature'){\n groupCommits.find(item => item.name === 'feat').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'refactor'){\n groupCommits.find(item => item.name === 'refactor').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'docs' || type == 'doc'){\n groupCommits.find(item => item.name === 'docs').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else if(type == 'test' || type == 'tests' || type == 'ci'){\n groupCommits.find(item => item.name === 'test').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n } else {\n groupCommits.find(item => item.name === 'other').list.push({\n type: type,\n scope: scope,\n title: title,\n commit: commit\n })\n }\n }\n\n\n});\n\n\nvar listCommits = function(list, key){\n\nlist.forEach(function (ct) {\n\n var type = ct.type;\n var scope = ct.scope;\n var title = '';\n var commit = ct.commit;\n\n if(type){\n if(key != 'other'){\n title = (scope? '__'+scope+'__: ':'') + ct.title;\n }else{\n title = '__' + type + (scope? '('+scope+')':'') + '__ : ' + ct.title;\n }\n }else{\n title = commit.title;\n }\n%> - <% if(typeof commitHref === 'undefined' || commitHref === '') { %>[```<%=commit.sha1.slice(0,8)%>```]<% } else { %>[[```<%=commit.sha1.slice(0,8)%>```](<%=commitHref%><%=commit.sha1%>)]<%}%> __-__ <%=title%> (*<%= commit.authorName %>*)\n<% })} %>\n\n๐ We'd like to thank all contributors for this new release! In particular,\n<% Object.keys(all_commiters).forEach(function (key) { %> <%= key %>, <% }) %> ๐\n\n<%\n for(var i of groupCommits){\n if(i.list.length == 0) continue;\n\nif (i.name === 'breaking' && i.show) { %>\n### ๐ฅ Breaking changes\n\n<%\t} else if (i.name === 'fix' && i.show) { %>\n### ๐ Bug fixes\n\n<%\t} else if( i.name === 'feat' && i.show) { %>\n### ๐ New Features\n\n<%\t} else if(i.name === 'perf' && i.show) { %>\n### โก Performance Improvements\n\n<%\t} else if(i.name === 'refactor' && i.show) { %>\n### ๐งผ Code Refactoring\n\n<%\t} else if(i.name === 'docs' && i.show) { %>\n### ๐ Documentation\n\n<%\t} else if(i.name === 'test' && i.show) { %>\n### ๐ Unit Test and CICD\n\n<%\t} else if (i.name === 'other' && i.show) { %>\n### ๐น Other Improvements\n\n<%\t}\n i.show && listCommits(i.list, i);\n} %>\n"
},
{
"path": ".github/requirements-test.txt",
"content": "pytest\npytest-custom_exit_code\n"
},
{
"path": ".github/workflows/cd.yml",
"content": "name: CD\n\non:\n push:\n branches:\n - main\n\njobs:\n prep-testbed:\n if: |\n !startsWith(github.event.head_commit.message, 'chore') &&\n !startsWith(github.event.head_commit.message, 'build: hotfix') &&\n !endsWith(github.event.head_commit.message, 'reformatted by jina-dev-bot')\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n - id: set-matrix\n run: |\n sudo apt-get install jq\n echo \"::set-output name=matrix::$(bash scripts/get-all-test-paths.sh)\"\n outputs:\n matrix: ${{ steps.set-matrix.outputs.matrix }}\n\n build-wheels:\n needs: [prep-testbed]\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest, windows-latest, macos-latest]\n cibw_arch: [\"auto64\"]\n python-version: [['3.7', \"cp37-*\"], ['3.8', \"cp38-*\"], ['3.9', \"cp39-*\"], ['3.10', \"cp310-*\"]]\n steps:\n - uses: actions/checkout@v3\n with:\n fetch-depth: 100\n - name: Set up Python ${{ matrix.python-version[0] }}\n uses: actions/setup-python@v4\n with:\n python-version: ${{ matrix.python-version[0] }}\n - name: Update version\n shell: bash\n run: |\n git fetch --depth=1 origin +refs/tags/*:refs/tags/*\n ./scripts/update-version.sh\n - name: Build sdist\n if: ${{ matrix.os == 'ubuntu-latest' && matrix.python-version[0] == '3.7' }}\n run: |\n pip install build\n python -m build --sdist\n - name: Build wheels\n uses: pypa/cibuildwheel@v2.10.2\n with:\n package-dir: ./\n env:\n CIBW_ENVIRONMENT: >\n STAN_BACKEND=\"${{ env.STAN_BACKEND }}\"\n CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014\n CIBW_BUILD: ${{ matrix.python-version[1] }}\n CIBW_SKIP: \"*musllinux*\"\n CIBW_ARCHS: ${{ matrix.cibw_arch }}\n # CIBW_ARCH_MACOS: x86_64 arm64\n CIBW_BUILD_FRONTEND: build\n - uses: actions/upload-artifact@v3\n with:\n path: |\n ./wheelhouse/*.whl\n ./dist/*.tar.gz\n\n core-test:\n needs: [ prep-testbed, build-wheels ]\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ ubuntu-latest, windows-latest, macos-latest ]\n cibw_arch: [ \"auto64\" ]\n python-version: [ [ '3.7', \"cp37-*\" ] ]\n test-path: ${{fromJson(needs.prep-testbed.outputs.matrix)}}\n env:\n JINA_HIDE_SURVEY: \"1\"\n steps:\n - uses: actions/checkout@v3\n with:\n fetch-depth: 100\n - name: Set up Python ${{ matrix.python-version[0] }}\n uses: actions/setup-python@v4\n with:\n python-version: ${{ matrix.python-version[0] }}\n - name: Prepare enviroment\n run: |\n python -m pip install --upgrade pip\n pip install jina\n pip install --pre docarray\n pip install pytest pytest-html pytest-cov pytest-mock pytest-repeat pytest-custom-exit-code pytest-timeout pytest-reraise\n - uses: actions/download-artifact@v3\n with:\n name: artifact\n - name: Install annlite linux\n if: ${{ matrix.os == 'ubuntu-latest' }}\n run: |\n pip install wheelhouse/*${{ matrix.python-version[1] }}**linux*.whl\n - name: Install annlite macos\n if: ${{ matrix.os == 'macos-latest' }}\n run: |\n pip install wheelhouse/*${{ matrix.python-version[1] }}**macos**x86_64*.whl\n - name: Install annlite win\n if: ${{ matrix.os == 'windows-latest'}}\n run: |\n pip install --pre --find-links=wheelhouse/ annlite\n\n - name: Test unix\n id: test_unix\n if: ${{ matrix.os == 'ubuntu-latest' || matrix.os == 'macos-latest' }}\n run: |\n cd ..\n mv annlite/tests/ ./\n pytest --suppress-no-test-exit-code --cov=annlite --cov-report=xml \\\n -v -s -m \"not gpu\" ${{ matrix.test-path }}\n echo \"::set-output name=codecov_flag::annlite\"\n timeout-minutes: 30\n - name: Test win\n id: test_win\n if: ${{ matrix.os == 'windows-latest'}}\n env:\n PYTHONIOENCODING: 'utf-8'\n run: |\n cd ..\n move annlite/tests/ ./\n cd tests/\n pytest -v -s -m \"not gpu\" -k \"test\"\n timeout-minutes: 30\n - name: Check codecov file\n id: check_files\n uses: andstor/file-existence-action@v1\n with:\n files: \"coverage.xml\"\n - name: Upload coverage from test to Codecov\n uses: codecov/codecov-action@v2\n if: steps.check_files.outputs.files_exists == 'true' && ${{ matrix.python-version[0] }} == '3.7' && ${{ matrix.matrix.os }} == 'ubuntu-latest'\n with:\n file: coverage.xml\n flags: ${{ steps.test.outputs.codecov_flag }}\n fail_ci_if_error: false\n token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos\n\n prerelease:\n needs: [core-test]\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n with:\n fetch-depth: 100\n - uses: actions/download-artifact@v3\n with:\n name: artifact\n - name: Pre-release (.devN)\n run: |\n git fetch --depth=1 origin +refs/tags/*:refs/tags/*\n pip install twine\n ./scripts/release.sh\n env:\n TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }}\n TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}\n JINA_SLACK_WEBHOOK: ${{ secrets.JINA_SLACK_WEBHOOK }}\n"
},
{
"path": ".github/workflows/ci.yml",
"content": "name: CI\n\non:\n pull_request:\n\njobs:\n commit-lint:\n runs-on: ubuntu-latest\n steps:\n - name: find the prev warning if exist\n uses: peter-evans/find-comment@v1\n id: fc\n with:\n issue-number: ${{ github.event.pull_request.number }}\n comment-author: \"github-actions[bot]\"\n body-includes: \"bad commit message\"\n - name: Delete comment if exist\n if: ${{ steps.fc.outputs.comment-id != 0 }}\n uses: actions/github-script@v3\n with:\n github-token: ${{ secrets.GITHUB_TOKEN }}\n script: |\n github.issues.deleteComment({\n owner: context.repo.owner,\n repo: context.repo.repo,\n comment_id: ${{ steps.fc.outputs.comment-id }},\n })\n - uses: actions/checkout@v2\n with:\n fetch-depth: 0\n - run: 'echo \"module.exports = {extends: [''@commitlint/config-conventional'']}\" > commitlint.config.js'\n - uses: wagoid/commitlint-github-action@v1\n env:\n GITHUB_TOKEN: \"${{ secrets.GITHUB_TOKEN }}\"\n - name: if lint failed\n if: ${{ failure() }}\n uses: peter-evans/create-or-update-comment@v1\n with:\n issue-number: ${{ github.event.pull_request.number }}\n body: |\n Thanks for your contribution :heart:\n :broken_heart: Unfortunately, this PR has one ore more **bad commit messages**, it can not be merged. To fix this problem, please refer to:\n - [Commit Message Guideline for the First Time Contributor](https://github.com/jina-ai/jina/issues/553)\n - [Contributing Guideline](https://github.com/jina-ai/jina/blob/master/CONTRIBUTING.md)\n\n Note, other CI tests will *not* *start* until the commit messages get fixed.\n\n This message will be deleted automatically when the commit messages get fixed.\n reaction-type: \"eyes\"\n\n# lint-flake-8:\n# runs-on: ubuntu-latest\n# steps:\n# - uses: actions/checkout@v2\n# - name: Set up Python 3.7\n# uses: actions/setup-python@v2\n# with:\n# python-version: 3.7\n# - name: Lint with flake8\n# run: |\n# pip install flake8\n# # stop the build if there are Python syntax errors or undefined names\n# flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude .git,__pycache__,docs/source/conf.py,old,build,dist,tests/,jina/resources/,bindings\n# # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide\n# flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude .git,__pycache__,docs/source/conf.py,old,build,dist,tests/,jina/resources/\n\n check-black:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n with:\n fetch-depth: 0\n - name: Set up Python 3.7\n uses: actions/setup-python@v2\n with:\n python-version: 3.7\n - id: file_changes\n uses: Ana06/get-changed-files@v1.2\n - name: check black\n run: ./scripts/black.sh\n env:\n CHANGED_FILES: ${{ steps.file_changes.outputs.added_modified }}\n\n prep-testbed:\n runs-on: ubuntu-latest\n needs: [commit-lint, check-black]\n steps:\n - uses: actions/checkout@v2\n - id: set-matrix\n run: |\n sudo apt-get install jq\n echo \"::set-output name=matrix::$(bash scripts/get-all-test-paths.sh)\"\n outputs:\n matrix: ${{ steps.set-matrix.outputs.matrix }}\n\n core-test:\n needs: [prep-testbed]\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest, windows-latest]\n cpython-version: [\"cp37-*\"]\n python-version: [3.7]\n cibw_arch: [\"auto64\"]\n test-path: ${{fromJson(needs.prep-testbed.outputs.matrix)}}\n env:\n JINA_HIDE_SURVEY: \"1\"\n steps:\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@v3\n with:\n python-version: ${{ matrix.python-version }}\n - uses: actions/checkout@v3\n - name: Prepare enviroment\n run: |\n python -m pip install --upgrade pip\n pip install jina\n pip install --pre docarray\n pip install pytest pytest-html pytest-cov pytest-mock pytest-repeat pytest-custom-exit-code pytest-timeout pytest-reraise\n - name: Build wheel\n uses: pypa/cibuildwheel@v2.10.2\n with:\n package-dir: ./\n env:\n CIBW_ENVIRONMENT: >\n STAN_BACKEND=\"${{ env.STAN_BACKEND }}\"\n CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014\n CIBW_BUILD: ${{ matrix.cpython-version }}\n CIBW_SKIP: \"*musllinux*\"\n CIBW_ARCHS: ${{ matrix.cibw_arch }}\n CIBW_BUILD_FRONTEND: build\n - name: Install annlite unix\n if: ${{ matrix.os == 'ubuntu-latest' }}\n run: |\n pip install wheelhouse/*linux*.whl\n - name: Install annlite win\n if: ${{ matrix.os == 'windows-latest'}}\n run: |\n pip install --pre --find-links=wheelhouse/ annlite\n - name: Test unix\n id: test_unix\n if: ${{ matrix.os == 'ubuntu-latest' || matrix.os == 'macos-latest' }}\n run: |\n cd ..\n mv annlite/tests/ ./\n pytest --suppress-no-test-exit-code --cov=annlite --cov-report=xml \\\n -v -s -m \"not gpu\" ${{ matrix.test-path }}\n echo \"::set-output name=codecov_flag::annlite\"\n timeout-minutes: 30\n - name: Test win\n id: test_win\n if: ${{ matrix.os == 'windows-latest'}}\n env:\n PYTHONIOENCODING: 'utf-8'\n run: |\n cd ..\n move annlite/tests/ ./\n cd tests/\n pytest -v -s -m \"not gpu\" -k \"test\"\n echo \"::set-output name=codecov_flag::annlite\"\n timeout-minutes: 30\n - name: Check codecov file\n id: check_files\n uses: andstor/file-existence-action@v1\n with:\n files: \"coverage.xml\"\n - name: Upload coverage from test to Codecov\n uses: codecov/codecov-action@v2\n if: steps.check_files.outputs.files_exists == 'true' && ${{ matrix.python-version }} == '3.7'\n with:\n file: coverage.xml\n flags: ${{ steps.test.outputs.codecov_flag }}\n fail_ci_if_error: false\n token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos\n\n # just for blocking the merge until all parallel core-test are successful\n success-all-test:\n needs: [core-test]\n if: always()\n runs-on: ubuntu-latest\n steps:\n - uses: technote-space/workflow-conclusion-action@v2\n - name: Check Failure\n if: env.WORKFLOW_CONCLUSION == 'failure'\n run: exit 1\n - name: Success\n if: ${{ success() }}\n run: echo \"All Done\"\n"
},
{
"path": ".github/workflows/force-release.yml",
"content": "name: Manual Release\n\non:\n workflow_dispatch:\n inputs:\n release_token:\n description: 'Your release token'\n required: true\n release_reason:\n description: 'Short reason for this manual release'\n required: true\n\njobs:\n token-check:\n runs-on: ubuntu-latest\n steps:\n - run: echo \"success!\"\n if: \"${{ github.event.inputs.release_token }} == ${{ env.release_token }}\"\n env:\n release_token: ${{ secrets.ANNLITE_RELEASE_TOKEN }}\n\n build-wheels:\n needs: [token-check]\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ ubuntu-latest, windows-latest, macos-latest ]\n cibw_arch: [ \"auto64\" ]\n python-version: [ [ '3.7', \"cp37-*\" ], [ '3.8', \"cp38-*\" ], [ '3.9', \"cp39-*\" ], [ '3.10', \"cp310-*\" ] ]\n steps:\n - uses: actions/checkout@v3\n with:\n fetch-depth: 100\n - name: Set up Python ${{ matrix.python-version[0] }}\n uses: actions/setup-python@v4\n with:\n python-version: ${{ matrix.python-version[0] }}\n - name: Build sdist\n if: ${{ matrix.os == 'ubuntu-latest' && matrix.python-version[0] == '3.7' }}\n run: |\n pip install build\n python -m build --sdist\n - name: Build wheels\n uses: pypa/cibuildwheel@v2.10.2\n with:\n package-dir: ./\n env:\n CIBW_ENVIRONMENT: >\n STAN_BACKEND=\"${{ env.STAN_BACKEND }}\"\n CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014\n CIBW_BUILD: ${{ matrix.python-version[1] }}\n CIBW_SKIP: \"*musllinux*\"\n CIBW_ARCHS: ${{ matrix.cibw_arch }}\n # CIBW_ARCH_MACOS: x86_64 arm64\n CIBW_BUILD_FRONTEND: build\n - uses: actions/upload-artifact@v3\n with:\n path: |\n ./wheelhouse/*.whl\n ./dist/*.tar.gz\n\n regular-release:\n needs: build-wheels\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n with:\n token: ${{ secrets.JINA_DEV_BOT }}\n fetch-depth: 100 # means max contribute history is limited to 100 lines\n# submodules: true\n - uses: actions/setup-python@v2\n with:\n python-version: 3.7\n - uses: actions/download-artifact@v3\n with:\n name: artifact\n - run: |\n git fetch --depth=1 origin +refs/tags/*:refs/tags/*\n npm install git-release-notes\n pip install twine wheel\n ./scripts/release.sh final \"${{ github.event.inputs.release_reason }}\" \"${{github.actor}}\"\n env:\n TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }}\n TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}\n JINA_SLACK_WEBHOOK: ${{ secrets.JINA_SLACK_WEBHOOK }}\n - if: failure()\n run: echo \"nothing to release\"\n - name: bumping master version\n uses: ad-m/github-push-action@v0.6.0\n with:\n github_token: ${{ secrets.JINA_DEV_BOT }}\n tags: true\n branch: main\n"
},
{
"path": ".github/workflows/tag.yml",
"content": "name: Release CD\n\non:\n push:\n tags:\n - \"v*\" # push to version tags trigger the build\n\n#on:\n# push:\n# branches-ignore:\n# - '**' # temporally disable this action\n\njobs:\n create-release:\n runs-on: ubuntu-latest\n steps:\n - name: Checkout code\n uses: actions/checkout@v2\n with:\n ref: 'main'\n - uses: actions/setup-python@v2\n with:\n python-version: 3.7\n - run: |\n python scripts/get-last-release-note.py\n - name: Create Release\n id: create_release\n uses: actions/create-release@v1\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # This token is provided by Actions, you do not need to create your own token\n with:\n tag_name: ${{ github.ref }}\n release_name: ๐ซ Patch ${{ github.ref }}\n body_path: 'tmp.md'\n draft: false\n prerelease: false\n"
},
{
"path": ".gitignore",
"content": "# Initially taken from Github's Python gitignore file\n# local dev\n.vscode\nbindings/pq_bindings.cpp\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\ndocs/.python-version\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n.idea/\ntoy*.py\n.DS_Store\npost/\ntoy*.ipynb\ndata/\n*.c\n.nes_cache\ntoy*.yml\n*.tmp\n\n/junit/\n/tests/junit/\n/docs/chapters/proto/docs.md\n\n# IntelliJ IDEA\n*.iml\n.idea\n\n# Rust\n/target\nCargo.lock\n\ntoy*.py\n"
},
{
"path": ".pre-commit-config.yaml",
"content": "repos:\n #- repo: https://github.com/terrencepreilly/darglint\n # rev: v1.5.8\n # hooks:\n # - id: darglint\n # files: annlite/\n # exclude: docs/\n # args:\n # - --message-template={path}:{line} {msg_id} {msg}\n # - -s=sphinx\n # - -z=full\n # - -v=2\n #- repo: https://github.com/pycqa/pydocstyle\n # rev: 5.1.1 # pick a git hash / tag to point to\n # hooks:\n # - id: pydocstyle\n # files: annlite/\n # exclude: docs/\n # args:\n # - --select=D101,D102,D103\n - repo: https://github.com/timothycrosley/isort\n rev: 5.12.0\n hooks:\n - id: isort\n args: [\"--profile\", \"black\"]\n - repo: https://github.com/psf/black\n rev: 22.3.0\n hooks:\n - id: black\n types: [python]\n exclude: docs/\n args:\n - -S\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v4.0.1\n hooks:\n - id: trailing-whitespace\n - id: check-yaml\n - id: end-of-file-fixer\n - id: requirements-txt-fixer\n - id: double-quote-string-fixer\n - id: check-merge-conflict\n - id: fix-encoding-pragma\n args: [\"--remove\"]\n - id: mixed-line-ending\n args: [\"--fix=lf\"]\n - repo: https://github.com/pre-commit/mirrors-clang-format\n rev: \"v14.0.3\"\n hooks:\n - id: clang-format\n"
},
{
"path": "CHANGELOG.md",
"content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n## Release Note (`0.3.0`)\n\n> Release time: 2022-03-04 11:49:06\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n numb3r3, felix-wang, David Buchaca Prats, ๐\n\n\n### ๐ New Features\n\n - [[```e4afbe17```](https://github.com/jina-ai/docarray/commit/e4afbe17c724f48de19bfbe6d6c127e62a8bcd5a)] __-__ add cd workflow (*numb3r3*)\n\n### ๐ Bug fixes\n\n - [[```2aeb5377```](https://github.com/jina-ai/docarray/commit/2aeb537779df6649fab143affa5c5decf8b2e364)] __-__ format (*numb3r3*)\n - [[```27ab16c2```](https://github.com/jina-ai/docarray/commit/27ab16c2622606e7bf77af5286620f0cdec4abb0)] __-__ build dep (*numb3r3*)\n - [[```7887ac95```](https://github.com/jina-ai/docarray/commit/7887ac95a7954c783264151cd26ae47319a47baa)] __-__ turn on upload pypi (#109) (*felix-wang*)\n - [[```2e4739ac```](https://github.com/jina-ai/docarray/commit/2e4739ac2f19e78bc23889ee2a1baef79be65e5d)] __-__ version (*numb3r3*)\n - [[```293a0dbf```](https://github.com/jina-ai/docarray/commit/293a0dbfe8b063c66ec4b8b710f44eb1c8ff4041)] __-__ release script (*numb3r3*)\n - [[```e77c95d9```](https://github.com/jina-ai/docarray/commit/e77c95d97fe54bb409ce39931de86fb45fdc6f41)] __-__ cov (*numb3r3*)\n - [[```287de379```](https://github.com/jina-ai/docarray/commit/287de379686fc4eb91afd99f32a85e3dcbc6bf27)] __-__ install in cd (*numb3r3*)\n - [[```88b8e0b6```](https://github.com/jina-ai/docarray/commit/88b8e0b6c033f33d76e5ec3e7e95f3de14d7ed70)] __-__ pip install in ci (*numb3r3*)\n - [[```348e0444```](https://github.com/jina-ai/docarray/commit/348e044463039c1a4ed60c43c1f83353378a86f7)] __-__ release (*numb3r3*)\n - [[```beb35f29```](https://github.com/jina-ai/docarray/commit/beb35f29c1f695728e7c43b235abb97a71f797ca)] __-__ setup pytest (*numb3r3*)\n - [[```e9ea4b89```](https://github.com/jina-ai/docarray/commit/e9ea4b8942d48034a7a471e2b91a6b542913f7c9)] __-__ workflow yml (*numb3r3*)\n - [[```5f0e21ec```](https://github.com/jina-ai/docarray/commit/5f0e21ec85a09a1dabc618841eb9fcc4c56c34bb)] __-__ ci yml (*numb3r3*)\n - [[```db58e283```](https://github.com/jina-ai/docarray/commit/db58e2835f758c94ec73fa3ccc795c286fdb4e86)] __-__ setup.py (*numb3r3*)\n - [[```bdc1c7a2```](https://github.com/jina-ai/docarray/commit/bdc1c7a276086b702c807112d0603c245476944a)] __-__ __cicd__: add scripts (*numb3r3*)\n\n### ๐งผ Code Refactoring\n\n - [[```9c693adb```](https://github.com/jina-ai/docarray/commit/9c693adb438742b633d04f554e442ceff9da923d)] __-__ remove pqlite namings (*David Buchaca Prats*)\n - [[```793f4711```](https://github.com/jina-ai/docarray/commit/793f4711c42d3a717e9e60e3548c222f1ecc7de2)] __-__ rename project (*numb3r3*)\n\n### ๐น Other Improvements\n\n - [[```f0c6d809```](https://github.com/jina-ai/docarray/commit/f0c6d809d47a6142f5ccf40c7cc7cf24045ad2e5)] __-__ bump version (*numb3r3*)\n\n\n## Release Note (`0.3.1`)\n\n> Release time: 2022-03-04 16:29:52\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```44298038```](https://github.com/jina-ai/docarray/commit/44298038e4e79f3dd0c3532ba6ad5b64d1a35caf)] __-__ add tag release workflow (#110) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```d15289f9```](https://github.com/jina-ai/docarray/commit/d15289f9912f158c7b9a7ec7d5f58188ecb19cf2)] __-__ __version__: the next version will be 0.3.1 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.2`)\n\n> Release time: 2022-06-09 10:04:30\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n David Buchaca Prats, Gustavo Ye, felix-wang, Han Xiao, numb3r3, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```a7804baf```](https://github.com/jina-ai/docarray/commit/a7804bafc90064f929afff62192f19e44406b5d1)] __-__ add filter method (#121) (*David Buchaca Prats*)\n\n### ๐ Bug fixes\n\n - [[```9448fcab```](https://github.com/jina-ai/docarray/commit/9448fcab060efe5b59108f3cac183aeb432a81d6)] __-__ only apply black with changed py files (#118) (*felix-wang*)\n - [[```e4e706e3```](https://github.com/jina-ai/docarray/commit/e4e706e313ba5cbfb7083a5dea9e75b8d2813394)] __-__ upgrade traverse path (#113) (*felix-wang*)\n - [[```a0fc7cac```](https://github.com/jina-ai/docarray/commit/a0fc7cac7722d108143f6590c3435c77376f76bd)] __-__ tag yml (*numb3r3*)\n\n### ๐น Other Improvements\n\n - [[```2ce1ec22```](https://github.com/jina-ai/docarray/commit/2ce1ec2283b381f5153ea60141a6bb474bbf0f0c)] __-__ __cpp/h__: clang-format (#117) (*Gustavo Ye*)\n - [[```a6168400```](https://github.com/jina-ai/docarray/commit/a61684000af4518aafa41d1d9dea47766bedf247)] __-__ update readme (*Han Xiao*)\n - [[```aec33c75```](https://github.com/jina-ai/docarray/commit/aec33c75b8ce883e044a828fbae7142f71dbb05a)] __-__ __version__: the next version will be 0.3.2 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.3`)\n\n> Release time: 2022-07-04 13:57:38\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Gustavo Ye, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```e240128f```](https://github.com/jina-ai/docarray/commit/e240128f403bde04bc21e1ac5ce3baa5507db687)] __-__ proto version bump (#126) (*felix-wang*)\n - [[```f62809b5```](https://github.com/jina-ai/docarray/commit/f62809b5d90f4fca1762c6a1c13beee40e972212)] __-__ change input type to `data_t` (#123) (*Gustavo Ye*)\n\n### ๐น Other Improvements\n\n - [[```f7f7b751```](https://github.com/jina-ai/docarray/commit/f7f7b75104a28f1860d3cdc94c87f526e742db51)] __-__ __version__: the next version will be 0.3.3 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.4`)\n\n> Release time: 2022-07-20 07:12:22\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```1a9de05c```](https://github.com/jina-ai/docarray/commit/1a9de05cdcbb861cd036b2ed9a3a66fa84e707f9)] __-__ copy executor.py to annlite folder (#129) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```5225e6cb```](https://github.com/jina-ai/docarray/commit/5225e6cbec7df3fbf4dbd7ab404d34bb7d899a29)] __-__ __version__: the next version will be 0.3.4 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.5`)\n\n> Release time: 2022-07-29 08:17:02\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```534eed81```](https://github.com/jina-ai/docarray/commit/534eed810e6b56238f0b71cc4123783236eac7fd)] __-__ support load and save hnsw (#133) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```1f76fb5d```](https://github.com/jina-ai/docarray/commit/1f76fb5dedc79d331dc655d6409718457c6c05d7)] __-__ Feat task queue (#131) (*Jie Fu*)\n - [[```103b5bf3```](https://github.com/jina-ai/docarray/commit/103b5bf31b59ca041d7ca70b8e0442421bcdf279)] __-__ __version__: the next version will be 0.3.5 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.6`)\n\n> Release time: 2022-08-23 10:45:36\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jie Fu, Gustavo Ye, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```690098be```](https://github.com/jina-ai/docarray/commit/690098be576fd5d6149e0502b9d2bf100b726f27)] __-__ add pq support for hnsw searching (#122) (*Gustavo Ye*)\n - [[```ad5a5fe3```](https://github.com/jina-ai/docarray/commit/ad5a5fe39293f771b9945102b45f05bebfaf3ad6)] __-__ indexer dumploader (#137) (*felix-wang*)\n - [[```29019e3d```](https://github.com/jina-ai/docarray/commit/29019e3da94d0a6025c91919a8978966f073b608)] __-__ integrate annlite with projector (#136) (*Jie Fu*)\n - [[```14c95986```](https://github.com/jina-ai/docarray/commit/14c9598602741a5558c02a98bd123b34cddc32b8)] __-__ implement pca (#135) (*Jie Fu*)\n\n### ๐ Bug fixes\n\n - [[```9f74de8f```](https://github.com/jina-ai/docarray/commit/9f74de8f34bc51ab740a05e9bfabf0be4ec77b87)] __-__ reload duplicate storage (#143) (*felix-wang*)\n - [[```7010d778```](https://github.com/jina-ai/docarray/commit/7010d77869dfd2a44ebaeb4a697bbac3e01d6970)] __-__ fix update/delete (#140) (*Jie Fu*)\n - [[```8a05887d```](https://github.com/jina-ai/docarray/commit/8a05887dc68d4219190de8767c98fc5a1740e3a2)] __-__ composite pca and pq (#139) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```157a1a9c```](https://github.com/jina-ai/docarray/commit/157a1a9c189f73f5757d2ec60c4427e41cd130bd)] __-__ __version__: the next version will be 0.3.6 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.7`)\n\n> Release time: 2022-08-26 04:48:47\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```01cf7a89```](https://github.com/jina-ai/docarray/commit/01cf7a8960f27abcbc414faf2d274e0d8da470db)] __-__ fix np.int64 issue (#146) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```b602acce```](https://github.com/jina-ai/docarray/commit/b602acce8f7f2ac085f9cbccae6fadeee1ebcb85)] __-__ __version__: the next version will be 0.3.7 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.8`)\n\n> Release time: 2022-08-26 11:12:19\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```a5eaa53e```](https://github.com/jina-ai/docarray/commit/a5eaa53e84aa374e061ca0bfddc7797fe369f8ec)] __-__ insert when not found at update (#151) (*felix-wang*)\n - [[```18af2d4c```](https://github.com/jina-ai/docarray/commit/18af2d4c42ced2ca9d393456ba9722953edce383)] __-__ normalize training data for cosine metric (#150) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```f0fb1c4b```](https://github.com/jina-ai/docarray/commit/f0fb1c4b37c89fea31fc24fb47b60181f0cd3218)] __-__ __version__: the next version will be 0.3.8 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.9`)\n\n> Release time: 2022-08-31 03:46:53\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Gustavo Ye, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```27589456```](https://github.com/jina-ai/docarray/commit/27589456fe72c3e9a1fe7d6c8d198e0e71fd9a89)] __-__ reload indexer (#154) (*felix-wang*)\n - [[```cdbe7256```](https://github.com/jina-ai/docarray/commit/cdbe725691cf5cc19db5df250d6d9d384bc94437)] __-__ __py__: roll back the original decorator (#142) (*Gustavo Ye*)\n\n### ๐น Other Improvements\n\n - [[```f34062b2```](https://github.com/jina-ai/docarray/commit/f34062b2249a3fce8cb51dc9e8b8497d33a54a2f)] __-__ __version__: the next version will be 0.3.9 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.10`)\n\n> Release time: 2022-09-06 09:40:07\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, numb3r3, Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```bc7f5ae2```](https://github.com/jina-ai/docarray/commit/bc7f5ae2e8e83cef919052a9d8a3a954a1cd813f)] __-__ allow columns as dict format (#160) (*felix-wang*)\n - [[```25c3ec3d```](https://github.com/jina-ai/docarray/commit/25c3ec3d616a9ece53c710d734ba71d0eb6fc19c)] __-__ thread safe index and search (#157) (*felix-wang*)\n - [[```d235cb83```](https://github.com/jina-ai/docarray/commit/d235cb83a2d8399f1120b3ebf6e54392ecb37efa)] __-__ executor unittest (#156) (*felix-wang*)\n\n### ๐ Unit Test and CICD\n\n - [[```e6543a28```](https://github.com/jina-ai/docarray/commit/e6543a28894aad5481d1a6a5bbd8f070ca979a90)] __-__ add docarray tests (#153) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```6a774648```](https://github.com/jina-ai/docarray/commit/6a774648d772ad98f42641c6094192bf4e7d1877)] __-__ update logo (*numb3r3*)\n - [[```33fd66be```](https://github.com/jina-ai/docarray/commit/33fd66be1135ebe63f2a1e1366196c8712e2bc63)] __-__ update readme (*numb3r3*)\n - [[```e1ad3a55```](https://github.com/jina-ai/docarray/commit/e1ad3a55d4732a115e6a925075b7b5ed768cd9f8)] __-__ update log (*numb3r3*)\n - [[```c2630257```](https://github.com/jina-ai/docarray/commit/c26302572de4e755d64ffee71b017e798f790e9e)] __-__ update readme (#149) (*felix-wang*)\n - [[```c1fd8cfe```](https://github.com/jina-ai/docarray/commit/c1fd8cfe746d94de2ebef843f6e13384f266773b)] __-__ __version__: the next version will be 0.3.10 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.11`)\n\n> Release time: 2022-09-08 06:17:25\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```0b435eb6```](https://github.com/jina-ai/docarray/commit/0b435eb6c5a737d370b389ba3c2395796febe95f)] __-__ bump pybind11 (#162) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```9c43f02c```](https://github.com/jina-ai/docarray/commit/9c43f02c5ec9face32ee58858b8ab25d65edbf12)] __-__ __version__: the next version will be 0.3.11 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.12`)\n\n> Release time: 2022-09-19 05:15:49\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, YangXiuyu, numb3r3, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```4eb61d3d```](https://github.com/jina-ai/docarray/commit/4eb61d3dea5a86d5212dafaaa0c1ecd87abc399a)] __-__ add limit and offset for filtering (#167) (*felix-wang*)\n\n### ๐ Documentation\n\n - [[```6ecfa833```](https://github.com/jina-ai/docarray/commit/6ecfa833e3d0f69a1a1b4992724753c89fa8d96f)] __-__ fix typo (#168) (*YangXiuyu*)\n\n### ๐น Other Improvements\n\n - [[```6902a8b1```](https://github.com/jina-ai/docarray/commit/6902a8b15763e73fd1cd90827ccfdc758f2fce32)] __-__ add description about hnsw parameters (#169) (*felix-wang*)\n - [[```aec2d605```](https://github.com/jina-ai/docarray/commit/aec2d605d43394638168282c78e13f7f2150200a)] __-__ update readme (*numb3r3*)\n - [[```95541bad```](https://github.com/jina-ai/docarray/commit/95541bad119dd2bb03e9894f32f57578cd3d8a7a)] __-__ __version__: the next version will be 0.3.12 (*Jina Dev Bot*)\n\n\n## Release Note (`0.3.13`)\n\n> Release time: 2022-09-26 02:04:59\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jie Fu, numb3r3, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```f98a8336```](https://github.com/jina-ai/docarray/commit/f98a83368e9fc81d82abb2be9c8a6569ee45e177)] __-__ change sqlite in_memory to false (#173) (*Jie Fu*)\n\n### ๐ Bug fixes\n\n - [[```b5f8694f```](https://github.com/jina-ai/docarray/commit/b5f8694fb5c33af15cc9b4a140a0165bf3aa01f4)] __-__ __hnswlib__: cannot find sufficient data with filtering (#176) (*felix-wang*)\n - [[```0aa07863```](https://github.com/jina-ai/docarray/commit/0aa078633d7b6b850a3ea0ffe651dbad06240cc0)] __-__ __ci__: unittest (#175) (*felix-wang*)\n\n### ๐งผ Code Refactoring\n\n - [[```896c5006```](https://github.com/jina-ai/docarray/commit/896c5006a6cac5f6b144bcb452ec5b2169ca4c88)] __-__ use rockdb as the storage backend (#171) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```66257b41```](https://github.com/jina-ai/docarray/commit/66257b4147836a8e8b7c17dd12537143509e78a2)] __-__ Revert "refactor: use rockdb as the storage backend (#171)" (*numb3r3*)\n - [[```e56aae75```](https://github.com/jina-ai/docarray/commit/e56aae757e65e5f8294a15b6ba998e00bda97e1a)] __-__ __version__: the next version will be 0.3.13 (*Jina Dev Bot*)\n\n\n## Release Note (`0.4.0`)\n\n> Release time: 2022-10-24 09:53:20\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n numb3r3, Jie Fu, YangXiuyu, felix-wang, Ziniu Yu, Gustavo Ye, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```1cd07f2c```](https://github.com/jina-ai/docarray/commit/1cd07f2c33e4f1f4e756baab4a6cd0a1452996cf)] __-__ hub as remote storage (#177) (*YangXiuyu*)\n - [[```90db6006```](https://github.com/jina-ai/docarray/commit/90db60067205d8e64a63c0c7eb4e782ac04d1033)] __-__ add flag for close (#191) (*Jie Fu*)\n - [[```3b782900```](https://github.com/jina-ai/docarray/commit/3b7829004edfb4f6cf4b77048a747a0d338c21b3)] __-__ cibuildwheels (#186) (*YangXiuyu*)\n\n### ๐ Bug fixes\n\n - [[```8c807937```](https://github.com/jina-ai/docarray/commit/8c80793733e96680d2ca846371ebc3c684f1d5d7)] __-__ restore when initializing annlite (#201) (*Jie Fu*)\n - [[```98f94d0f```](https://github.com/jina-ai/docarray/commit/98f94d0f4ff1bd0b1bc1b99d828d6fd21932dc2b)] __-__ set is_train=true after loading pca (#199) (*Jie Fu*)\n - [[```0b9bb413```](https://github.com/jina-ai/docarray/commit/0b9bb413abe631ea07bd7358160060b16cafa827)] __-__ unittest (#198) (*YangXiuyu*)\n - [[```0d1d5a2e```](https://github.com/jina-ai/docarray/commit/0d1d5a2e58eb043397f24a1d6efcee07c22e2cb1)] __-__ pre-release (*numb3r3*)\n - [[```fdcf3818```](https://github.com/jina-ai/docarray/commit/fdcf3818f1b8a2b11df96f4c6c72e972bab91238)] __-__ cd release version (#197) (*YangXiuyu*)\n - [[```f4bbc495```](https://github.com/jina-ai/docarray/commit/f4bbc49529a2ac4174bbb0cc7c88ab8666bd034c)] __-__ clean codes (*numb3r3*)\n - [[```6bedafe4```](https://github.com/jina-ai/docarray/commit/6bedafe4c49ce8c4831b4cc01c1b8702924e4b8f)] __-__ test-paths (*numb3r3*)\n - [[```beb79b23```](https://github.com/jina-ai/docarray/commit/beb79b23a1b1412bad402410a0b90393f0e535b6)] __-__ __cd__: combine build and test (*numb3r3*)\n - [[```19652ed0```](https://github.com/jina-ai/docarray/commit/19652ed0b6ef0ffc9dc1c3d2dc2640075e4d6d13)] __-__ cd workflow (#196) (*YangXiuyu*)\n - [[```4fe2b3be```](https://github.com/jina-ai/docarray/commit/4fe2b3be51ed0b571c6f0b9c9f20a68c18b7486c)] __-__ cd workflow (#195) (*YangXiuyu*)\n - [[```4fdacde2```](https://github.com/jina-ai/docarray/commit/4fdacde2a22a5e6fb0edc15bb6f77bf4d8f25834)] __-__ cd workflow (#193) (*YangXiuyu*)\n - [[```494fbfcb```](https://github.com/jina-ai/docarray/commit/494fbfcbb96257168542f56ec617a9dc9082c084)] __-__ cd release (#192) (*felix-wang*)\n - [[```d2803ce0```](https://github.com/jina-ai/docarray/commit/d2803ce00e41a69470a3001e4677927b407c3282)] __-__ cd workflow (#190) (*felix-wang*)\n - [[```fe9db3d9```](https://github.com/jina-ai/docarray/commit/fe9db3d9f7e04cbad96832409dd1d3159195c060)] __-__ build on apple silicon (#188) (*Ziniu Yu*)\n - [[```bb402ae9```](https://github.com/jina-ai/docarray/commit/bb402ae9674a8410686b8dde3aea82f1b86fc10b)] __-__ __bindings__: build on windows (#183) (*YangXiuyu*)\n - [[```efc18f80```](https://github.com/jina-ai/docarray/commit/efc18f80228ef2371c19fbb205cabb94afd47385)] __-__ update executor parameter (#180) (*YangXiuyu*)\n - [[```f142089a```](https://github.com/jina-ai/docarray/commit/f142089a419bba3d7e23b2c34b16455ddbcb805d)] __-__ executor tests (#179) (*felix-wang*)\n - [[```79e171d4```](https://github.com/jina-ai/docarray/commit/79e171d411f31dad7231bce1682ac78ff7b3e1b2)] __-__ __annliteindexer__: annlite executor integration (#170) (*YangXiuyu*)\n\n### ๐งผ Code Refactoring\n\n - [[```e8c59907```](https://github.com/jina-ai/docarray/commit/e8c59907540bada5282a413b43d64f11596678d3)] __-__ use rocksdb as the docs storage engine (#178) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```e6dfcd2a```](https://github.com/jina-ai/docarray/commit/e6dfcd2a7ff0fdd149dd606487f7aa0107775c63)] __-__ bump to v0.4.0 (*numb3r3*)\n - [[```6ce03fd0```](https://github.com/jina-ai/docarray/commit/6ce03fd04e5bb63da3c3d2cf1c620ac77daf6c2a)] __-__ release (*numb3r3*)\n - [[```02857ec4```](https://github.com/jina-ai/docarray/commit/02857ec4103e32139f695adac4f7a40c4e65c67a)] __-__ Add pq dist table support (#158) (*Gustavo Ye*)\n - [[```ff54290b```](https://github.com/jina-ai/docarray/commit/ff54290bfb19d07b953aa00d5d52cb1e8f805d3a)] __-__ __version__: the next version will be 0.3.14 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.0`)\n\n> Release time: 2022-10-26 07:12:58\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n numb3r3, felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```de5273ab```](https://github.com/jina-ai/docarray/commit/de5273ab7cf81ba26bb4307cd66f3fc50974b8ef)] __-__ undo wheels uploading (#203) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```5faf1dfa```](https://github.com/jina-ai/docarray/commit/5faf1dfa3326213c3b9e20de0d8b7aca4dbdf6ba)] __-__ regular release (*numb3r3*)\n - [[```95e51f26```](https://github.com/jina-ai/docarray/commit/95e51f269ac130b9d54b004734ff68afc07a9cf1)] __-__ __version__: the next version will be 0.4.1 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.1`)\n\n> Release time: 2022-11-03 03:29:55\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n numb3r3, Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ New Features\n\n - [[```dd08f6f2```](https://github.com/jina-ai/docarray/commit/dd08f6f28008aaf020cb972f45e154afaf99cb28)] __-__ add file splitter and merger to support big file transfer to hubble (#202) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```83e45001```](https://github.com/jina-ai/docarray/commit/83e4500180e98bb967ab2ec17ca202c1702fb411)] __-__ bump 0.5.1 (*numb3r3*)\n - [[```557c624b```](https://github.com/jina-ai/docarray/commit/557c624b3419b8bf70386e66a75d1958ef427738)] __-__ fix force release (*numb3r3*)\n - [[```77312214```](https://github.com/jina-ai/docarray/commit/773122144a9ec39663829c64090d6116585fb6e1)] __-__ __version__: the next version will be 0.5.1 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.2`)\n\n> Release time: 2022-11-03 07:10:10\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```3f886f10```](https://github.com/jina-ai/docarray/commit/3f886f10d17195913cea839f3e4ac8631426ad5c)] __-__ restore from latest snapshot (#204) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```2462afa4```](https://github.com/jina-ai/docarray/commit/2462afa4451c66603cc1e25bf755c97a97b7eec1)] __-__ __version__: the next version will be 0.5.2 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.4`)\n\n> Release time: 2022-11-04 04:27:28\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```f2ce3a9a```](https://github.com/jina-ai/docarray/commit/f2ce3a9a4b35a71a492fbe63203114e0d7fa6b52)] __-__ bump rocksdict (#206) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```c84635c9```](https://github.com/jina-ai/docarray/commit/c84635c9854b7b3c2b5718941979720b673ed7eb)] __-__ __version__: the next version will be 0.5.3 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.5`)\n\n> Release time: 2022-12-22 03:27:38\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```61e0a57f```](https://github.com/jina-ai/docarray/commit/61e0a57f901b087fbe6cab16aefb1d46d5fd54a2)] __-__ fix offset factory (#208) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```f6fdc9d3```](https://github.com/jina-ai/docarray/commit/f6fdc9d34fd8e8de47e7ac68e6813956c2aeebfe)] __-__ __version__: the next version will be 0.5.5 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.6`)\n\n> Release time: 2023-02-08 09:06:21\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```906ed7d4```](https://github.com/jina-ai/docarray/commit/906ed7d4159527e8b5f00cdc9106433925f04a43)] __-__ save and load offset2ids for list_like feature (#210) (*Jie Fu*)\n\n### ๐ Unit Test and CICD\n\n - [[```a2519f0f```](https://github.com/jina-ai/docarray/commit/a2519f0f323d0c3cb2d68ad1cd3eeed4beb40ce6)] __-__ add test for offset2ids saving and loading (#212) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```68215807```](https://github.com/jina-ai/docarray/commit/68215807953eed366e1d6f48fbbbbc1abebf13ff)] __-__ __version__: the next version will be 0.5.6 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.7`)\n\n> Release time: 2023-02-21 12:34:55\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```4a66d0c0```](https://github.com/jina-ai/docarray/commit/4a66d0c06e7497cb94cd18da505a2e309594b20c)] __-__ remove delete tags for cell_table and meta_table (#215) (*Jie Fu*)\n\n### ๐น Other Improvements\n\n - [[```eee752fe```](https://github.com/jina-ai/docarray/commit/eee752fe091dc5b5b519bb2aa866cb7d5dc48822)] __-__ __version__: the next version will be 0.5.7 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.8`)\n\n> Release time: 2023-03-24 11:13:35\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, YangXiuyu, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```bfeb7e94```](https://github.com/jina-ai/docarray/commit/bfeb7e94c879db263dc4931b54b6b87dac09df81)] __-__ delete and update (#222) (*felix-wang*)\n - [[```7783b30f```](https://github.com/jina-ai/docarray/commit/7783b30fd55a10bd2dd0c75dc0451f0fe1834927)] __-__ __hnsw__: bump hnswlib to v0.6.2 (#185) (*YangXiuyu*)\n\n### ๐น Other Improvements\n\n - [[```4c145ddd```](https://github.com/jina-ai/docarray/commit/4c145ddd19abb4caec479941d1c0ffb03c4cfcf3)] __-__ __version__: the next version will be 0.5.8 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.9`)\n\n> Release time: 2023-04-07 04:16:28\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n felix-wang, Jina Dev Bot, ๐\n\n\n### ๐ Bug fixes\n\n - [[```ab807ff6```](https://github.com/jina-ai/docarray/commit/ab807ff69e533ec90e0b0f6c782749267324d177)] __-__ bump rocksdb >= 0.3.9 (#225) (*felix-wang*)\n\n### ๐น Other Improvements\n\n - [[```3668c3fa```](https://github.com/jina-ai/docarray/commit/3668c3fa6f7a2d2af39549987637e11e385a75bb)] __-__ __version__: the next version will be 0.5.9 (*Jina Dev Bot*)\n\n\n## Release Note (`0.5.10`)\n\n> Release time: 2023-04-19 03:12:50\n\n\n\n๐ We'd like to thank all contributors for this new release! In particular,\n Jie Fu, Jina Dev Bot, ๐\n\n\n### ๐น Other Improvements\n\n - [[```d55d544e```](https://github.com/jina-ai/docarray/commit/d55d544e876653bb47e79e2dd00f859c3ef00ffa)] __-__ bumping docarray (#227) (*Jie Fu*)\n - [[```86057c69```](https://github.com/jina-ai/docarray/commit/86057c69401e2b6d63822be2c3c3a0f63d4661a6)] __-__ __version__: the next version will be 0.5.10 (*Jina Dev Bot*)\n\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "MANIFEST.in",
"content": "include setup.py\ninclude requirements.txt\ninclude pyproject.toml\nglobal-include *.pyx\nrecursive-include include/hnswlib/ *\n"
},
{
"path": "Makefile",
"content": "pypi: dist\n\ttwine upload dist/*\n\ndist:\n\trm -rf dist/*\n\tpip install build\n\tpython -m build --sdist\n\ntest:\n\tpython -m unittest discover --start-directory tests --pattern \"*_test*.py\"\n\nclean:\n\trm -rf *.egg-info build dist tmp var tests/__pycache__ annlite/*.so\n\n.PHONY: dist\n"
},
{
"path": "README.md",
"content": "\n
\n
\n
\n![\"AnnLite](\"https://github.com/jina-ai/annlite/blob/main/.github/assets/logo.svg?raw=true\")
\n
\n
\nA fast embedded library for approximate nearest neighbor search\n
\n\n\n![\"GitHub\"](\"https://img.shields.io/github/license/jina-ai/annlite?style=flat-square\")
\n![\"PyPI\"](\"https://img.shields.io/pypi/v/annlite?label=Release&style=flat-square\")
\n![\"Codecov](\"https://img.shields.io/codecov/c/github/jina-ai/annlite/main?logo=Codecov&logoColor=white&style=flat-square\")
\n![](\"https://img.shields.io/badge/Slack-3.1k-blueviolet?logo=slack&logoColor=white&style=flat-square\")
\n
\n\n\n\n\n## What is AnnLite?\n\n`AnnLite` is a *lightweight* and *embeddable* library for **fast** and **filterable** *approximate nearest neighbor search* (ANNS).\nIt allows to search for nearest neighbors in a dataset of millions of points with a Pythonic API.\n\n\n**Highlighted features:**\n\n- ๐ฅ **Easy-to-use**: a simple API is designed to be used with Python. It is easy to use and intuitive to set up to production.\n\n- ๐ **Fast**: the library uses a highly optimized approximate nearest neighbor search algorithm (*HNSW*) to search for nearest neighbors.\n\n- ๐ **Filterable**: the library allows you to search for nearest neighbors within a subset of the dataset.\n\n- ๐ฑ **Integration**: Smooth integration with neural search ecosystem including [Jina](https://github.com/jina-ai/jina) and [DocArray](https://github.com/jina-ai/docarray),\n so that users can easily expose search API with **gRPC** and/or **HTTP**.\n\nThe library is easy to install and use. It is designed to be used with Python.\n\n\n\n## Installation\n\nTo use AnnLite, you need to first install it. The easiest way to install AnnLite is using `pip`:\n\n```bash\npip install -U annlite\n```\n\nor install from source:\n\n```bash\npython setup.py install\n```\n\n## Quick start\n\nBefore you start, you need to know some experience about [DocArray](https://github.com/jina-ai/docarray).\n`AnnLite` is designed to be used with [DocArray](https://github.com/jina-ai/docarray), so you need to know how to use `DocArray` first.\n\nFor example, you can create a `DocArray` with `1000` random vectors with `128` dimensions:\n\n```python\nfrom docarray import DocumentArray\nimport numpy as np\n\ndocs = DocumentArray.empty(1000)\ndocs.embeddings = np.random.random([1000, 128]).astype(np.float32)\n```\n\n\n### Index\n\nThen you can create an `AnnIndexer` to index the created `docs` and search for nearest neighbors:\n\n```python\nfrom annlite import AnnLite\n\nann = AnnLite(128, metric='cosine', data_path=\"/tmp/annlite_data\")\nann.index(docs)\n```\n\nNote that this will create a directory `/tmp/annlite_data` to persist the documents indexed.\nIf this directory already exists, the index will be loaded from the directory.\nAnd if you want to create a new index, you can delete the directory first.\n\n### Search\n\nThen you can search for nearest neighbors for some query docs with `ann.search()`:\n\n```python\nquery = DocumentArray.empty(5)\nquery.embeddings = np.random.random([5, 128]).astype(np.float32)\n\nresult = ann.search(query)\n```\n\nThen, you can inspect the retrieved docs for each query doc inside `query` matches:\n```python\nfor q in query:\n print(f'Query {q.id}')\n for k, m in enumerate(q.matches):\n print(f'{k}: {m.id} {m.scores[\"cosine\"]}')\n```\n\n```bash\nQuery ddbae2073416527bad66ff186543eff8\n0: 47dcf7f3fdbe3f0b8d73b87d2a1b266f {'value': 0.17575037}\n1: 7f2cbb8a6c2a3ec7be024b750964f317 {'value': 0.17735684}\n2: 2e7eed87f45a87d3c65c306256566abb {'value': 0.17917466}\nQuery dda90782f6514ebe4be4705054f74452\n0: 6616eecba99bd10d9581d0d5092d59ce {'value': 0.14570713}\n1: d4e3147fc430de1a57c9883615c252c6 {'value': 0.15338594}\n2: 5c7b8b969d4381f405b8f07bc68f8148 {'value': 0.15743542}\n...\n```\n\nOr shorten the loop as one-liner using the element & attribute selector:\n\n```python\nprint(query['@m', ('id', 'scores__cosine')])\n```\n\n### Query\n\nYou can get specific document by its id:\n\n```python\ndoc = ann.get_doc_by_id('')\n```\n\nAnd you can also get the documents with `limit` and `offset`, which is useful for pagination:\n\n```python\ndocs = ann.get_docs(limit=10, offset=0)\n```\n\nFurthermore, you can also get the documents ordered by a specific column from the index:\n\n```python\ndocs = ann.get_docs(limit=10, offset=0, order_by='x', ascending=True)\n```\n\n**Note**: the `order_by` column must be one of the `columns` in the index.\n\n### Update\n\nAfter you have indexed the `docs`, you can update the docs in the index by calling `ann.update()`:\n\n```python\nupdated_docs = docs.sample(10)\nupdated_docs.embeddings = np.random.random([10, 128]).astype(np.float32)\n\nann.update(updated_docs)\n```\n\n\n### Delete\n\nAnd finally, you can delete the docs from the index by calling `ann.delete()`:\n\n```python\nto_delete = docs.sample(10)\nann.delete(to_delete)\n```\n\n## Search with filters\n\nTo support search with filters, the annlite must be created with `colums` parameter, which is a series of fields you want to filter by.\nAt the query time, the annlite will filter the dataset by providing `conditions` for certain fields.\n\n```python\nimport annlite\n\n# the column schema: (name:str, dtype:type, create_index: bool)\nann = annlite.AnnLite(128, columns=[('price', float)], data_path=\"/tmp/annlite_data\")\n```\n\nThen you can insert the docs, in which each doc has a field `price` with a float value contained in the `tags`:\n\n\n```python\nimport random\n\ndocs = DocumentArray.empty(1000)\ndocs = DocumentArray(\n [\n Document(id=f'{i}', tags={'price': random.random()})\n for i in range(1000)\n ]\n)\n\ndocs.embeddings = np.random.random([1000, 128]).astype(np.float32)\n\nann.index(docs)\n```\n\nThen you can search for nearest neighbors with filtering conditions as:\n\n```python\nquery = DocumentArray.empty(5)\nquery.embeddings = np.random.random([5, 128]).astype(np.float32)\n\nann.search(query, filter={\"price\": {\"$lte\": 50}}, limit=10)\nprint(f'the result with filtering:')\nfor i, q in enumerate(query):\n print(f'query [{i}]:')\n for m in q.matches:\n print(f'\\t{m.id} {m.scores[\"euclidean\"].value} (price={m.tags[\"price\"]})')\n```\n\nThe `conditions` parameter is a dictionary of conditions. The key is the field name, and the value is a dictionary of conditions.\nThe query language is the same as [MongoDB Query Language](https://docs.mongodb.com/manual/reference/operator/query/).\nWe currently support a subset of those selectors.\n\n- `$eq` - Equal to (number, string)\n- `$ne` - Not equal to (number, string)\n- `$gt` - Greater than (number)\n- `$gte` - Greater than or equal to (number)\n- `$lt` - Less than (number)\n- `$lte` - Less than or equal to (number)\n- `$in` - Included in an array\n- `$nin` - Not included in an array\n\n\nThe query will be performed on the field if the condition is satisfied. The following is an example of a query:\n\n1. A Nike shoes with white color\n\n ```python\n {\n \"brand\": {\"$eq\": \"Nike\"},\n \"category\": {\"$eq\": \"Shoes\"},\n \"color\": {\"$eq\": \"White\"}\n }\n ```\n\n We also support boolean operators `$or` and `$and`:\n\n ```python\n {\n \"$and\":\n {\n \"brand\": {\"$eq\": \"Nike\"},\n \"category\": {\"$eq\": \"Shoes\"},\n \"color\": {\"$eq\": \"White\"}\n }\n }\n ```\n\n2. A Nike shoes or price less than `100$`:\n\n ```python\n {\n \"$or\":\n {\n \"brand\": {\"$eq\": \"Nike\"},\n \"price\": {\"$lte\": 100}\n }\n }\n ```\n\n## Dump and Load\n\nBy default, the hnsw index is in memory. You can dump the index to `data_path` by calling `.dump()`:\n\n```python\n\nfrom annlite import AnnLite\n\nann = AnnLite(128, metric='cosine', data_path=\"/path/to/data_path\")\nann.index(docs)\nann.dump()\n```\n\nAnd you can restore the hnsw index from `data_path` if it exists:\n\n```python\nnew_ann = AnnLite(128, metric='cosine', data_path=\"/path/to/data_path\")\n```\n\nIf you didn't dump the hnsw index, the index will be rebuilt from scratch. This will take a while.\n\n## Supported distance metrics\n\nThe annlite supports the following distance metrics:\n\n#### Supported distances:\n\n| Distance | parameter | Equation |\n|----------------------------------------------------------------------|----------------:|--------------------------------------------------------:|\n| [Euclidean](https://en.wikipedia.org/wiki/Euclidean_distance) | `euclidean` | d = sqrt(sum((Ai-Bi)^2)) |\n| [Inner product](https://en.wikipedia.org/wiki/Inner_product_space) | `inner_product` | d = 1.0 - sum(Ai\\*Bi) |\n| [Cosine similarity](https://en.wikipedia.org/wiki/Cosine_similarity) | `cosine` | d = 1.0 - sum(Ai\\*Bi) / sqrt(sum(Ai\\*Ai) * sum(Bi\\*Bi)) |\n\nNote that inner product is not an actual metric. An element can be closer to some other element than to itself.\nThat allows some speedup if you remove all elements that are not the closest to themselves from the index, e.g.,\n`inner_product([1.0, 1.0], [1.0. 1.0]) < inner_product([1.0, 1.0], [2.0, 2.0])`\n\n\n## HNSW algorithm parameters\n\nThe HNSW algorithm has several parameters that can be tuned to improve the search performance.\n\n### Search parameters\n\n- `ef_search` - The size of the dynamic list for the nearest neighbors during search (default: `50`).\nThe larger the value, the more accurate the search results, but the slower the search speed.\nThe `ef_search` must be larger than `limit` parameter in `search(..., limit)`.\n\n- `limit` - The maximum number of results to return (default: `10`).\n\n## Construction parameters\n\n- `max_connection` - The number of bi-directional links created for every new element during construction (default: `16`).\nReasonable range is from `2` to `100`. Higher values works better for dataset with higher dimensionality and/or high recall.\nThis parameter also affects the memory consumption during construction, which is roughly `max_connection * 8-10` bytes per stored element.\n\n As an example for `n_dim=4` random vectors optimal `max_connection` for search is somewhere around `6`,\nwhile for high dimensional datasets, higher `max_connection` are required (e.g. `M=48-64`) for optimal performance at high recall.\nThe range `max_connection=12-48` is ok for the most of the use cases.\nWhen `max_connection` is changed one has to update the other parameters.\nNonetheless, `ef_search` and `ef_construction` parameters can be roughly estimated by assuming that `max_connection * ef_{construction}` is a constant.\n\n\n- `ef_construction`: The size of the dynamic list for the nearest neighbors during construction (default: `200`).\nHigher values give better accuracy, but increase construction time and memory consumption.\nAt some point, increasing `ef_construction` does not give any more accuracy.\nTo set `ef_construction` to a reasonable value, one can measure the recall: if the recall is lower than 0.9, then increase `ef_construction` and re-run the search.\n\nTo set the parameters, you can define them when creating the annlite:\n\n```python\nfrom annlite import AnnLite\n\nann = AnnLite(128, columns=[('price', float)], data_path=\"/tmp/annlite_data\", ef_construction=200, max_connection=16)\n```\n\n## Benchmark\n\nOne can run `executor/benchmark.py` to get a quick performance overview.\n\n|Stored data| Indexing time | Query size=1 | Query size=8 | Query size=64|\n|---|---|---|---|---|\n|10000 | 2.970 | 0.002 | 0.013 | 0.100|\n|100000 | 76.474 | 0.011 | 0.078 | 0.649|\n|500000 | 467.936 | 0.046 | 0.356 | 2.823|\n|1000000 | 1025.506 | 0.091 | 0.695 | 5.778|\n\nResults with filtering can be generated from `examples/benchmark_with_filtering.py`. This script should produce a table similar to:\n\n| Stored data |% same filter| Indexing time | Query size=1 | Query size=8 | Query size=64|\n|-----|-----|-----|-----|-----|-----|\n| 10000 | 5 | 2.869 | 0.004 | 0.030 | 0.270 |\n| 10000 | 15 | 2.869 | 0.004 | 0.035 | 0.294 |\n| 10000 | 20 | 3.506 | 0.005 | 0.038 | 0.287 |\n| 10000 | 30 | 3.506 | 0.005 | 0.044 | 0.356 |\n| 10000 | 50 | 3.506 | 0.008 | 0.064 | 0.484 |\n| 10000 | 80 | 2.869 | 0.013 | 0.098 | 0.910 |\n| 100000 | 5 | 75.960 | 0.018 | 0.134 | 1.092 |\n| 100000 | 15 | 75.960 | 0.026 | 0.211 | 1.736 |\n| 100000 | 20 | 78.475 | 0.034 | 0.265 | 2.097 |\n| 100000 | 30 | 78.475 | 0.044 | 0.357 | 2.887 |\n| 100000 | 50 | 78.475 | 0.068 | 0.565 | 4.383 |\n| 100000 | 80 | 75.960 | 0.111 | 0.878 | 6.815 |\n| 500000 | 5 | 497.744 | 0.069 | 0.561 | 4.439 |\n| 500000 | 15 | 497.744 | 0.134 | 1.064 | 8.469 |\n| 500000 | 20 | 440.108 | 0.152 | 1.199 | 9.472 |\n| 500000 | 30 | 440.108 | 0.212 | 1.650 | 13.267 |\n| 500000 | 50 | 440.108 | 0.328 | 2.637 | 21.961 |\n| 500000 | 80 | 497.744 | 0.580 | 4.602 | 36.986 |\n| 1000000 | 5 | 1052.388 | 0.131 | 1.031 | 8.212 |\n| 1000000 | 15 | 1052.388 | 0.263 | 2.191 | 16.643 |\n| 1000000 | 20 | 980.598 | 0.351 | 2.659 | 21.193 |\n| 1000000 | 30 | 980.598 | 0.461 | 3.713 | 29.794 |\n| 1000000 | 50 | 980.598 | 0.732 | 5.975 | 47.356 |\n| 1000000 | 80 | 1052.388 | 1.151 | 9.255 | 73.552 |\n\n\nNote that:\n- query times presented are represented in seconds.\n- `% same filter` indicates the amount of data that verifies a filter in the database.\n - For example, if `% same filter = 10` and `Stored data = 1_000_000` then it means `100_000` example verify the filter.\n\n\n## Next steps\n\nIf you already have experience with Jina and DocArray, you can start using `AnnLite` right away.\n\nOtherwise, you can check out this advanced tutorial to learn how to use `AnnLite`: [here]() in practice.\n\n\n## ๐ FAQ\n\n**1. Why should I use `AnnLite`?**\n\n`AnnLite` is easy to use and intuitive to set up in production. It is also very fast and memory efficient, making it a great choice for approximate nearest neighbor search.\n\n**2. How do I use `AnnLite` with Jina?**\n\nWe have implemented an executor for `AnnLite` that can be used with Jina.\n\n```python\nfrom jina import Flow\n\nwith Flow().add(uses='jinahub://AnnLiteIndexer', uses_with={'n_dim': 128}) as f:\n f.post('/index', inputs=docs)\n```\n\n3. Does `AnnLite` support search with filters?\n\n```text\nYes.\n```\n\n\n## Documentation\n\nYou can find the documentation on [Github]() and [ReadTheDocs]()\n\n## ๐ค Contribute and spread the word\n\nWe are also looking for contributors who want to help us improve: code, documentation, issues, feedback! Here is how you can get started:\n\n- Have a look through GitHub issues labeled \"Good first issue\".\n- Read our Contributor Covenant Code of Conduct\n- Open an issue or submit your pull request!\n\n\n## License\n\n`AnnLite` is licensed under the [Apache License 2.0]().\n"
},
{
"path": "annlite/__init__.py",
"content": "__version__ = '0.5.11'\n\nfrom .index import AnnLite\n"
},
{
"path": "annlite/container.py",
"content": "import warnings\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Dict, List, Optional, Tuple\n\nimport numpy as np\nfrom docarray import Document, DocumentArray\nfrom loguru import logger\n\nif TYPE_CHECKING:\n from .core.codec.pq import PQCodec\n from .core.codec.projector import ProjectorCodec\n\nfrom .core.index.hnsw import HnswIndex\nfrom .enums import Metric\nfrom .storage.base import ExpandMode\nfrom .storage.kv import DocStorage\nfrom .storage.table import CellTable, MetaTable\n\nVALID_FILTERABLE_DATA_TYPES = [int, str, float]\n\n\nclass CellContainer:\n def __init__(\n self,\n n_dim: int,\n metric: Metric = Metric.COSINE,\n n_cells: int = 1,\n projector_codec: Optional['ProjectorCodec'] = None,\n pq_codec: Optional['PQCodec'] = None,\n initial_size: Optional[int] = None,\n expand_step_size: int = 50000,\n expand_mode: 'ExpandMode' = ExpandMode.STEP,\n filterable_attrs: Optional[Dict] = None,\n serialize_config: Optional[Dict] = None,\n data_path: 'Path' = Path('./data'),\n **kwargs,\n ):\n self.n_dim = n_dim\n self.metric = metric\n self.n_cells = n_cells\n self.n_components = projector_codec.n_components if projector_codec else None\n self.data_path = data_path\n self.serialize_config = serialize_config\n\n self._pq_codec = pq_codec\n self._projector_codec = projector_codec\n\n self._vec_indexes = [\n HnswIndex(\n dim=self.n_components or n_dim,\n metric=metric,\n initial_size=initial_size,\n expand_step_size=expand_step_size,\n expand_mode=expand_mode,\n pq_codec=pq_codec,\n **kwargs,\n )\n for _ in range(n_cells)\n ]\n\n self._doc_stores = [\n DocStorage(\n data_path / f'cell_{_}',\n serialize_config=serialize_config or {},\n lock=True,\n )\n for _ in range(n_cells)\n ]\n\n columns = []\n if filterable_attrs:\n for attr_name, attr_type in filterable_attrs.items():\n if isinstance(attr_type, str):\n attr_type = eval(attr_type)\n\n if attr_type not in VALID_FILTERABLE_DATA_TYPES:\n raise ValueError(\n f'Invalid filterable attribute type `{attr_type}` for attribute `{attr_name}`. '\n )\n columns.append((attr_name, attr_type))\n\n self._cell_tables = [\n CellTable(f'table_{c}', columns=columns) for c in range(n_cells)\n ]\n\n self._meta_table = MetaTable('metas', data_path=data_path, in_memory=True)\n\n def ivf_search(\n self,\n x: 'np.ndarray',\n cells: 'np.ndarray',\n where_clause: str = '',\n where_params: Tuple = (),\n limit: int = 10,\n ):\n dists = []\n\n doc_idx = []\n cell_ids = []\n count = 0\n for cell_id in cells:\n cell_table = self.cell_table(cell_id)\n cell_size = cell_table.count()\n if cell_size == 0:\n continue\n\n indices = None\n if where_clause:\n indices = cell_table.query(\n where_clause=where_clause, where_params=where_params\n )\n\n if len(indices) == 0:\n continue\n\n indices = np.array(indices, dtype=np.int64)\n\n _dists, _doc_idx = self.vec_index(cell_id).search(\n x, limit=min(limit, cell_size), indices=indices\n )\n\n if count >= limit and _dists[0] > dists[-1][-1]:\n continue\n\n dists.append(_dists)\n doc_idx.append(_doc_idx)\n cell_ids.extend([cell_id] * len(_dists))\n count += len(_dists)\n\n cell_ids = np.array(cell_ids, dtype=np.int64)\n if len(dists) != 0:\n dists = np.hstack(dists)\n doc_idx = np.hstack(doc_idx)\n\n indices = dists.argsort(axis=0)[:limit]\n dists = dists[indices]\n cell_ids = cell_ids[indices]\n doc_idx = doc_idx[indices]\n\n doc_ids = []\n for cell_id, offset in zip(cell_ids, doc_idx):\n doc_id = self.cell_table(cell_id).get_docid_by_offset(offset)\n doc_ids.append(doc_id)\n return dists, doc_ids, cell_ids\n\n def filter_cells(\n self,\n cells: 'np.ndarray',\n where_clause: str = '',\n where_params: Tuple = (),\n limit: int = -1,\n offset: int = 0,\n order_by: Optional[str] = None,\n ascending: bool = True,\n include_metadata: bool = False,\n ):\n result = DocumentArray()\n if len(cells) > 1 and offset > 0:\n raise ValueError('Offset is not supported for multiple cells')\n\n for cell_id in cells:\n cell_table = self.cell_table(cell_id)\n cell_size = cell_table.count()\n if cell_size == 0:\n continue\n\n indices = cell_table.query(\n where_clause=where_clause,\n where_params=where_params,\n order_by=order_by,\n limit=limit,\n offset=offset,\n ascending=ascending,\n )\n\n if len(indices) == 0:\n continue\n\n for offset in indices:\n doc_id = self.cell_table(cell_id).get_docid_by_offset(offset)\n doc = Document(id=doc_id)\n if include_metadata or (len(cells) > 1 and order_by):\n doc = self.doc_store(cell_id).get([doc_id])[0]\n\n result.append(doc)\n\n if not order_by and len(result) >= limit > 0:\n break\n\n # reordering the results from multiple cells\n if order_by and len(cells) > 1:\n result = sorted(\n result, key=lambda d: d.tags.get(order_by), reverse=not ascending\n )\n if limit > 0:\n result = result[:limit]\n result = DocumentArray(result)\n\n return result\n\n def search_cells(\n self,\n query: 'np.ndarray',\n cells: 'np.ndarray',\n where_clause: str = '',\n where_params: Tuple = (),\n limit: int = 10,\n include_metadata: bool = False,\n ):\n if self._projector_codec:\n query = self._projector_codec.encode(query)\n\n topk_dists, topk_docs = [], []\n for x, cell_idx in zip(query, cells):\n # x.shape = (self.n_dim,)\n dists, doc_ids, cells = self.ivf_search(\n x,\n cells=cell_idx,\n where_clause=where_clause,\n where_params=where_params,\n limit=limit,\n )\n\n topk_dists.append(dists)\n match_docs = DocumentArray()\n for dist, doc_id, cell_id in zip(dists, doc_ids, cells):\n doc = Document(id=doc_id)\n if include_metadata:\n doc = self.doc_store(cell_id).get([doc_id])[0]\n\n doc.scores[self.metric.name.lower()].value = dist\n match_docs.append(doc)\n topk_docs.append(match_docs)\n\n return topk_dists, topk_docs\n\n def _search_cells(\n self,\n query: 'np.ndarray',\n cells: 'np.ndarray',\n where_clause: str = '',\n where_params: Tuple = (),\n limit: int = 10,\n ):\n if self._projector_codec:\n query = self._projector_codec.encode(query)\n\n topk_dists, topk_ids = [], []\n for x, cell_idx in zip(query, cells):\n dists, ids, cells = self.ivf_search(\n x,\n cells=cell_idx,\n where_clause=where_clause,\n where_params=where_params,\n limit=limit,\n )\n topk_dists.append(dists)\n topk_ids.append(ids)\n\n return topk_dists, [np.array(ids, dtype=int) for ids in topk_ids]\n\n def insert(\n self,\n data: 'np.ndarray',\n cells: 'np.ndarray',\n docs: 'DocumentArray',\n only_index: bool = False,\n ):\n assert len(docs) == len(data)\n\n if self._projector_codec:\n data = self._projector_codec.encode(data)\n\n unique_cells, unique_cell_counts = np.unique(cells, return_counts=True)\n\n if len(unique_cells) == 1:\n cell_id = unique_cells[0]\n\n offsets = self.cell_table(cell_id).insert(docs)\n offsets = np.array(offsets, dtype=np.int64)\n\n self.vec_index(cell_id).add_with_ids(data, offsets)\n self._meta_table.bulk_add_address([d.id for d in docs], cells, offsets)\n\n if not only_index:\n self.doc_store(cell_id).insert(docs)\n\n else:\n for cell_id, cell_count in zip(unique_cells, unique_cell_counts):\n # TODO: Jina should allow boolean filtering in docarray to avoid this\n # and simply use cells == cell_index\n indices = np.where(cells == cell_id)[0]\n cell_docs = docs[indices.tolist()]\n\n cell_offsets = self.cell_table(cell_id).insert(cell_docs)\n cell_offsets = np.array(cell_offsets, dtype=np.int64)\n\n cell_data = data[indices, :]\n\n self.vec_index(cell_id).add_with_ids(cell_data, cell_offsets)\n self._meta_table.bulk_add_address(\n [d.id for d in cell_docs], [cell_id] * cell_count, cell_offsets\n )\n\n if not only_index:\n self.doc_store(cell_id).insert(cell_docs)\n\n logger.debug(f'{len(docs)} new docs added')\n\n def _add_vecs(self, data: 'np.ndarray', cells: 'np.ndarray', offsets: 'np.ndarray'):\n assert data.shape[0] == cells.shape[0]\n assert data.shape[1] == self.n_dim\n\n unique_cells, _ = np.unique(cells, return_counts=True)\n\n for cell_id in unique_cells:\n indices = cells == cell_id\n x = data[indices, :]\n ids = offsets[indices]\n\n self.vec_index(cell_id).add_with_ids(x, ids)\n\n def update(\n self,\n data: 'np.ndarray',\n cells: 'np.ndarray',\n docs: 'DocumentArray',\n insert_if_not_found: bool = True,\n raise_errors_on_not_found: bool = False,\n ):\n update_success = 0\n\n new_data = []\n new_cells = []\n new_docs = []\n\n for (\n x,\n doc,\n cell_id,\n ) in zip(data, docs, cells):\n _cell_id, _offset = self._meta_table.get_address(doc.id)\n if cell_id == _cell_id:\n self.vec_index(cell_id).add_with_ids(x.reshape(1, -1), [_offset])\n self.doc_store(cell_id).update([doc])\n self.meta_table.add_address(doc.id, cell_id, _offset)\n update_success += 1\n elif _cell_id is None:\n if raise_errors_on_not_found and not insert_if_not_found:\n raise Exception(\n f'The document (id={doc.id}) cannot be updated as'\n f'it is not found in the index'\n )\n elif not (raise_errors_on_not_found or insert_if_not_found):\n warnings.warn(\n f'The document (id={doc.id}) cannot be updated as '\n f'it is not found in the index',\n RuntimeWarning,\n )\n elif insert_if_not_found:\n new_data.append(x)\n new_cells.append(cell_id)\n new_docs.append(doc)\n update_success += 1\n else:\n continue\n else:\n # DELETE and INSERT\n self.vec_index(_cell_id).delete(_offset)\n self.cell_table(_cell_id).delete_by_offset(_offset)\n self.doc_store(_cell_id).delete([doc.id])\n\n new_data.append(x)\n new_cells.append(cell_id)\n new_docs.append(doc)\n update_success += 1\n\n if len(new_data) > 0:\n new_data = np.stack(new_data)\n new_cells = np.array(new_cells, dtype=np.int64)\n\n self.insert(new_data, new_cells, new_docs)\n\n logger.debug(\n f'total items for updating: {len(docs)}, ' f'success: {update_success}'\n )\n\n def delete(\n self,\n ids: List[str],\n raise_errors_on_not_found: bool = False,\n ):\n delete_success = 0\n\n for doc_id in ids:\n cell_id, offset = self._meta_table.get_address(doc_id)\n if cell_id is not None:\n self.vec_index(cell_id).delete([offset])\n self.cell_table(cell_id).delete_by_offset(offset)\n self.doc_store(cell_id).delete([doc_id])\n self.meta_table.delete_address(doc_id)\n delete_success += 1\n else:\n if raise_errors_on_not_found:\n raise Exception(\n f'The document (id={doc_id}) cannot be updated as'\n f'it is not found in the index'\n )\n else:\n continue\n\n logger.debug(\n f'total items for updating: {len(ids)}, ' f'success: {delete_success}'\n )\n\n def _rebuild_database(self):\n \"\"\"rebuild doc_store and meta_table after annlite download databse from hubble\"\"\"\n\n self._doc_stores = [\n DocStorage(\n self.data_path / f'cell_{_}',\n serialize_config=self.serialize_config or {},\n lock=True,\n )\n for _ in range(self.n_cells)\n ]\n # self._meta_table = MetaTable('metas', data_path=self.data_path, in_memory=False)\n\n def _get_doc_by_id(self, doc_id: str):\n cell_id = 0\n if self.n_cells > 1:\n cell_id, _ = self._meta_table.get_address(doc_id)\n\n da = self.doc_store(cell_id).get([doc_id])\n return da[0] if len(da) > 0 else None\n\n def documents_generator(self, cell_id: int, batch_size: int = 1000):\n for docs in self.doc_store(cell_id).batched_iterator(batch_size=batch_size):\n yield docs\n\n @property\n def cell_tables(self):\n return self._cell_tables\n\n @property\n def cell_indexes(self):\n return self._vec_indexes\n\n def cell_table(self, cell_id: int):\n return self._cell_tables[cell_id]\n\n def doc_store(self, cell_id: int):\n return self._doc_stores[cell_id]\n\n def vec_index(self, cell_id: int):\n return self._vec_indexes[cell_id]\n\n @property\n def meta_table(self):\n return self._meta_table\n\n @property\n def total_docs(self):\n return sum([store.size for store in self._doc_stores])\n\n @property\n def index_size(self):\n return sum([table.size for table in self._cell_tables])\n"
},
{
"path": "annlite/core/__init__.py",
"content": "from .codec import PQCodec, ProjectorCodec, VQCodec\n"
},
{
"path": "annlite/core/codec/__init__.py",
"content": "from .pq import PQCodec\nfrom .projector import ProjectorCodec\nfrom .vq import VQCodec\n"
},
{
"path": "annlite/core/codec/base.py",
"content": "import pickle\nfrom abc import ABC, abstractmethod\nfrom typing import TYPE_CHECKING\n\nif TYPE_CHECKING:\n from pathlib import Path\n\n\nclass BaseCodec(ABC):\n def __init__(self, require_train: bool = True):\n self.require_train = require_train\n self._is_trained = False if require_train else True\n\n @abstractmethod\n def fit(self, *args, **kwargs):\n pass\n\n @abstractmethod\n def encode(self):\n pass\n\n @abstractmethod\n def decode(self):\n pass\n\n def dump(self, target_path: 'Path'):\n pickle.dump(self, target_path.open('wb'), protocol=4)\n\n @staticmethod\n def load(from_path: 'Path'):\n return pickle.load(from_path.open('rb'))\n\n @property\n def is_trained(self):\n return self._is_trained\n\n def _check_trained(self):\n assert self.is_trained is True, f'{self.__class__.__name__} requires training'\n"
},
{
"path": "annlite/core/codec/pq.py",
"content": "from argparse import ArgumentError\n\nimport numpy as np\nfrom scipy.cluster.vq import vq\n\nfrom annlite import pq_bind\n\nfrom ...enums import Metric\nfrom ...math import l2_normalize\nfrom ...profile import time_profile\nfrom .base import BaseCodec\n\n# from pqlite.pq_bind import precompute_adc_table, dist_pqcodes_to_codebooks\n\n\nclass PQCodec(BaseCodec):\n \"\"\"Implementation of Product Quantization (PQ) [Jegou11]_.\n\n For the indexing phase of database vectors,\n a `D`-dim input vector is divided into `M` `D`/`M`-dim sub-vectors.\n Each sub-vector is quantized into a small integer via `Ks` codewords.\n\n For the querying phase, given a new `D`-dim query vector, the distance between the query\n and the database PQ-codes are efficiently approximated via Asymmetric Distance.\n All vectors must be np.ndarray with np.float32\n\n .. [Jegou11] H. Jegou et al., \"Product Quantization for Nearest Neighbor Search\", IEEE TPAMI 2011\n\n :param d_vector: the dimensionality of input vectors\n :param n_subvectors: The number of sub-space\n :param n_clusters: The number of codewords for each subspace\n (typically 256, so that each sub-vector is quantized\n into 256 bits pqlite.utils.asymmetric_distance= 1 byte = uint8)\n :param n_init: Number of times K-Means is trained with different centroid seeds. Best result of\n the `n_init` consecutive runs is selected.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n n_subvectors: int = 8,\n n_clusters: int = 256,\n metric: Metric = Metric.EUCLIDEAN,\n n_init: int = 4,\n ):\n super(PQCodec, self).__init__(require_train=True)\n self.dim = dim\n self.n_subvectors = n_subvectors\n self.n_clusters = n_clusters\n\n assert (\n dim % n_subvectors == 0\n ), 'input dimension must be dividable by number of sub-space'\n self.d_subvector = dim // n_subvectors\n\n self.code_dtype = (\n np.uint8\n if n_clusters <= 2**8\n else (np.uint16 if n_clusters <= 2**16 else np.uint32)\n )\n\n # assert (\n # metric == Metric.EUCLIDEAN\n # ), f'The distance metric `{metric.name}` is not supported yet!'\n self.metric = metric\n\n self.normalize_input = False\n if self.metric == Metric.COSINE:\n self.normalize_input = True\n\n self._codebooks = np.zeros(\n (self.n_subvectors, self.n_clusters, self.d_subvector), dtype=np.float32\n )\n self.kmeans = []\n self.n_init = n_init\n\n def __hash__(self):\n return hash(\n (\n self.__class__.__name__,\n self.dim,\n self.n_subvectors,\n self.n_clusters,\n self.metric,\n self.code_dtype,\n )\n )\n\n def fit(self, x: 'np.ndarray', iter: int = 100):\n \"\"\"Train the K-Means for each cartesian product\n\n :param x: Training vectors with shape=(N, D)\n :param iter: Number of iterations in Kmeans\n \"\"\"\n from sklearn.cluster import KMeans\n\n assert x.dtype == np.float32\n assert x.ndim == 2\n\n if self.normalize_input:\n x = l2_normalize(x)\n\n # [m][ks][ds]: m-th subspace, ks-the codeword, ds-th dim\n self._codebooks = np.zeros(\n (self.n_subvectors, self.n_clusters, self.d_subvector), dtype=np.float32\n )\n for m in range(self.n_subvectors):\n kmeans = KMeans(\n n_clusters=self.n_clusters, max_iter=iter, n_init=self.n_init\n )\n self.kmeans.append(kmeans)\n self.kmeans[m].fit(x[:, m * self.d_subvector : (m + 1) * self.d_subvector])\n self._codebooks[m] = self.kmeans[m].cluster_centers_\n\n self._is_trained = True\n\n def partial_fit(self, x: 'np.ndarray'):\n \"\"\"Given a batch of training vectors, update the internal MiniBatchKMeans.\n This method is specially designed to be used when data does not fit in memory.\n\n :param x: Training vectors with shape=(N, D)\n \"\"\"\n assert x.ndim == 2\n\n if self.normalize_input:\n x = l2_normalize(x)\n\n if len(self.kmeans) > 0:\n for m in range(self.n_subvectors):\n self.kmeans[m].partial_fit(\n x[:, m * self.d_subvector : (m + 1) * self.d_subvector]\n )\n else:\n from sklearn.cluster import MiniBatchKMeans\n\n for m in range(self.n_subvectors):\n self.kmeans.append(MiniBatchKMeans(n_clusters=self.n_clusters))\n\n for m in range(self.n_subvectors):\n self.kmeans[m].partial_fit(\n x[:, m * self.d_subvector : (m + 1) * self.d_subvector]\n )\n\n def build_codebook(self):\n \"\"\"Constructs sub-codebooks from the current parameters of the models in `self.kmeans`\n This step is not necessary if full KMeans is trained used calling `.fit`.\n \"\"\"\n\n self._codebooks = np.zeros(\n (self.n_subvectors, self.n_clusters, self.d_subvector), dtype=np.float32\n )\n\n for m in range(self.n_subvectors):\n self._codebooks[m] = self.kmeans[m].cluster_centers_\n\n self._is_trained = True\n\n def encode(self, x: 'np.ndarray'):\n \"\"\"Encode input vectors into PQ-codes.\n\n :param x: Input vectors with shape=(N, D) and dtype=np.float32.\n :return: np.ndarray: PQ codes with shape=(N, M) and dtype=self.code_dtype\n \"\"\"\n assert x.dtype == np.float32\n assert x.ndim == 2\n N, D = x.shape\n assert (\n D == self.d_subvector * self.n_subvectors\n ), 'input dimension must be Ds * M'\n\n # codes[n][m] : code of n-th vec, m-th subspace\n codes = np.empty((N, self.n_subvectors), dtype=self.code_dtype)\n for m in range(self.n_subvectors):\n sub_vecs = x[:, m * self.d_subvector : (m + 1) * self.d_subvector]\n codes[:, m], _ = vq(sub_vecs, self.codebooks[m])\n\n return codes\n\n def decode(self, codes: 'np.ndarray'):\n \"\"\"Given PQ-codes, reconstruct original D-dimensional vectors\n approximately by fetching the codewords.\n\n :param codes: PQ-cdoes with shape=(N, M) and dtype=self.code_dtype.\n Each row is a PQ-code\n :return: Reconstructed vectors with shape=(N, D) and dtype=np.float32\n \"\"\"\n assert codes.ndim == 2\n N, M = codes.shape\n assert M == self.n_subvectors\n assert codes.dtype == self.code_dtype\n\n vecs = np.empty((N, self.d_subvector * self.n_subvectors), dtype=np.float32)\n for m in range(self.n_subvectors):\n vecs[:, m * self.d_subvector : (m + 1) * self.d_subvector] = self.codebooks[\n m\n ][codes[:, m], :]\n\n return vecs\n\n def precompute_adc(self, query: object) -> object:\n \"\"\"Compute a distance table for a query vector.\n The distances are computed by comparing each sub-vector of the query\n to the codewords for each sub-subspace.\n `dtable[m][ks]` contains the squared Euclidean distance between\n the `m`-th sub-vector of the query and the `ks`-th codeword\n for the `m`-th sub-space (`self.codewords[m][ks]`).\n\n :param query: Input vector with shape=(D, ) and dtype=np.float32\n :return: Distance table. which contains dtable with shape=(M, Ks)\n and dtype=np.float32\n \"\"\"\n assert query.dtype == np.float32\n assert query.ndim == 1, 'input must be a single vector'\n\n # dtable[m] : distance between m-th subvec and m-th codewords (m-th subspace)\n # dtable[m][ks] : distance between m-th subvec and ks-th codeword of m-th codewords\n\n # Warning: the following line produces `ValueError: buffer source array is read-only`\n # if no `const` is used in the cython implementation using a memoryview\n dtable = pq_bind.precompute_adc_table(\n query, self.d_subvector, self.n_clusters, self.codebooks\n )\n\n return DistanceTable(dtable)\n\n @property\n def codebooks(self):\n return self._codebooks\n\n # trained pq interface ----------------\n def get_codebook(self) -> 'np.ndarray':\n \"\"\"Return the codebook parameters.\n\n Expect a 3-dimensional matrix is returned,\n with shape (`n_subvectors`, `n_clusters`, `d_subvector`) and dtype float32\n \"\"\"\n return np.ascontiguousarray(self.codebooks, dtype='float32')\n\n def get_subspace_splitting(self):\n \"\"\"Return subspace splitting setting\n\n :return: tuple of (`n_subvectors`, `n_clusters`, `d_subvector`)\n \"\"\"\n return (self.n_subvectors, self.n_clusters, self.d_subvector)\n\n # def get_dist_mat(self, x: np.ndarray):\n # \"\"\"Return the distance tables in form of matrix for multiple queries\n\n # :param query: shape('N', 'D'),\n\n # :return: ndarray with shape('N', `n_subvectors`, `n_clusters`)\n\n # .. note::\n # _description_\n # \"\"\"\n # assert x.dtype == np.float32\n # assert x.ndim == 2\n # N, D = x.shape\n # assert (\n # D == self.d_subvector * self.n_subvectors\n # ), 'input dimension must be Ds * M'\n # if self.normalize_input:\n # x = l2_normalize(x)\n\n # x = x.reshape(\n # N,\n # self.n_subvectors,\n # 1,\n # self.d_subvector,\n # )\n # if self.metric == Metric.EUCLIDEAN:\n # # (1, n_subvectors, n_clusters, d_subvector)\n # codebook = self.codebooks[np.newaxis, ...]\n\n # # broadcast to (N, n_subvectors, n_clusters, d_subvector)\n # dist_vector = (x - codebook) ** 2\n\n # # reduce to (N, n_subvectors, n_clusters)\n # dist_mat = np.sum(dist_vector, axis=3)\n # elif self.metric in [Metric.INNER_PRODUCT, Metric.COSINE]:\n # # (1, n_subvectors, n_clusters, d_subvector)\n # codebook = self.codebooks[np.newaxis, ...]\n\n # # broadcast to (N, n_subvectors, n_clusters, d_subvector)\n # dist_vector = x * codebook\n\n # # reduce to (N, n_subvectors, n_clusters)\n # dist_mat = 1 / self.n_clusters - np.sum(dist_vector, axis=3)\n # else:\n # raise ArgumentError(f'Unable support metrics {self.metric}')\n # return np.ascontiguousarray(dist_mat, dtype='float32')\n\n def get_dist_mat(self, x: np.ndarray):\n \"\"\"Return the distance tables in form of matrix for multiple queries\n\n :param query: shape('N', 'D'),\n\n :return: ndarray with shape('N', `n_subvectors`, `n_clusters`)\n\n .. note::\n _description_\n \"\"\"\n assert x.dtype == np.float32\n assert x.ndim == 2\n N, D = x.shape\n assert (\n D == self.d_subvector * self.n_subvectors\n ), 'input dimension must be Ds * M'\n if self.normalize_input:\n x = l2_normalize(x)\n\n if self.metric == Metric.EUCLIDEAN:\n dist_mat = pq_bind.batch_precompute_adc_table(\n x, self.d_subvector, self.n_clusters, self.codebooks\n )\n elif self.metric in [Metric.INNER_PRODUCT, Metric.COSINE]:\n dist_mat = 1 / self.n_clusters - np.array(\n pq_bind.batch_precompute_adc_table_ip(\n x, self.d_subvector, self.n_clusters, self.codebooks\n ),\n dtype='float32',\n )\n else:\n raise ArgumentError(f'Unable support metrics {self.metric}')\n return np.ascontiguousarray(dist_mat, dtype='float32')\n\n # -------------------------------------\n\n\nclass DistanceTable(object):\n \"\"\"Distance table from query to codeworkds.\n Given a query vector, a PQ/OPQ instance compute this DistanceTable class\n using :func:`PQ.dtable` or :func:`OPQ.dtable`.\n The Asymmetric Distance from query to each database codes can be computed\n by :func:`DistanceTable.adist`.\n Args:\n dtable (np.ndarray): Distance table with shape=(M, Ks) and dtype=np.float32\n computed by :func:`PQ.dtable` or :func:`OPQ.dtable`\n Attributes:\n dtable (np.ndarray): Distance table with shape=(M, Ks) and dtype=np.float32.\n Note that dtable[m][ks] contains the squared Euclidean distance between\n (1) m-th sub-vector of query and (2) ks-th codeword for m-th subspace.\n \"\"\"\n\n def __init__(self, dtable: 'np.ndarray'):\n\n assert dtable.ndim == 2\n self.dtable = dtable\n\n def adist(self, codes):\n \"\"\"Given PQ-codes, compute Asymmetric Distances between the query (self.dtable)\n and the PQ-codes.\n Args:\n codes (np.ndarray): PQ codes with shape=(N, M) and\n dtype=pq.code_dtype where pq is a pq instance that creates the codes\n Returns:\n np.ndarray: Asymmetric Distances with shape=(N, ) and dtype=np.float32\n \"\"\"\n\n assert codes.ndim == 2\n dists = pq_bind.dist_pqcodes_to_codebooks(self.dtable, codes)\n\n # The above line is equivalent to the followings:\n # dists = np.zeros((N, )).astype(np.float32)\n # for n in range(N):\n # for m in range(M):\n # dists[n] += self.dtable[m][codes[n][m]]\n return dists\n"
},
{
"path": "annlite/core/codec/projector.py",
"content": "from typing import Optional\n\nimport numpy as np\n\nfrom .base import BaseCodec\n\n\nclass ProjectorCodec(BaseCodec):\n \"\"\"Implementation of Projector.\n\n :param n_components: number of components to keep.\n :param whiten: when True (False by default) the components_ vectors are multiplied\n by the square root of n_samples and then divided by the singular\n values to ensure uncorrelated outputs with unit component-wise variances.\n :param svd_solver:\n If auto: The solver is selected by a default policy based on X.shape and\n n_components: if the input data is larger than 500x500 and the number of\n components to extract is lower than 80% of the smallest dimension of the\n data, then the more efficient โrandomizedโ method is enabled. Otherwise\n the exact full SVD is computed and optionally truncated afterwards.\n\n If full: run exact full SVD calling the standard LAPACK solver via scipy.\n linalg.svd and select the components by postprocessing.\n\n If arpack: run SVD truncated to n_components calling ARPACK solver via\n scipy.sparse.linalg.svds. It requires strictly 0 < n_components < min(X.shape).\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n n_components: int = 128,\n whiten: Optional[bool] = False,\n svd_solver: Optional[str] = 'auto',\n ):\n super(ProjectorCodec, self).__init__(require_train=True)\n self.dim = dim\n self.n_components = n_components\n assert self.dim >= self.n_components, (\n f'the dimension after projector should be less than original dimension, got '\n f'original dimension: {self.dim} and projector dimension: {self.n_components}'\n )\n\n self.whiten = whiten\n self.svd_solver = svd_solver\n\n self.pca = None\n\n def __hash__(self):\n return hash(\n (\n self.__class__.__name__,\n self.dim,\n self.n_components,\n self.whiten,\n self.svd_solver,\n )\n )\n\n def fit(self, x: 'np.ndarray'):\n \"\"\"Train projector model\n\n :param x: Training vectors with shape=(N, D)\n \"\"\"\n assert x.ndim == 2\n assert (\n x.shape[1] == self.dim,\n ), 'dimension of input data must be equal to \"dim\"'\n assert (\n x.shape[0] > self.n_components\n ), 'number of input data must be larger than or equal to n_components'\n\n if self.pca is None:\n from sklearn.decomposition import PCA\n\n self.pca = PCA(\n n_components=self.n_components,\n whiten=self.whiten,\n svd_solver=self.svd_solver,\n )\n\n self.pca.fit(x)\n self._is_trained = True\n\n def partial_fit(self, x: 'np.ndarray'):\n \"\"\"Given a batch of training vectors, update the internal projector.\n This method is specially designed to be used when data does not fit in memory.\n\n :param x: Training vectors with shape=(N, D)\n \"\"\"\n\n assert x.ndim == 2\n assert x.shape[1] == self.dim, 'dimension of input data must be equal to \"dim\"'\n assert (\n x.shape[0] > self.n_components\n ), 'number of input data must be larger than or equal to n_components'\n\n if self.pca is None:\n from sklearn.decomposition import IncrementalPCA\n\n self.pca = IncrementalPCA(\n n_components=self.n_components,\n whiten=self.whiten,\n )\n\n self.pca.partial_fit(x)\n self._is_trained = True\n\n def encode(self, x: 'np.ndarray'):\n \"\"\"Encode input vectors using projector.\n\n :param x: Input vectors with shape=(N, D)\n :return: np.ndarray: transformed vectors using projector.\n \"\"\"\n assert x.ndim == 2\n assert x.shape[1] == self.dim, 'dimension of input data must be equal to \"dim\"'\n\n return self.pca.transform(x)\n\n def decode(self, x: 'np.ndarray'):\n \"\"\"Given transformed vectors, reconstruct original D-dimensional vectors\n approximately.\n\n :param x: vectors with shape=(N, self.n_components).\n :return: Reconstructed vectors with shape=(N, D)\n \"\"\"\n assert x.ndim == 2\n assert x.shape[1] == self.n_components\n\n return self.pca.inverse_transform(x)\n\n @property\n def components(self):\n \"\"\"Principal axes in feature space, representing the directions of maximum\n variance in the data.\n \"\"\"\n self._check_trained()\n return self.pca.components_\n\n @property\n def explained_variance_ratio(self):\n \"\"\"Percentage of variance explained by each of the selected components.\"\"\"\n self._check_trained()\n return self.pca.explained_variance_ratio_\n\n @property\n def mean(self):\n \"\"\"Per-feature empirical mean.\"\"\"\n self._check_trained()\n return self.pca.mean_\n\n @property\n def var(self):\n \"\"\"Per-feature empirical variance\"\"\"\n self._check_trained()\n return self.pca.var_\n"
},
{
"path": "annlite/core/codec/vq.py",
"content": "import numpy as np\nfrom scipy.cluster.vq import vq\n\nfrom ...enums import Metric\nfrom .base import BaseCodec\n\n\nclass VQCodec(BaseCodec):\n def __init__(\n self,\n n_clusters: int,\n metric: Metric = Metric.EUCLIDEAN,\n iter: int = 100,\n n_init: int = 4,\n *args,\n **kwargs\n ):\n super(VQCodec, self).__init__(require_train=True)\n self.n_clusters = n_clusters\n\n # assert (\n # metric == Metric.EUCLIDEAN\n # ), f'The distance metric `{metric.name}` is not supported yet!'\n self.metric = metric\n self._codebook = None\n self.iter = iter\n self.kmeans = None\n self.n_init = n_init\n\n def __hash__(self):\n return hash((self.__class__.__name__, self.n_clusters, self.metric))\n\n def fit(self, x: 'np.ndarray'):\n \"\"\"Given training vectors, run k-means for each sub-space and create\n codewords for each sub-space.\n\n :param x: Training vectors with shape=(N, D) and dtype=np.float32.\n :param iter: The number of iteration for k-means\n \"\"\"\n\n from sklearn.cluster import KMeans\n\n assert x.dtype == np.float32\n assert x.ndim == 2\n\n self.kmeans = KMeans(self.n_clusters, max_iter=self.iter, n_init=self.n_init)\n self.kmeans.fit(x)\n self._codebook = self.kmeans.cluster_centers_\n self._is_trained = True\n\n def partial_fit(self, x: 'np.ndarray'):\n \"\"\"Given a batch of training vectors, update the internal MiniBatchKMeans.\n This method is specially designed to be used when data does not fit in memory.\n\n :param x: Training vectors with shape=(N, D)\n \"\"\"\n assert x.ndim == 2\n if self.kmeans:\n self.kmeans.partial_fit(x)\n else:\n from sklearn.cluster import MiniBatchKMeans\n\n self.kmeans = MiniBatchKMeans(\n n_clusters=self.n_clusters, max_iter=self.iter\n )\n self.kmeans.partial_fit(x)\n\n def build_codebook(self):\n \"\"\"Constructs a codebook from the current MiniBatchKmeans\n This step is not necessary if full KMeans is trained used calling `.fit`.\n \"\"\"\n self._codebook = self.kmeans.cluster_centers_\n self._is_trained = True\n\n def encode(self, x: 'np.ndarray'):\n \"\"\"Encodes each row of the input array `x` it's closest cluster id.\"\"\"\n self._check_trained()\n assert x.dtype == np.float32\n assert x.ndim == 2\n\n codes, _ = vq(x, self.codebook)\n return codes\n\n def decode(self, x: 'np.ndarray'):\n return None\n\n @property\n def codebook(self):\n self._check_trained()\n return self._codebook\n"
},
{
"path": "annlite/core/index/__init__.py",
"content": ""
},
{
"path": "annlite/core/index/base.py",
"content": "import abc\nfrom typing import List, Optional, Union\n\nimport numpy as np\n\nfrom ...enums import ExpandMode, Metric\nfrom ...helper import str2dtype\n\n\nclass BaseIndex(abc.ABC):\n def __init__(\n self,\n dim: int,\n dtype: Union[np.dtype, str] = np.float32,\n metric: Metric = Metric.COSINE,\n initial_size: Optional[int] = None,\n expand_step_size: int = 10240,\n expand_mode: ExpandMode = ExpandMode.STEP,\n *args,\n **kwargs\n ):\n assert expand_step_size > 0\n self.initial_size = initial_size or expand_step_size\n\n self.expand_step_size = expand_step_size\n self.expand_mode = expand_mode\n\n self.dim = dim\n self.dtype = str2dtype(dtype) if isinstance(dtype, str) else dtype\n self.metric = metric\n\n self._size = 0\n self._capacity = self.initial_size\n\n @property\n def capacity(self) -> int:\n return self._capacity\n\n @property\n def size(self):\n return self._size\n\n @abc.abstractmethod\n def add_with_ids(self, x: np.ndarray, ids: List[int], **kwargs):\n ...\n\n @abc.abstractmethod\n def delete(self, ids: List[int]):\n ...\n\n @abc.abstractmethod\n def update_with_ids(self, x: np.ndarray, ids: List[int], **kwargs):\n ...\n\n def reset(self, capacity: Optional[int] = None):\n self._size = 0\n self._capacity = capacity or self.initial_size\n"
},
{
"path": "annlite/core/index/flat_index.py",
"content": "from typing import List, Optional\n\nimport numpy as np\nfrom loguru import logger\n\nfrom ...math import cdist, top_k\nfrom .base import BaseIndex\n\n\nclass FlatIndex(BaseIndex):\n def __init__(self, *args, **kwargs):\n super(FlatIndex, self).__init__(*args, **kwargs)\n self._data = np.zeros((self.initial_size, self.dim), dtype=self.dtype)\n\n def search(\n self, x: np.ndarray, limit: int = 10, indices: Optional[np.ndarray] = None\n ):\n _dim = x.shape[-1]\n assert (\n _dim == self.dim\n ), f'the query embedding dimension does not match with index dimension: {_dim} vs {self.dim}'\n\n x = x.reshape((-1, self.dim))\n\n data = self._data[: self.size]\n data_ids = np.arange(self.size)\n\n if indices is not None:\n data = self._data[indices]\n data_ids = data_ids[indices]\n\n dists = cdist(x, data, metric=self.metric.name.lower())\n dists, idx = top_k(dists, limit, descending=False)\n\n # TODO: change the shape of return\n dists = dists[0]\n data_ids = data_ids[idx[0]]\n\n return dists, data_ids\n\n def add_with_ids(self, x: np.ndarray, ids: List[int]):\n for idx in ids:\n if idx >= self._capacity:\n self._expand_capacity()\n\n start = self._size\n end = start + len(x)\n\n self._data[ids, :] = x\n self._size = end\n\n def _expand_capacity(self):\n new_block = np.zeros((self.expand_step_size, self.dim), dtype=self.dtype)\n self._data = np.concatenate((self._data, new_block), axis=0)\n\n self._capacity += self.expand_step_size\n logger.debug(\n f'total storage capacity is expanded by {self.expand_step_size}',\n )\n\n def reset(self, capacity: Optional[int] = None):\n super().reset(capacity=capacity)\n self._data = np.zeros((self.capacity, self.dim), dtype=self.dtype)\n\n def delete(self, ids: List[int]):\n raise RuntimeError(\n f'the deletion operation is not allowed for {self.__class__.__name__}!'\n )\n\n def update_with_ids(self, x: np.ndarray, ids: List[int], **kwargs):\n self._data[ids, :] = x\n"
},
{
"path": "annlite/core/index/hnsw/__init__.py",
"content": "from .index import HnswIndex\n"
},
{
"path": "annlite/core/index/hnsw/index.py",
"content": "import math\nimport os.path\nfrom functools import wraps\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, List, Optional, Union\n\nimport numpy as np\nfrom loguru import logger\n\nfrom annlite.hnsw_bind import Index\n\nfrom ....enums import Metric\nfrom ....math import l2_normalize\nfrom ..base import BaseIndex\n\nif TYPE_CHECKING:\n from ...codec.base import BaseCodec\n\n\ndef pre_process(f):\n @wraps(f)\n def pre_processed(self: 'HnswIndex', x: np.ndarray, *args, **kwargs):\n if x.ndim == 1:\n x = x.reshape((1, -1))\n if x.dtype != self.dtype:\n x = x.astype(self.dtype)\n\n if self.normalization_enable:\n x = l2_normalize(x)\n\n if self.pq_enable:\n if not self.pq_codec.is_trained:\n raise RuntimeError(\n 'Please train the PQ before using HNSW quantization backend'\n )\n elif not self._set_backend_pq:\n self._index.loadPQ(self.pq_codec)\n self._set_backend_pq = True\n kwargs['pre_process_dtables'] = self.pq_codec.get_dist_mat(x)\n x = self.pq_codec.encode(x)\n\n assert kwargs['pre_process_dtables'].dtype == 'float32'\n assert kwargs['pre_process_dtables'].flags['C_CONTIGUOUS']\n return f(self, x, *args, **kwargs)\n else:\n return f(self, x, *args, **kwargs)\n\n return pre_processed\n\n\nclass HnswIndex(BaseIndex):\n def __init__(\n self,\n dim: int,\n dtype: np.dtype = np.float32,\n metric: Metric = Metric.COSINE,\n ef_construction: int = 200,\n ef_search: int = 50,\n max_connection: int = 16,\n pq_codec: Optional['BaseCodec'] = None,\n index_file: Optional[Union[str, Path]] = None,\n **kwargs,\n ):\n \"\"\"\n :param dim: The dimensionality of vectors to index\n :param index_file: A file-like object or a string containing a file name.\n :param metric: Distance metric type, can be 'euclidean', 'inner_product', or 'cosine'\n :param ef_construction: the size of the dynamic list for the nearest neighbors (used during the building).\n :param ef_search: the size of the dynamic list for the nearest neighbors (used during the search).\n :param max_connection: The number of bi-directional links created for every new element during construction.\n Reasonable range for M is 2-100.\n \"\"\"\n super().__init__(dim, dtype=dtype, metric=metric, **kwargs)\n\n self.ef_construction = ef_construction\n self.ef_search = ef_search\n self.max_connection = max_connection\n self.pq_codec = pq_codec\n self._set_backend_pq = False\n self.index_file = index_file\n\n self._init_hnsw_index()\n\n def _init_hnsw_index(self):\n self._index = Index(space=self.space_name, dim=self.dim)\n if self.index_file:\n if os.path.exists(self.index_file):\n logger.info(\n f'indexer will be loaded from {self.index_file}',\n )\n self.load(self.index_file)\n else:\n raise FileNotFoundError(\n f'index path: {self.index_file} does not exist',\n )\n else:\n if self.pq_codec is not None and self.pq_codec.is_trained:\n self._index.init_index(\n max_elements=self.capacity,\n ef_construction=self.ef_construction,\n M=self.max_connection,\n pq_codec=self.pq_codec,\n )\n self._set_backend_pq = True\n else:\n self._index.init_index(\n max_elements=self.capacity,\n ef_construction=self.ef_construction,\n M=self.max_connection,\n pq_codec=None,\n )\n self._set_backend_pq = False\n\n self._index.set_ef(self.ef_search)\n\n def load(self, index_file: Union[str, Path]):\n self._index.load_index(str(index_file))\n if self.pq_codec:\n self._index.loadPQ(self.pq_codec)\n\n def dump(self, index_file: Union[str, Path]):\n self._index.save_index(str(index_file))\n\n @pre_process\n def add_with_ids(\n self,\n x: 'np.ndarray',\n ids: List[int],\n # kwargs maybe used by pre_process\n pre_process_dtables=None,\n ):\n max_id = max(ids) + 1\n if max_id > self.capacity:\n expand_steps = math.ceil(max_id / self.expand_step_size)\n self._expand_capacity(expand_steps * self.expand_step_size)\n\n self._index.add_items(x, ids=ids, dtables=pre_process_dtables)\n\n @pre_process\n def search(\n self,\n query: 'np.ndarray',\n limit: int = 10,\n indices: Optional['np.ndarray'] = None,\n # kwargs maybe used by pre_process\n pre_process_dtables=None,\n ):\n ef_search = max(self.ef_search, limit)\n self._index.set_ef(ef_search)\n\n if indices is not None:\n # TODO: add a smart strategy to speed up this case (bruteforce search would be better)\n if len(indices) < limit:\n limit = len(indices)\n ids, dists = self._index.knn_query_with_filter(\n query, filters=indices, k=limit, dtables=pre_process_dtables\n )\n else:\n ids, dists = self._index.knn_query(\n query, k=limit, dtables=pre_process_dtables\n )\n\n # convert squared l2 into euclidean distance\n if self.metric == Metric.EUCLIDEAN:\n dists = np.sqrt(dists)\n\n return dists[0], ids[0]\n\n def delete(self, ids: List[int]):\n for i in ids:\n self._index.mark_deleted(i)\n\n def update_with_ids(self, x: 'np.ndarray', ids: List[int], **kwargs):\n raise RuntimeError(\n f'the update operation is not allowed for {self.__class__.__name__}!'\n )\n\n def _expand_capacity(self, new_capacity: int):\n self._capacity = new_capacity\n self._index.resize_index(new_capacity)\n logger.debug(\n f'HNSW index capacity is expanded by {self.expand_step_size}',\n )\n\n def reset(self, capacity: Optional[int] = None):\n super().reset(capacity=capacity)\n self._init_hnsw_index()\n\n @property\n def size(self):\n return self._index.element_count\n\n @property\n def space_name(self):\n if self.metric == Metric.EUCLIDEAN:\n return 'l2'\n elif self.metric == Metric.INNER_PRODUCT:\n return 'ip'\n return 'cosine'\n\n @property\n def pq_enable(self):\n return self.pq_codec is not None\n\n @property\n def normalization_enable(self):\n return self.metric == Metric.COSINE\n"
},
{
"path": "annlite/core/index/pq_index.py",
"content": "from typing import List, Optional\n\nimport numpy as np\n\nfrom ...math import top_k\nfrom ..codec.pq import PQCodec\nfrom .flat_index import FlatIndex\n\n\n# TODO: deprecated this index\nclass PQIndex(FlatIndex): # pragma: no cover\n def __init__(\n self,\n dim: int,\n pq_codec: PQCodec,\n **kwargs,\n ):\n assert pq_codec is not None\n self._dense_dim = dim\n super(PQIndex, self).__init__(\n pq_codec.n_subvectors, dtype=pq_codec.code_dtype, **kwargs\n )\n self._pq_codec = pq_codec\n\n def add_with_ids(self, x: np.ndarray, ids: List[int]):\n x = self._pq_codec.encode(x)\n super(PQIndex, self).add_with_ids(x, ids)\n\n def search(\n self, x: np.ndarray, limit: int = 10, indices: Optional[np.ndarray] = None\n ):\n _dim = x.shape[-1]\n assert (\n _dim == self._pq_codec.dim\n ), f'the query embedding dimension does not match with index dimension: {_dim} vs {self.dim}'\n\n precomputed = self._pq_codec.precompute_adc(x)\n\n codes = self._data\n data_idx = np.arange(self._capacity)\n\n if indices is not None:\n codes = self._data[indices]\n data_idx = data_idx[indices]\n\n dists = precomputed.adist(codes) # (10000, )\n dists = np.expand_dims(dists, axis=0)\n\n dists, ids = top_k(dists, limit, descending=False)\n\n # TODO: change the shape of return\n ids = ids[0]\n if indices is not None:\n ids = data_idx[ids]\n\n return dists[0], ids\n"
},
{
"path": "annlite/enums.py",
"content": "from enum import IntEnum\n\n\nclass BetterEnum(IntEnum):\n \"\"\"The base class of Enum.\"\"\"\n\n def __str__(self):\n return self.name\n\n @classmethod\n def from_string(cls, s: str):\n \"\"\"\n Parse the enum from a string.\n :param s: string representation of the enum value\n :return: enum value\n \"\"\"\n try:\n return cls[s.upper()]\n except KeyError:\n raise ValueError(\n f'{s.upper()} is not a valid enum for {cls!r}, must be one of {list(cls)}'\n )\n\n\nclass Metric(BetterEnum):\n EUCLIDEAN = 1\n INNER_PRODUCT = 2\n COSINE = 3\n\n\nclass ExpandMode(BetterEnum):\n STEP = 1\n DOUBLE = 2\n ADAPTIVE = 3\n"
},
{
"path": "annlite/executor.py",
"content": "import threading\nimport time\nimport traceback\nimport warnings\nfrom threading import Thread\nfrom typing import Dict, List, Optional, Tuple, Union\n\nfrom docarray import Document, DocumentArray\nfrom jina import Executor, requests\nfrom jina.logging.logger import JinaLogger\n\nINDEX_BATCH_SIZE = 1024\n\n\nclass AnnLiteIndexer(Executor):\n \"\"\"A simple indexer that wraps the AnnLite indexer and adds a simple interface for indexing and searching.\n\n :param n_dim: Dimensionality of vectors to index\n :param metric: Distance metric type. Can be 'euclidean', 'inner_product', or 'cosine'\n :param limit: Number of results to get for each query document in search\n :param n_components: Number of components to use for dimensionality reduction\n :param match_args: the arguments to `DocumentArray`'s match function\n :param data_path: the workspace of the AnnLiteIndexer but not support when shards > 1.\n :param ef_construction: The construction time/accuracy trade-off\n :param ef_search: The query time accuracy/speed trade-off\n :param max_connection: The maximum number of outgoing connections in the\n graph (the \"M\" parameter)\n :param include_metadata: If True, return the document metadata in response\n :param index_access_paths: Default traversal paths on docs\n (used for indexing, delete and update), e.g. '@r', '@c', '@r,c'\n :param search_access_paths: Default traversal paths on docs\n (used for search), e.g. '@r', '@c', '@r,c'\n :param columns: A list or dict of column names to index.\n :param dim: Deprecated, use n_dim instead\n \"\"\"\n\n def __init__(\n self,\n n_dim: int = 0,\n metric: str = 'cosine',\n limit: int = 10,\n n_components: Optional[int] = None,\n match_args: Optional[Dict] = None,\n data_path: Optional[str] = None,\n ef_construction: Optional[int] = None,\n ef_search: Optional[int] = None,\n max_connection: Optional[int] = None,\n include_metadata: bool = True,\n index_access_paths: str = '@r',\n search_access_paths: str = '@r',\n columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None,\n list_like: Optional[bool] = False,\n dim: int = None,\n *args,\n **kwargs,\n ):\n\n super().__init__(*args, **kwargs)\n self.logger = JinaLogger(self.__class__.__name__)\n\n n_dim = n_dim or dim\n if not n_dim:\n raise ValueError('Please specify the dimension of the vectors to index!')\n\n self.n_components = n_components\n self.metric = metric\n self.match_args = match_args or {}\n self.include_metadata = include_metadata\n if limit:\n self.match_args.update({'limit': limit})\n\n self.index_access_paths = index_access_paths\n if 'index_traversal_paths' in kwargs:\n warnings.warn(\n f'`index_traversal_paths` is deprecated. Use `index_access_paths` instead.'\n )\n self.index_access_paths = kwargs['index_traversal_paths']\n\n self.search_access_paths = search_access_paths\n if 'search_traversal_paths' in kwargs:\n warnings.warn(\n f'`search_traversal_paths` is deprecated. Use `search_access_paths` instead.'\n )\n self.search_access_paths = kwargs['search_traversal_paths']\n\n self._data_buffer = DocumentArray()\n self._index_batch_size = INDEX_BATCH_SIZE\n self._max_length_queue = 2 * self._index_batch_size\n self._index_lock = threading.Lock()\n\n self.logger = JinaLogger(getattr(self.metas, 'name', self.__class__.__name__))\n\n if getattr(self.runtime_args, 'shards', 1) > 1 and data_path:\n raise ValueError(\n '`data_path` is not supported when shards > 1, please use `workspace` instead'\n )\n\n config = {\n 'n_dim': n_dim,\n 'n_components': n_components,\n 'metric': metric,\n 'ef_construction': ef_construction,\n 'ef_search': ef_search,\n 'max_connection': max_connection,\n 'data_path': data_path or self.workspace or './workspace',\n 'columns': columns,\n 'list_like': list_like,\n }\n self._index = DocumentArray(storage='annlite', config=config)\n\n # start indexing thread in background to group indexing requests\n # together and perform batch indexing at once\n self._start_index_loop()\n\n @requests(on='/index')\n def index(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Index new documents\n\n :param docs: the Documents to index\n :param parameters: dictionary with options for indexing\n Keys accepted:\n - 'access_paths': traversal paths on docs, e.g. '@r', '@c', '@r,c'\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.index_access_paths)\n flat_docs = docs[access_paths]\n if len(flat_docs) == 0:\n return\n\n while len(self._data_buffer) >= self._max_length_queue:\n time.sleep(0.001)\n\n with self._index_lock:\n self._data_buffer.extend(flat_docs)\n\n def _start_index_loop(self):\n \"\"\"Start the indexing loop in background.\n\n This loop is responsible for batch indexing the documents in the buffer.\n \"\"\"\n\n def _index_loop():\n try:\n while True:\n # if the buffer is none, will break the loop\n if self._data_buffer is None:\n break\n\n # if the buffer is empty, will wait for new documents to be added\n if len(self._data_buffer) == 0:\n time.sleep(0.1) # sleep for 100ms\n continue\n\n # acquire the lock to prevent threading issues\n with self._index_lock:\n batch_docs = self._data_buffer.pop(\n range(\n self._index_batch_size\n if len(self._data_buffer) > self._index_batch_size\n else len(self._data_buffer)\n )\n )\n self._index.extend(batch_docs)\n self.logger.debug(f'indexing {len(batch_docs)} docs done...')\n except Exception as e:\n self.logger.error(traceback.format_exc())\n raise e\n\n self._index_thread = Thread(target=_index_loop, daemon=False)\n self._index_thread.start()\n\n @requests(on='/update')\n def update(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Update existing documents\n\n :param docs: the Documents to update\n :param parameters: dictionary with options for updating\n Keys accepted:\n - 'access_paths': traversal paths on docs, e.g. '@r', '@c', '@r,c'\n - 'raise_errors_on_not_found': if True, raise an error if a document is not found. Default is False.\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.index_access_paths)\n raise_errors_on_not_found = parameters.get('raise_errors_on_not_found', False)\n flat_docs = docs[access_paths]\n if len(flat_docs) == 0:\n return\n\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot update documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n for doc in flat_docs:\n try:\n self._index[doc.id] = doc\n except IndexError:\n if raise_errors_on_not_found:\n raise Exception(\n f'The document (id={doc.id}) cannot be updated as'\n f'it is not found in the index'\n )\n else:\n self.logger.warning(\n f'cannot update doc {doc.id} as it does not exist in storage'\n )\n\n @requests(on='/delete')\n def delete(self, parameters: dict = {}, **kwargs):\n \"\"\"Delete existing documents\n\n Delete entries from the index by id\n :param parameters: parameters to the request\n \"\"\"\n\n delete_ids = parameters.get('ids', [])\n if len(delete_ids) == 0:\n return\n\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot delete documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n\n del self._index[delete_ids]\n\n @requests(on='/search')\n def search(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Perform a vector similarity search and retrieve Document matches\n\n Search can be performed with candidate filtering. Filters are a triplet (column,operator,value).\n More than a filter can be applied during search. Therefore, conditions for a filter are specified as a list triplets.\n Each triplet contains:\n\n - column: Column used to filter.\n - operator: Binary operation between two values. Some supported operators include `['>','<','=','<=','>=']`.\n - value: value used to compare a candidate.\n\n :param docs: the Documents to search with\n :param parameters: dictionary for parameters for the search operation\n Keys accepted:\n - 'access_paths' (str): traversal paths on docs, e.g. '@r', '@c', '@r,c'\n - 'filter' (dict): the filtering conditions on document tags\n - 'limit' (int): nr of matches to get per Document\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.search_access_paths)\n flat_docs = docs[access_paths]\n match_args = (\n {**self.match_args, **parameters}\n if parameters is not None\n else self.match_args\n )\n\n with self._index_lock:\n # if len(self._data_buffer) > 0:\n # raise RuntimeError(\n # f'Cannot search documents while the pending documents in the buffer are not indexed yet. '\n # 'Please wait for the pending documents to be indexed.'\n # )\n\n flat_docs.match(self._index, **match_args)\n\n @requests(on='/backup')\n def backup(self, parameters: Optional[Dict] = {}, **kwargs):\n \"\"\"\n Backup data to local or remote.\n Use api of \n\n Keys accepted:\n - 'target' (str): the name of indexer you want to backup as\n \"\"\"\n\n target_name = parameters.get('target_name', None)\n token = parameters.get('token', None)\n if target_name:\n target_name = f'{target_name}_{self.runtime_args.shard_id}'\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot backup documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n self._index._annlite.backup(target_name, token)\n if self._index._list_like:\n self._index._save_offset2ids()\n\n @requests(on='/restore')\n def restore(self, parameters: Optional[Dict] = {}, **kwargs):\n \"\"\"\n Restore data from local or remote.\n Use api of \n \"\"\"\n source_name = parameters.get('source_name', None)\n token = parameters.get('token', None)\n if source_name:\n source_name = f'{source_name}_{self.runtime_args.shard_id}'\n self._index._annlite.restore(source_name, token)\n if self._index._list_like:\n self._index._load_offset2ids()\n\n @requests(on='/filter')\n def filter(self, parameters: Dict, **kwargs):\n \"\"\"\n Query documents from the indexer by the filter `query` object in parameters. The `query` object must follow the\n specifications in the `find` method of `DocumentArray` using annlite: https://docarray.jina.ai/fundamentals/documentarray/find/#filter-with-query-operators\n :param parameters: Dictionary to define the `filter` that you want to use.\n \"\"\"\n\n return self._index.find(parameters.get('filter', None))\n\n @requests(on='/fill_embedding')\n def fill_embedding(self, docs: DocumentArray, **kwargs):\n \"\"\"\n retrieve embedding of Documents by id\n :param docs: DocumentArray to search with\n \"\"\"\n\n for doc in docs:\n doc.embedding = self._index[doc.id].embedding\n\n @requests(on='/status')\n def status(self, **kwargs) -> DocumentArray:\n \"\"\"Return the document containing status information about the indexer.\n\n The status will contain information on the total number of indexed and deleted\n documents, and on the number of (searchable) documents currently in the index.\n \"\"\"\n\n status = Document(\n tags={\n 'appending_size': len(self._data_buffer),\n 'total_docs': len(self._index),\n 'index_size': len(self._index),\n }\n )\n return DocumentArray([status])\n\n def flush(self):\n \"\"\"Flush all the data in the buffer to the index\"\"\"\n while len(self._data_buffer) > 0:\n time.sleep(0.1)\n\n @requests(on='/clear')\n def clear(self, **kwargs):\n \"\"\"Clear the index of all entries.\"\"\"\n self.flush()\n\n with self._index_lock:\n self._data_buffer = None\n self._index_thread.join()\n\n self._data_buffer = DocumentArray()\n self._index.clear()\n\n self._start_index_loop()\n\n def close(self, **kwargs):\n \"\"\"Close the index.\"\"\"\n super().close()\n\n self.flush()\n\n # wait for the index thread to finish\n with self._index_lock:\n self._data_buffer = None\n self._index_thread.join()\n\n # WARNING: the commented code below hangs the close in pytest `pytest tests/test_*.py`\n # But don't know why. It works fine in `pytest tests/test_executor.py` and normal python execution\n del self._index\n"
},
{
"path": "annlite/filter.py",
"content": "from typing import Dict\n\nLOGICAL_OPERATORS = {'$and': 'AND', '$or': 'OR'}\n\nCOMPARISON_OPERATORS = {\n '$lt': '<',\n '$gt': '>',\n '$lte': '<=',\n '$gte': '>=',\n '$eq': '=',\n '$neq': '!=',\n}\n\nMEMBERSHIP_OPERATORS = {'$in': 'IN', '$nin': 'NOT IN'}\n\n\ndef _sql_parsing(data, default_logic: str = 'AND'):\n \"\"\"\n :param data: JSON Object (dict).\n :param parameters: dict.\n :return: where clause (str) built from data\n \"\"\"\n where_clause = ''\n parameters = []\n\n if isinstance(data, dict):\n for i, (key, value) in enumerate(data.items()):\n if key in LOGICAL_OPERATORS:\n clause, params = _sql_parsing(\n value, default_logic=LOGICAL_OPERATORS[key]\n )\n if i == 0:\n where_clause += clause\n else:\n where_clause += f' {LOGICAL_OPERATORS[key]} {clause}'\n parameters.extend(params)\n elif key.startswith('$'):\n raise ValueError(\n f'The operator {key} is not supported yet, please double check the given filters!'\n )\n else:\n if i > 0:\n where_clause += f' {default_logic} '\n\n items = list(value.items())\n\n if len(items) == 0:\n raise ValueError(f'The query express is illegal: {data}')\n elif len(items) > 1:\n clause_list, params_list = [], []\n\n for op, val in items:\n _clause, _params = _sql_parsing({key: {op: val}})\n clause_list.append(_clause)\n params_list.extend(_params)\n\n where_clause += f' AND '.join(clause_list)\n parameters.extend(params_list)\n else:\n op, val = items[0]\n if op in LOGICAL_OPERATORS:\n clause, params = _sql_parsing(\n val, default_logic=LOGICAL_OPERATORS[op]\n )\n where_clause += clause\n parameters.extend(params)\n elif op in COMPARISON_OPERATORS:\n parameters.append(val)\n where_clause += f'({key} {COMPARISON_OPERATORS[op]} ?)'\n elif op in MEMBERSHIP_OPERATORS:\n parameters.extend(val)\n where_clause += f'({key} {MEMBERSHIP_OPERATORS[op]}({\", \".join([\"?\"]*len(val))}))'\n else:\n raise ValueError(\n f'The operator {op} is not supported yet, please double check the given filters!'\n )\n elif isinstance(data, list):\n clause_list, params_list = [], []\n for d in data:\n _clause, _params = _sql_parsing(d)\n clause_list.append(_clause)\n params_list.extend(_params)\n where_clause += '(' + f' {default_logic} '.join(clause_list) + ')'\n parameters.extend(params_list)\n\n elif isinstance(data, str):\n return data, parameters\n else:\n raise ValueError(f'The query express is illegal: {data}')\n return where_clause, tuple(parameters)\n\n\nclass Filter(object):\n \"\"\"A class to parse query language to SQL where clause.\"\"\"\n\n def __init__(self, tree_data: Dict = {}):\n self.tree_data = tree_data\n\n def parse_where_clause(self):\n return _sql_parsing(self.tree_data or {})\n"
},
{
"path": "annlite/helper.py",
"content": "import sys\n\nimport numpy as np\nfrom loguru import logger\n\n\ndef setup_logging(debug: bool):\n \"\"\"\n Setup the log formatter for AnnLite.\n \"\"\"\n\n log_level = 'INFO'\n if debug:\n log_level = 'DEBUG'\n\n logger.remove()\n logger.add(\n sys.stdout,\n colorize=True,\n level=log_level,\n )\n\n\ndef str2dtype(dtype_str: str):\n if dtype_str in ['double', 'float64']:\n dtype = np.float64\n elif dtype_str in ['half', 'float16']:\n dtype = np.float16\n elif dtype_str in ['float', 'float32']:\n dtype = np.float32\n elif dtype_str in ['bfloat16']:\n dtype = np.bfloat16\n elif dtype_str in ['long', 'int64']:\n dtype = np.int64\n elif dtype_str in ['int', 'int32']:\n dtype = np.int32\n elif dtype_str in ['int16']:\n dtype = np.int16\n elif dtype_str in ['int8']:\n dtype = np.int8\n elif dtype_str in ['uint8']:\n dtype = np.uint8\n elif dtype_str in ['bool']:\n dtype = np.bool\n else:\n raise TypeError(f'Unrecognized dtype string: {dtype_str}')\n return dtype\n"
},
{
"path": "annlite/hubble_tools.py",
"content": "import os\nimport platform\nimport shutil\nimport time\nfrom pathlib import Path\nfrom typing import Optional, Union\n\nfrom filesplit.merge import Merge\nfrom filesplit.split import Split\nfrom loguru import logger\n\nignored_extn = ['.DS_Store']\n\n\ndef get_size(input: Path) -> float:\n import os\n\n return os.stat(str(input)).st_size / (1024 * 1024)\n\n\ndef make_archive(input: Path, output_name: str) -> Path:\n \"\"\"\n This function will create a zip archive of the input file (tmp.zip) at the\n same folder of input path.\n \"\"\"\n output_path = shutil.make_archive(\n os.path.join(str(input.parent), output_name),\n 'zip',\n str(input.parent),\n str(input.name),\n )\n return Path(output_path)\n\n\nclass Uploader:\n def __init__(self, size_limit=1024, client=None):\n \"\"\"\n This class create a filesplit object to split the file into small pieces and\n upload them on to hubble.\n :params size_limit: The max size of split files.\n :params client: hubble client used for uploading.\n \"\"\"\n self.size_limit = size_limit\n self.client = client\n\n def upload_file(\n self, input: Path, target_name: str, type: str, cell_id: Union[int, str]\n ):\n logger.info(f'Start to upload single file: {input} to hubble ...')\n size = get_size(input)\n if size > self.size_limit:\n split_list = self._split_file(input)\n self.upload_directory(split_list, target_name, type, cell_id, merge=False)\n shutil.rmtree(split_list)\n else:\n if self._check_exists(target_name, type, input.name):\n return\n self._upload_hubble(input, target_name, type, input.name, cell_id)\n\n def upload_directory(\n self,\n input: Path,\n target_name: str,\n type: str,\n cell_id: Union[int, str],\n merge: bool = True,\n ):\n def _upload():\n if self._check_exists(target_name, type, str(idx) + '.zip'):\n return\n\n Path.mkdir(input.parent / str(idx))\n for f in split_list:\n shutil.copy(f, input.parent / str(idx))\n output_path = make_archive(input.parent / str(idx), str(idx) + '.zip')\n self._upload_hubble(\n output_path, target_name, type, str(idx) + '.zip', cell_id\n )\n Path(output_path).unlink()\n shutil.rmtree(input.parent / str(idx))\n\n logger.info(f'Start to upload directory: {input} to hubble ...')\n if merge:\n size_list = list(\n zip(list(input.iterdir()), [get_size(f) for f in list(input.iterdir())])\n )\n sorted_size_list = sorted(size_list, key=lambda x: x[1])\n\n split_list = []\n total_size = 0\n idx = 0\n\n for file_name, file_size in sorted_size_list:\n for extn in ignored_extn:\n if extn in str(file_name):\n continue\n if total_size + file_size > self.size_limit:\n if len(split_list) == 0:\n raise Exception(\n f'The smallest file: {file_size} is bigger '\n f'than size_limit. Please set a larger value '\n f'of size_limit, now is {self.size_limit}MB.'\n )\n _upload()\n idx += 1\n total_size = 0\n split_list = [file_name]\n else:\n split_list.append(file_name)\n total_size += file_size\n if len(split_list) > 0:\n _upload()\n else:\n for idx, file_name in enumerate(list(input.glob('*'))):\n if self._check_exists(target_name, type, str(file_name.name)):\n continue\n self._upload_hubble(\n file_name, target_name, type, str(file_name.name), cell_id\n )\n\n def archive_and_upload(\n self,\n target_name: str,\n type: str,\n file_name: str,\n cell_id: Union[int, str],\n root_path: Path,\n upload_folder: str,\n ):\n if self._check_exists(target_name, type, file_name):\n return\n upload_file = shutil.make_archive(\n os.path.join(str(root_path), f'{target_name}_{type}'),\n 'zip',\n str(root_path),\n upload_folder,\n )\n\n logger.info(\n f'Start to upload: {upload_file} to hubble. '\n f'[target_name: {target_name}, '\n f'type: {type}, '\n f'file_name: {file_name}, '\n f'cell_id: {cell_id}].'\n )\n\n self.client.upload_artifact(\n f=upload_file,\n metadata={\n 'name': target_name,\n 'type': type,\n 'file_name': file_name,\n 'cell': cell_id,\n },\n )\n Path(upload_file).unlink()\n\n def _check_exists(self, target_name: str, type: str, file_name: str) -> bool:\n art_list = self.client.list_artifacts(\n filter={\n 'metaData.name': target_name,\n 'metaData.type': f'{type}',\n 'metaData.file_name': f'{file_name}',\n }\n )\n if len(art_list['data']) != 0:\n logger.info(\n f'[target_name: {target_name}, type: {type}, file_name: {file_name}] '\n f'already exists on hubble, will skip it ...'\n )\n return True\n else:\n return False\n\n def _split_file(self, input: Path) -> Path:\n output_dir = input / f'{input}_split'\n if output_dir.exists():\n logger.info(\n f'Origin file: {str(input)} has already been split to: {output_dir}, will skip ...'\n )\n return output_dir\n Path.mkdir(output_dir)\n Split(str(input), str(output_dir)).bysize(size=self.size_limit * 1024 * 1024)\n num_files = len(list(output_dir.glob('*')))\n logger.info(\n f'Origin file: {str(input)} has been split '\n f'into {num_files} parts. Output file: {output_dir}'\n )\n return output_dir\n\n def _upload_hubble(\n self,\n upload_file: Path,\n target_name: str,\n type: str,\n file_name: str,\n cell_id: Union[str, int],\n ):\n\n logger.info(\n f'Start to upload: {upload_file} to hubble. '\n f'[target_name: {target_name}, '\n f'type: {type}, '\n f'file_name: {file_name}, '\n f'cell_id: {cell_id}].'\n )\n\n start_time = time.time()\n failed_times = 0\n\n while True:\n try:\n self.client.upload_artifact(\n f=str(upload_file),\n metadata={\n 'name': target_name,\n 'type': type,\n 'file_name': file_name,\n 'cell': cell_id,\n },\n show_progress=True,\n )\n break\n except Exception as e:\n logger.info(e)\n failed_times += 1\n if failed_times == 3:\n logger.info(\n f'Tried more than 3 times to upload {upload_file}, type is: {type}, will exist...'\n )\n return\n else:\n continue\n\n logger.info(\n f'Takes {time.time() - start_time} seconds to upload {upload_file}.'\n )\n\n\nclass Merger:\n def __init__(self, restore_path, client):\n \"\"\"\n This class creates an object to download and merge the split files from hubble.\n :param restore_path: tmp directory for downloading and merging files.\n :param client: hubble client used for merging files.\n \"\"\"\n self.restore_path = restore_path\n self.restore_path.mkdir(parents=True)\n self.client = client\n\n def merge_file(self, inputdir: Path, outputdir: Path, outputfilename: Path):\n Merge(\n inputdir=str(inputdir),\n outputdir=str(outputdir),\n outputfilename=str(outputfilename),\n ).merge()\n\n def get_artifact_ids(self, art_list, type: str, cell_id: Optional[int] = None):\n ids = [\n [\n art['_id'],\n art['metaData']['type'],\n art['metaData']['file_name'],\n art['metaData']['cell'],\n ]\n for art in art_list['data']\n if type == art['metaData']['type']\n ]\n if cell_id:\n ids = [item for item in ids if int(item[3]) == cell_id]\n ids = [[item[0], item[1], item[2]] for item in ids]\n return ids\n\n def download(self, ids, download_folder):\n Path.mkdir(self.restore_path / download_folder)\n\n for ids, type, file_name in ids:\n self.client.download_artifact(\n id=ids,\n f=str(self.restore_path / download_folder / file_name),\n show_progress=True,\n )\n"
},
{
"path": "annlite/index.py",
"content": "import hashlib\nimport logging\nimport os\nimport platform\nimport warnings\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Dict, List, Optional, Union\n\nimport numpy as np\nfrom docarray.math.ndarray import to_numpy_array\nfrom loguru import logger\n\nif TYPE_CHECKING:\n from docarray import DocumentArray\n\nfrom .container import CellContainer\nfrom .core import PQCodec, ProjectorCodec, VQCodec\nfrom .enums import Metric\nfrom .filter import Filter\nfrom .helper import setup_logging\nfrom .math import cdist, top_k\n\nMAX_TRAINING_DATA_SIZE = 10240\n\n\nclass AnnLite(CellContainer):\n \"\"\":class:`AnnLite` is an approximate nearest neighbor search library.\n\n To create a :class:`AnnLite` object, simply:\n\n .. highlight:: python\n .. code-block:: python\n ann = AnnLite(256, metric='cosine')\n\n :param n_dim: dimensionality of input vectors. there are 2 constraints on dim:\n (1) it needs to be divisible by n_subvectors; (2) it needs to be a multiple of 4.*\n :param metric: distance metric type, can be 'euclidean', 'inner_product', or 'cosine'.\n :param n_subvectors: number of sub-quantizers, essentially this is the byte size of\n each quantized vector, default is None.\n :param n_cells: number of coarse quantizer clusters, default is 1.\n :param n_probe: number of cells to search for each query, default is 16.\n :param n_components: number of components to keep.\n :param initial_size: initial capacity assigned to each voronoi cell of coarse quantizer.\n ``n_cells * initial_size`` is the number of vectors that can be stored initially.\n if any cell has reached its capacity, that cell will be automatically expanded.\n If you need to add vectors frequently, a larger value for init_size is recommended.\n :param columns: the columns to be indexed for fast filtering, default is None.\n :param filterable_attrs: a dict of attributes to be indexed for fast filtering, default is None.\n The key is the attribute name, and the value is the attribute type. And it only works when ``columns`` is None.\n :param data_path: path to the directory where the data is stored.\n :param create_if_missing: if False, do not create the directory path if it is missing.\n :param read_only: if True, the index is not writable.\n :param verbose: if True, will print the debug logging info.\n\n .. note::\n Remember that the shape of any tensor that contains data points has to be `[n_data, dim]`.\n \"\"\"\n\n def __init__(\n self,\n n_dim: int,\n metric: Union[str, Metric] = 'cosine',\n n_cells: int = 1,\n n_subvectors: Optional[int] = None,\n n_clusters: Optional[int] = 256,\n n_probe: int = 16,\n n_components: Optional[int] = None,\n initial_size: Optional[int] = None,\n expand_step_size: int = 10240,\n columns: Optional[Union[Dict, List]] = None,\n filterable_attrs: Optional[Dict] = None,\n data_path: Union[Path, str] = Path('./data'),\n create_if_missing: bool = True,\n read_only: bool = False,\n verbose: bool = False,\n **kwargs,\n ):\n setup_logging(verbose)\n\n if 'dim' in kwargs:\n warnings.warn(\n 'The argument `dim` will be deprecated, please use `n_dim` instead.'\n )\n n_dim = kwargs['dim']\n\n if n_subvectors:\n assert (\n n_dim % n_subvectors == 0\n ), '\"n_dim\" needs to be divisible by \"n_subvectors\"'\n self.n_dim = n_dim\n self.n_components = n_components\n self.n_subvectors = n_subvectors\n self.n_clusters = n_clusters\n self.n_probe = max(n_probe, n_cells)\n self.n_cells = n_cells\n self.size_limit = 2048\n\n if isinstance(metric, str):\n metric = Metric.from_string(metric)\n self.metric = metric\n\n self._use_smart_probing = True\n\n self.read_only = read_only\n\n data_path = Path(data_path)\n if create_if_missing:\n data_path.mkdir(parents=True, exist_ok=True)\n self.data_path = data_path\n\n self._projector_codec = None\n if self._projector_codec_path.exists():\n logger.info(\n f'Load pre-trained projector codec (n_components={self.n_components}) from {self.model_path}'\n )\n self._projector_codec = ProjectorCodec.load(self._projector_codec_path)\n elif n_components:\n logger.info(\n f'Initialize Projector codec (n_components={self.n_components})'\n )\n self._projector_codec = ProjectorCodec(\n n_dim, n_components=self.n_components\n )\n\n self._vq_codec = None\n if self._vq_codec_path.exists():\n logger.info(\n f'Load trained VQ codec (K={self.n_cells}) from {self.model_path}'\n )\n self._vq_codec = VQCodec.load(self._vq_codec_path)\n elif n_cells > 1:\n logger.info(f'Initialize VQ codec (K={self.n_cells})')\n self._vq_codec = VQCodec(self.n_cells, metric=self.metric)\n\n self._pq_codec = None\n if self._pq_codec_path.exists():\n logger.info(\n f'Load trained PQ codec (n_subvectors={self.n_subvectors}) from {self.model_path}'\n )\n self._pq_codec = PQCodec.load(self._pq_codec_path)\n elif n_subvectors:\n logger.info(f'Initialize PQ codec (n_subvectors={self.n_subvectors})')\n self._pq_codec = PQCodec(\n dim=n_dim\n if not self._projector_codec\n else self._projector_codec.n_components,\n n_subvectors=self.n_subvectors,\n n_clusters=self.n_clusters,\n metric=self.metric,\n )\n\n if columns is not None:\n if filterable_attrs:\n logger.warning('`filterable_attrs` will be overwritten by `columns`.')\n\n filterable_attrs = {}\n for n, t in columns.items() if isinstance(columns, dict) else columns:\n filterable_attrs[n] = t\n\n super(AnnLite, self).__init__(\n n_dim,\n metric=metric,\n projector_codec=self._projector_codec,\n pq_codec=self._pq_codec,\n n_cells=n_cells,\n initial_size=initial_size,\n expand_step_size=expand_step_size,\n filterable_attrs=filterable_attrs,\n data_path=data_path,\n **kwargs,\n )\n\n if not self.is_trained and self.total_docs > 0:\n # train the index from scratch based on the data in the data_path\n logger.info(f'Train the index by reading data from {self.data_path}')\n total_size = 0\n # TODO: add a progress bar\n for docs in self.documents_generator(0, batch_size=1024):\n x = to_numpy_array(docs.embeddings)\n total_size += x.shape[0]\n self.partial_train(x, auto_save=True, force_train=True)\n if total_size >= MAX_TRAINING_DATA_SIZE:\n break\n logger.info(f'Total training data size: {total_size}')\n\n if self.total_docs > 0:\n self.restore()\n\n def _sanity_check(self, x: 'np.ndarray'):\n assert x.ndim == 2, 'inputs must be a 2D array'\n assert (\n x.shape[1] == self.n_dim\n ), f'inputs must have the same dimension as the index , got {x.shape[1]}, expected {self.n_dim}'\n\n return x.shape\n\n def train(self, x: 'np.ndarray', auto_save: bool = True, force_train: bool = False):\n \"\"\"Train the index with the given data.\n\n :param x: the ndarray data for training.\n :param auto_save: if False, will not dump the trained model to ``model_path``.\n :param force_train: if True, enforce to retrain the model, and overwrite the model if ``auto_save=True``.\n \"\"\"\n n_data, _ = self._sanity_check(x)\n\n if self.is_trained and not force_train:\n logger.warning(\n 'The indexer has been trained or is not trainable. Please use ``force_train=True`` to retrain.'\n )\n return\n\n if self._projector_codec:\n logger.info(\n f'Start training Projector codec (n_components={self.n_components}) with {n_data} data...'\n )\n self._projector_codec.fit(x)\n\n if self._vq_codec:\n logger.info(\n f'Start training VQ codec (K={self.n_cells}) with {n_data} data...'\n )\n self._vq_codec.fit(x)\n\n if self._pq_codec:\n logger.info(\n f'Start training PQ codec (n_subvectors={self.n_subvectors}) with {n_data} data...'\n )\n self._pq_codec.fit(x)\n\n logger.info(f'The annlite is successfully trained!')\n\n if auto_save:\n self.dump_model()\n\n def partial_train(\n self, x: np.ndarray, auto_save: bool = True, force_train: bool = False\n ):\n \"\"\"Partially train the index with the given data.\n\n :param x: the ndarray data for training.\n :param auto_save: if False, will not dump the trained model to ``model_path``.\n :param force_train: if True, enforce to retrain the model, and overwrite the model if ``auto_save=True``.\n\n \"\"\"\n n_data, _ = self._sanity_check(x)\n\n if self.is_trained and not force_train:\n logger.warning(\n 'The annlite has been trained or is not trainable. Please use ``force_train=True`` to retrain.'\n )\n return\n\n if self._projector_codec:\n logging.info(\n f'Partial training Projector codec (n_components={self.n_components}) with {n_data} data...'\n )\n self._projector_codec.partial_fit(x)\n\n if self._vq_codec:\n logger.info(\n f'Partial training VQ codec (K={self.n_cells}) with {n_data} data...'\n )\n self._vq_codec.partial_fit(x)\n\n if self._pq_codec:\n logger.info(\n f'Partial training PQ codec (n_subvectors={self.n_subvectors}) with {n_data} data...'\n )\n self._pq_codec.partial_fit(x)\n\n if auto_save:\n self.dump_model()\n\n def index(self, docs: 'DocumentArray', **kwargs):\n \"\"\"Add the documents to the index.\n\n :param docs: the document array to be indexed.\n \"\"\"\n\n if self.read_only:\n logger.error('The indexer is readonly, cannot add new documents')\n return\n\n if not self.is_trained:\n raise RuntimeError(f'The indexer is not trained, cannot add new documents')\n\n x = to_numpy_array(docs.embeddings)\n n_data, _ = self._sanity_check(x)\n\n assigned_cells = (\n self._vq_codec.encode(x)\n if self._vq_codec\n else np.zeros(n_data, dtype=np.int64)\n )\n return super(AnnLite, self).insert(x, assigned_cells, docs)\n\n def update(\n self,\n docs: 'DocumentArray',\n raise_errors_on_not_found: bool = False,\n insert_if_not_found: bool = True,\n **kwargs,\n ):\n \"\"\"Update the documents in the index.\n\n :param insert_if_not_found: whether to raise error when updated id is not found.\n :param raise_errors_on_not_found: whether to raise exception when id not found.\n :param docs: the document array to be updated.\n \"\"\"\n if self.read_only:\n logger.error('The indexer is readonly, cannot update documents')\n return\n\n if not self.is_trained:\n raise RuntimeError(f'The indexer is not trained, cannot add new documents')\n\n x = to_numpy_array(docs.embeddings)\n n_data, _ = self._sanity_check(x)\n\n assigned_cells = (\n self._vq_codec.encode(x)\n if self._vq_codec\n else np.zeros(n_data, dtype=np.int64)\n )\n\n return super(AnnLite, self).update(\n x,\n assigned_cells,\n docs,\n raise_errors_on_not_found=raise_errors_on_not_found,\n insert_if_not_found=insert_if_not_found,\n )\n\n def search(\n self,\n docs: 'DocumentArray',\n filter: Optional[dict] = None,\n limit: int = 10,\n include_metadata: bool = True,\n **kwargs,\n ):\n \"\"\"Search the index, and attach matches to the query Documents in `docs`\n\n :param docs: the document array to be searched.\n :param filter: the filter to be applied to the search.\n :param limit: the number of results to get for each query document in search\n :param include_metadata: whether to return document metadata in response.\n \"\"\"\n if not self.is_trained:\n raise RuntimeError(f'The indexer is not trained, cannot add new documents')\n\n query_np = to_numpy_array(docs.embeddings)\n\n match_dists, match_docs = self.search_by_vectors(\n query_np, filter=filter, limit=limit, include_metadata=include_metadata\n )\n\n for doc, matches in zip(docs, match_docs):\n doc.matches = matches\n\n def search_by_vectors(\n self,\n query_np: 'np.ndarray',\n filter: Optional[dict] = None,\n limit: int = 10,\n include_metadata: bool = True,\n ):\n \"\"\"Search the index by vectors, and return the matches.\n\n :param query_np: the query vectors.\n :param filter: the filter to be applied to the search.\n :param limit: the number of results to get for each query document in search\n :param include_metadata: whether to return document metadata in response.\n \"\"\"\n\n cells = self._cell_selection(query_np, limit)\n where_clause, where_params = Filter(filter or {}).parse_where_clause()\n\n match_dists, match_docs = self.search_cells(\n query=query_np,\n cells=cells,\n where_clause=where_clause,\n where_params=where_params,\n limit=limit,\n include_metadata=include_metadata,\n )\n return match_dists, match_docs\n\n def filter(\n self,\n filter: Dict,\n limit: int = 10,\n offset: int = 0,\n order_by: Optional[str] = None,\n ascending: bool = True,\n include_metadata: bool = True,\n ):\n \"\"\"Find the documents by the filter.\n\n :param filter: the filter to be applied to the search.\n :param limit: the number of results.\n :param offset: the offset of the results.\n :param order_by: the field to order the results.\n :param ascending: whether to order the results in ascending order.\n :param include_metadata: whether to return document metadata in response.\n \"\"\"\n cells = [x for x in range(self.n_cells)]\n where_clause, where_params = Filter(filter or {}).parse_where_clause()\n\n match_docs = self.filter_cells(\n cells=cells,\n where_clause=where_clause,\n where_params=where_params,\n limit=limit,\n offset=offset,\n order_by=order_by,\n ascending=ascending,\n include_metadata=include_metadata,\n )\n if limit > 0:\n return match_docs[:limit]\n return match_docs\n\n def get_doc_by_id(self, doc_id: str):\n \"\"\"Get the document by id.\n\n :param doc_id: the document id.\n \"\"\"\n\n return self._get_doc_by_id(doc_id)\n\n def get_docs(\n self,\n filter: Optional[dict] = None,\n limit: int = 10,\n offset: int = 0,\n order_by: Optional[str] = None,\n ascending: bool = True,\n ):\n \"\"\"Get the documents.\n\n :param filter: the filter to be applied to the search.\n :param limit: the number of results.\n :param offset: the offset of the results.\n :param order_by: the field to order the results.\n :param ascending: whether to order the results in ascending order. It only works when `order_by` is specified.\n \"\"\"\n\n return self.filter(\n filter=filter,\n limit=limit,\n offset=offset,\n order_by=order_by,\n ascending=ascending,\n include_metadata=True,\n )\n\n def _cell_selection(self, query_np, limit):\n n_data, _ = self._sanity_check(query_np)\n\n if self._vq_codec:\n dists = cdist(\n query_np, self._vq_codec.codebook, metric=self.metric.name.lower()\n )\n dists, cells = top_k(dists, k=self.n_probe)\n else:\n cells = np.zeros((n_data, 1), dtype=np.int64)\n\n # if self.use_smart_probing and self.n_probe > 1:\n # p = -topk_sims.abs().sqrt()\n # p = torch.softmax(p / self.smart_probing_temperature, dim=-1)\n #\n # # p_norm = p.norm(dim=-1)\n # # sqrt_d = self.n_probe ** 0.5\n # # score = 1 - (p_norm * sqrt_d - 1) / (sqrt_d - 1) - 1e-6\n # # n_probe_list = torch.ceil(score * (self.n_probe) ).long()\n #\n # max_n_probe = torch.tensor(self.n_probe, device=self.device)\n # normalized_entropy = - torch.sum(p * torch.log2(p) / torch.log2(max_n_probe), dim=-1)\n # n_probe_list = torch.ceil(normalized_entropy * max_n_probe).long()\n # else:\n # n_probe_list = None\n return cells\n\n def search_numpy(\n self,\n query_np: 'np.ndarray',\n filter: Dict = {},\n limit: int = 10,\n **kwargs,\n ):\n \"\"\"Search the index and return distances to the query and ids of the closest documents.\n\n :param query_np: matrix containing query vectors as rows\n :param filter: the filtering conditions\n :param limit: the number of results to get for each query document in search\n \"\"\"\n\n if not self.is_trained:\n raise RuntimeError(f'The indexer is not trained, cannot add new documents')\n\n dists, doc_ids = self._search_numpy(query_np, filter, limit)\n return dists, doc_ids\n\n def _search_numpy(self, query_np: 'np.ndarray', filter: Dict = {}, limit: int = 10):\n \"\"\"Search approximate nearest vectors in different cells, returns distances and ids\n\n :param query_np: matrix containing query vectors as rows\n :param filter: the filtering conditions\n :param limit: the number of results to get for each query document in search\n \"\"\"\n cells = self._cell_selection(query_np, limit)\n where_clause, where_params = Filter(filter).parse_where_clause()\n\n dists, ids = self._search_cells(\n query=query_np,\n cells=cells,\n where_clause=where_clause,\n where_params=where_params,\n limit=limit,\n )\n return dists, ids\n\n def delete(\n self,\n docs: Union['DocumentArray', List[str]],\n raise_errors_on_not_found: bool = False,\n ):\n \"\"\"Delete entries from the index by id\n\n :param raise_errors_on_not_found: whether to raise exception when id not found.\n :param docs: the documents to delete\n \"\"\"\n\n super().delete(\n docs if isinstance(docs, list) else docs[:, 'id'], raise_errors_on_not_found\n )\n\n def clear(self):\n \"\"\"Clear the whole database\"\"\"\n for cell_id in range(self.n_cells):\n self.vec_index(cell_id).reset()\n self.cell_table(cell_id).clear()\n self.doc_store(cell_id).clear()\n self.meta_table.clear()\n\n def close(self):\n for cell_id in range(self.n_cells):\n self.doc_store(cell_id).close()\n\n def encode(self, x: 'np.ndarray'):\n n_data, _ = self._sanity_check(x)\n\n if self._projector_codec:\n x = self._projector_codec.encode(x)\n\n if self._vq_codec:\n x = self._pq_codec.encode(x)\n\n return x\n\n def decode(self, x: 'np.ndarray'):\n assert len(x.shape) == 2\n assert x.shape[1] == self.n_subvectors\n\n if self._pq_codec:\n x = self._pq_codec.decode(x)\n\n if self._projector_codec:\n x = self._projector_codec.decode(x)\n\n return x\n\n @property\n def params_hash(self):\n model_metas = (\n f'n_dim: {self.n_dim} '\n f'metric: {self.metric} '\n f'n_cells: {self.n_cells} '\n f'n_components: {self.n_components} '\n f'n_subvectors: {self.n_subvectors}'\n )\n return hashlib.md5(f'{model_metas}'.encode()).hexdigest()\n\n @property\n def model_path(self):\n return self.data_path / f'parameters-{self.params_hash}'\n\n @property\n def _vq_codec_path(self):\n return self.model_path / f'vq_codec.params'\n\n @property\n def _pq_codec_path(self):\n return self.model_path / f'pq_codec.params'\n\n @property\n def _projector_codec_path(self):\n return self.model_path / f'projector_codec.params'\n\n @property\n def index_hash(self):\n latest_commit = self.meta_table.get_latest_commit()\n date_time = latest_commit[-1] if latest_commit else None\n if date_time:\n if platform.system() == 'Windows':\n return date_time.isoformat('#', 'hours')\n return date_time.isoformat('#', 'seconds')\n else:\n import datetime\n\n return (\n datetime.datetime.utcnow().isoformat('#', 'hours')\n if platform.system() == 'Windows'\n else datetime.datetime.utcnow().isoformat('#', 'seconds')\n )\n\n @property\n def index_path(self):\n if self.index_hash:\n return (\n self.data_path\n / f'snapshot-{self.params_hash}'\n / f'{self.index_hash}-SNAPSHOT'\n )\n return None\n\n @property\n def snapshot_path(self):\n paths = list(\n (self.data_path / f'snapshot-{self.params_hash}').glob(f'*-SNAPSHOT')\n )\n\n if paths:\n paths = sorted(paths, key=lambda x: x.name)\n return paths[-1]\n return None\n\n @property\n def remote_store_client(self):\n try:\n import hubble\n\n os.environ['JINA_AUTH_TOKEN'] = self.token\n client = hubble.Client(max_retries=None, jsonify=True)\n client.get_user_info()\n return client\n except Exception as ex:\n logger.error(f'Not login to hubble yet.')\n raise ex\n\n def backup(self, target_name: Optional[str] = None, token: Optional[str] = None):\n # file lock will be released when backup to remote, this will\n # release the file lock. And it's only needed in Windows\n # since we need to release file lock before we can access rocksdb files.\n if not target_name:\n logger.info('dump to local ...')\n self.dump()\n else:\n if token is None:\n logger.error(f'back up to remote needs token')\n logger.info(f'dump to remote: {target_name}')\n self.close()\n self._backup_index_to_remote(target_name, token)\n\n def restore(self, source_name: Optional[str] = None, token: Optional[str] = None):\n # file lock will be released when restore from remote\n if not source_name:\n if self.total_docs > 0:\n logger.info(f'restore Annlite from local')\n self._rebuild_index_from_local()\n else:\n if token is None:\n logger.error(f'restore from remote needs token')\n logger.info(f'restore Annlite from artifact: {source_name}')\n self.close()\n self._rebuild_index_from_remote(source_name, token)\n\n def dump_model(self):\n logger.info(f'Save the parameters to {self.model_path}')\n self.model_path.mkdir(parents=True, exist_ok=True)\n if self._projector_codec:\n self._projector_codec.dump(self._projector_codec_path)\n if self._vq_codec:\n self._vq_codec.dump(self._vq_codec_path)\n if self._pq_codec:\n self._pq_codec.dump(self._pq_codec_path)\n\n def dump_index(self):\n import shutil\n\n logger.info(f'Save the indexer to {self.index_path}')\n try:\n if Path.exists(self.index_path):\n logger.info(\n f'Index path {self.index_path} already exists, will be '\n f'overwritten'\n )\n shutil.rmtree(self.index_path)\n self.index_path.mkdir(parents=True)\n\n for cell_id in range(self.n_cells):\n self.vec_index(cell_id).dump(self.index_path / f'cell_{cell_id}.hnsw')\n self.cell_table(cell_id).dump(self.index_path / f'cell_{cell_id}.db')\n self.meta_table.dump(self.index_path / f'meta.db')\n except Exception as ex:\n logger.error(f'Failed to dump the indexer, {ex!r}')\n\n if self.index_path:\n shutil.rmtree(self.index_path)\n\n def dump(self):\n self.dump_model()\n self.dump_index()\n\n def _backup_index_to_remote(self, target_name: str, token: str):\n\n self.dump()\n\n from .hubble_tools import Uploader\n\n self.token = token\n client = self.remote_store_client\n uploader = Uploader(size_limit=self.size_limit, client=client)\n\n for cell_id in range(self.n_cells):\n # upload database\n uploader.upload_directory(\n Path(self.data_path) / f'cell_{cell_id}',\n target_name=target_name,\n type='database',\n cell_id=cell_id,\n )\n\n # upload hnsw file\n uploader.upload_file(\n Path(self.index_path) / f'cell_{cell_id}.hnsw',\n target_name=target_name,\n type='hnsw',\n cell_id=cell_id,\n )\n\n # upload cell_table\n uploader.upload_file(\n Path(self.index_path) / f'cell_{cell_id}.db',\n target_name=target_name,\n type='cell_table',\n cell_id=cell_id,\n )\n\n # upload meta table\n uploader.upload_file(\n Path(self.index_path) / 'meta.db',\n target_name=target_name,\n type='meta_table',\n cell_id=0,\n )\n\n # upload training model\n uploader.archive_and_upload(\n target_name,\n 'model',\n 'model.zip',\n 'all',\n self.model_path.parent,\n str(self.model_path.name),\n )\n\n def _rebuild_index_from_local(self):\n if self.snapshot_path:\n logger.info(f'Load the indexer from snapshot {self.snapshot_path}')\n for cell_id in range(self.n_cells):\n self.vec_index(cell_id).load(\n self.snapshot_path / f'cell_{cell_id}.hnsw'\n )\n self.cell_table(cell_id).load(self.snapshot_path / f'cell_{cell_id}.db')\n self.meta_table.load(self.snapshot_path / f'meta.db')\n else:\n logger.info(f'Rebuild the indexer from scratch')\n for cell_id in range(self.n_cells):\n cell_size = self.doc_store(cell_id).size\n\n if cell_size == 0:\n continue # skip empty cell\n\n logger.debug(\n f'Rebuild the index of cell-{cell_id} ({cell_size} docs)...'\n )\n for docs in self.documents_generator(cell_id, batch_size=10240):\n x = to_numpy_array(docs.embeddings)\n\n assigned_cells = np.ones(len(docs), dtype=np.int64) * cell_id\n super().insert(x, assigned_cells, docs, only_index=True)\n logger.debug(f'Rebuild the index of cell-{cell_id} done')\n if self.model_path:\n logger.info(f'Load the model from {self.model_path}')\n self._reload_models()\n\n def _rebuild_index_from_remote(self, source_name: str, token: str):\n import shutil\n\n from .hubble_tools import Merger\n\n self.token = token\n client = self.remote_store_client\n art_list = client.list_artifacts(\n filter={'metaData.name': source_name}, pageSize=100\n )\n if len(art_list['data']) == 0:\n logger.info(f'The indexer `{source_name}` not found. ')\n else:\n logger.info(f'Load the indexer `{source_name}` from remote store')\n\n restore_path = self.data_path / 'restore'\n merger = Merger(restore_path=restore_path, client=client)\n\n for cell_id in range(self.n_cells):\n # download hnsw files and merge and load\n logger.info(f'Load the hnsw `{source_name}` from remote store')\n\n hnsw_ids = merger.get_artifact_ids(\n art_list, type='hnsw', cell_id=cell_id\n )\n merger.download(ids=hnsw_ids, download_folder=f'hnsw_{cell_id}')\n if len(hnsw_ids) > 1:\n merger.merge_file(\n inputdir=restore_path / f'hnsw_{cell_id}',\n outputdir=restore_path / f'hnsw_{cell_id}',\n outputfilename=Path(f'cell_{cell_id}.hnsw'),\n )\n self.vec_index(cell_id).load(\n restore_path / f'hnsw_{cell_id}' / f'cell_{cell_id}.hnsw'\n )\n shutil.rmtree(restore_path / f'hnsw_{cell_id}')\n\n # download cell_table files and merge and load\n logger.info(f'Load the cell_table `{source_name}` from remote store')\n\n cell_table_ids = merger.get_artifact_ids(\n art_list, type='cell_table', cell_id=cell_id\n )\n merger.download(\n ids=cell_table_ids, download_folder=f'cell_table_{cell_id}'\n )\n if len(cell_table_ids) > 1:\n merger.merge_file(\n inputdir=restore_path / f'cell_table_{cell_id}',\n outputdir=restore_path / f'cell_table_{cell_id}',\n outputfilename=Path(f'cell_{cell_id}.db'),\n )\n\n self.cell_table(cell_id).load(\n restore_path / f'cell_table_{cell_id}' / f'cell_{cell_id}.db'\n )\n shutil.rmtree(restore_path / f'cell_table_{cell_id}')\n\n # download database files and rebuild\n logger.info(f'Load the database `{source_name}` from remote store')\n\n database_ids = merger.get_artifact_ids(\n art_list, type='database', cell_id=cell_id\n )\n merger.download(ids=database_ids, download_folder='database')\n for zip_file in list((restore_path / 'database').iterdir()):\n # default has only one cell\n shutil.unpack_archive(zip_file, self.data_path / f'cell_{cell_id}')\n for f in list(\n (\n self.data_path\n / f'cell_{cell_id}'\n / zip_file.name.split('.zip')[0]\n ).iterdir()\n ):\n origin_database_path = (\n self.data_path / f'cell_{cell_id}' / f.name\n )\n if origin_database_path.exists():\n origin_database_path.unlink()\n f.rename(self.data_path / f'cell_{cell_id}' / f.name)\n shutil.rmtree(\n self.data_path\n / f'cell_{cell_id}'\n / zip_file.name.split('.zip')[0]\n )\n Path(zip_file).unlink()\n self._rebuild_database()\n\n # download meta_table files\n logger.info(f'Load the meta_table `{source_name}` from remote store')\n\n meta_table_ids = merger.get_artifact_ids(\n art_list, type='meta_table', cell_id=0\n )\n merger.download(ids=meta_table_ids, download_folder='meta_table')\n\n if len(meta_table_ids) > 1:\n merger.merge_file(\n inputdir=restore_path / 'meta_table',\n outputdir=restore_path / 'meta_table',\n outputfilename=Path('meta.db'),\n )\n self._meta_table.load(restore_path / 'meta_table' / 'meta.db')\n shutil.rmtree(restore_path / 'meta_table')\n\n # download model files\n logger.info(f'Load the model `{source_name}` from remote store')\n file_name = str(self.model_path.parent / f'{source_name}_model.zip')\n model_id = [\n art['_id']\n for art in art_list['data']\n if 'model' in art['metaData']['type']\n ]\n assert len(model_id) == 1\n client.download_artifact(\n id=model_id[0],\n f=file_name,\n show_progress=True,\n )\n shutil.unpack_archive(file_name, self.model_path.parent)\n self._reload_models()\n Path(file_name).unlink()\n\n shutil.rmtree(restore_path)\n\n @property\n def is_trained(self):\n if self._projector_codec and (not self._projector_codec.is_trained):\n return False\n if self._vq_codec and (not self._vq_codec.is_trained):\n return False\n if self._pq_codec and (not self._pq_codec.is_trained):\n return False\n return True\n\n def _reload_models(self):\n if self._projector_codec_path.exists():\n self._projector_codec = ProjectorCodec.load(self._projector_codec_path)\n if self._vq_codec_path.exists():\n self._vq_codec = VQCodec.load(self._vq_codec_path)\n if self._pq_codec_path.exists():\n self._pq_codec = PQCodec.load(self._pq_codec_path)\n\n @property\n def use_smart_probing(self):\n return self._use_smart_probing\n\n @use_smart_probing.setter\n def use_smart_probing(self, value):\n assert type(value) is bool\n self._use_smart_probing = value\n\n @property\n def stat(self):\n \"\"\"Get information on status of the indexer.\"\"\"\n return {\n 'total_docs': self.total_docs,\n 'index_size': self.index_size,\n 'n_cells': self.n_cells,\n 'n_dim': self.n_dim,\n 'n_components': self.n_components,\n 'metric': self.metric.name,\n 'is_trained': self.is_trained,\n }\n\n # @property\n # def smart_probing_temperature(self):\n # return self._smart_probing_temperature\n #\n # @smart_probing_temperature.setter\n # def smart_probing_temperature(self, value):\n # assert value > 0\n # assert self.use_smart_probing, 'set use_smart_probing to True first'\n # self._smart_probing_temperature = value\n"
},
{
"path": "annlite/math.py",
"content": "from typing import Tuple\n\nimport numpy as np\n\n\ndef l2_normalize(x: 'np.ndarray', eps: float = np.finfo(np.float32).eps):\n \"\"\"Scale input vectors individually to unit norm.\n\n :param x: The data to normalize\n :param eps: a small jitter to avoid divde by zero\n :return: Normalized input X\n \"\"\"\n\n norms = np.einsum('ij,ij->i', x, x)\n np.sqrt(norms, norms)\n constant_mask = norms < 10 * eps\n norms[constant_mask] = 1.0\n return x / norms[:, np.newaxis]\n\n\ndef cosine(\n x_mat: 'np.ndarray', y_mat: 'np.ndarray', eps: float = np.finfo(np.float32).eps\n) -> 'np.ndarray':\n \"\"\"Cosine distance between each row in x_mat and each row in y_mat.\n\n :param x_mat: np.ndarray with ndim=2\n :param y_mat: np.ndarray with ndim=2\n :param eps: a small jitter to avoid divde by zero\n :return: np.ndarray with ndim=2\n \"\"\"\n return 1 - np.clip(\n (np.dot(x_mat, y_mat.T) + eps)\n / (\n np.outer(np.linalg.norm(x_mat, axis=1), np.linalg.norm(y_mat, axis=1)) + eps\n ),\n -1,\n 1,\n )\n\n\ndef sqeuclidean(x_mat: 'np.ndarray', y_mat: 'np.ndarray') -> 'np.ndarray':\n \"\"\"Squared Euclidean distance between each row in x_mat and each row in y_mat.\n :param x_mat: np.ndarray with ndim=2\n :param y_mat: np.ndarray with ndim=2\n :return: np.ndarray with ndim=2\n \"\"\"\n return (\n np.sum(y_mat**2, axis=1)\n + np.sum(x_mat**2, axis=1)[:, np.newaxis]\n - 2 * np.dot(x_mat, y_mat.T)\n )\n\n\ndef euclidean(x_mat: 'np.ndarray', y_mat: 'np.ndarray') -> 'np.ndarray':\n \"\"\"Euclidean distance between each row in x_mat and each row in y_mat.\n\n :param x_mat: scipy.sparse like array with ndim=2\n :param y_mat: scipy.sparse like array with ndim=2\n :return: np.ndarray with ndim=2\n \"\"\"\n return np.sqrt(sqeuclidean(x_mat, y_mat))\n\n\ndef pdist(\n x_mat: 'np.ndarray',\n metric: str,\n) -> 'np.ndarray':\n \"\"\"Computes Pairwise distances between observations in n-dimensional space.\n\n :param x_mat: Union['np.ndarray','scipy.sparse.csr_matrix', 'scipy.sparse.coo_matrix'] of ndim 2\n :param metric: string describing the metric type\n :return: np.ndarray of ndim 2\n \"\"\"\n return cdist(x_mat, x_mat, metric)\n\n\ndef cdist(x_mat: 'np.ndarray', y_mat: 'np.ndarray', metric: str) -> 'np.ndarray':\n \"\"\"Computes the pairwise distance between each row of X and each row on Y according to `metric`.\n - Let `n_x = x_mat.shape[0]`\n - Let `n_y = y_mat.shape[0]`\n - Returns a matrix `dist` of shape `(n_x, n_y)` with `dist[i,j] = metric(x_mat[i], y_mat[j])`.\n :param x_mat: numpy or scipy array of ndim 2\n :param y_mat: numpy or scipy array of ndim 2\n :param metric: string describing the metric type\n :return: np.ndarray of ndim 2\n \"\"\"\n dists = {'cosine': cosine, 'sqeuclidean': sqeuclidean, 'euclidean': euclidean}[\n metric\n ](x_mat, y_mat)\n\n return dists\n\n\ndef top_k(\n values: 'np.ndarray', k: int, descending: bool = False\n) -> Tuple['np.ndarray', 'np.ndarray']:\n \"\"\"Finds values and indices of the k largest entries for the last dimension.\n\n :param values: array of distances\n :param k: number of values to retrieve\n :param descending: find top k biggest values\n :return: indices and distances\n \"\"\"\n if descending:\n values = -values\n\n if k >= values.shape[1]:\n idx = values.argsort(axis=1)[:, :k]\n values = np.take_along_axis(values, idx, axis=1)\n else:\n idx_ps = values.argpartition(kth=k, axis=1)[:, :k]\n values = np.take_along_axis(values, idx_ps, axis=1)\n idx_fs = values.argsort(axis=1)\n idx = np.take_along_axis(idx_ps, idx_fs, axis=1)\n values = np.take_along_axis(values, idx_fs, axis=1)\n\n if descending:\n values = -values\n\n return values, idx\n"
},
{
"path": "annlite/profile.py",
"content": "import cProfile\nimport pstats\nimport random\nfrom functools import wraps\n\nrandom.seed(20)\n\ntry:\n import builtins\n\n line_profile = builtins.profile\nexcept AttributeError:\n # No line profiler, provide a pass-through version\n def profile(func):\n return func\n\n line_profile = profile\n\n\ndef time_profile(\n output_file=None, sort_by='cumulative', lines_to_print=None, strip_dirs=False\n):\n \"\"\"A time profiler decorator.\n Inspired by and modified the profile decorator of Giampaolo Rodola:\n https://github.com/ekhoda/profile_decorator\n Args:\n output_file: str or None. Default is None\n Path of the output file. If only name of the file is given, it's\n saved in the current directory.\n If it's None, the name of the decorated function is used.\n sort_by: str or SortKey enum or tuple/list of str/SortKey enum\n Sorting criteria for the Stats object.\n For a list of valid string and SortKey refer to:\n https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats\n lines_to_print: int or None\n Number of lines to print. Default (None) is for all the lines.\n This is useful in reducing the size of the printout, especially\n that sorting by 'cumulative', the time consuming operations\n are printed toward the top of the file.\n strip_dirs: bool\n Whether to remove the leading path info from file names.\n This is also useful in reducing the size of the printout\n Returns:\n Profile of the decorated function\n \"\"\"\n\n def inner(func):\n @wraps(func)\n def wrapper(*args, **kwargs):\n _output_file = output_file or func.__name__ + '.prof'\n pr = cProfile.Profile()\n pr.enable()\n retval = func(*args, **kwargs)\n pr.disable()\n pr.dump_stats(_output_file)\n\n with open(_output_file, 'w') as f:\n ps = pstats.Stats(pr, stream=f)\n if strip_dirs:\n ps.strip_dirs()\n if isinstance(sort_by, (tuple, list)):\n ps.sort_stats(*sort_by)\n else:\n ps.sort_stats(sort_by)\n ps.print_stats(lines_to_print)\n return retval\n\n return wrapper\n\n return inner\n"
},
{
"path": "annlite/storage/__init__.py",
"content": ""
},
{
"path": "annlite/storage/base.py",
"content": "import abc\nfrom typing import TYPE_CHECKING, List, Optional\n\nif TYPE_CHECKING:\n import numpy as np\n\nfrom ..enums import ExpandMode\n\n\nclass Storage(abc.ABC):\n def __init__(\n self,\n initial_size: Optional[int] = None,\n expand_step_size: int = 10240,\n expand_mode: ExpandMode = ExpandMode.ADAPTIVE,\n ):\n if initial_size is None:\n initial_size = expand_step_size\n assert initial_size >= 0\n assert expand_step_size > 0\n\n self.initial_size = initial_size\n self.expand_step_size = expand_step_size\n self.expand_mode = expand_mode\n\n @property\n @abc.abstractmethod\n def capacity(self) -> int:\n ...\n\n @property\n @abc.abstractmethod\n def size(self):\n ...\n\n @abc.abstractmethod\n def clean(self):\n ...\n\n @abc.abstractmethod\n def add(\n self,\n data: 'np.ndarray',\n cells: 'np.ndarray',\n ids: List[str],\n doc_tags: Optional[List[dict]] = None,\n ):\n ...\n\n @abc.abstractmethod\n def delete(self, ids: List[str]):\n ...\n\n @abc.abstractmethod\n def update(\n self,\n data: 'np.ndarray',\n cells: 'np.ndarray',\n ids: List[str],\n doc_tags: Optional[List[dict]] = None,\n ):\n ...\n"
},
{
"path": "annlite/storage/kv.py",
"content": "import time\nimport warnings\nfrom pathlib import Path\nfrom typing import Dict, List, Union\n\nfrom docarray import Document, DocumentArray\nfrom rocksdict import Options, Rdict, ReadOptions, WriteBatch, WriteOptions\n\n\nclass DocStorage:\n \"\"\"The backend storage engine of Documents\"\"\"\n\n def __init__(\n self,\n path: Union[str, Path],\n serialize_config: Dict = {},\n create_if_missing: bool = True,\n **kwargs,\n ):\n self._path = str(path)\n self._serialize_config = serialize_config\n\n self._kwargs = kwargs\n\n self._init_db(create_if_missing=create_if_missing, **self._kwargs)\n\n def _init_db(self, create_if_missing: bool = True, **kwargs):\n opt = Options(raw_mode=True)\n\n opt.optimize_for_point_lookup(1024)\n\n opt.set_inplace_update_support(True)\n opt.set_allow_concurrent_memtable_write(False)\n\n # configure mem-table to a large value (256 MB)\n opt.set_write_buffer_size(0x10000000)\n\n # 256 MB file size\n opt.set_target_file_size_base(0x10000000)\n\n # # set to plain-table for better performance\n # opt.set_plain_table_factory(PlainTableFactoryOptions())\n\n opt.create_if_missing(create_if_missing)\n\n self._db = Rdict(path=self._path, options=opt)\n\n # get the size of the database, if it is not created, set it to 0\n self._size = len(list(self._db.keys()))\n\n self._is_closed = False\n\n def insert(self, docs: 'DocumentArray'):\n write_batch = WriteBatch(raw_mode=True)\n write_opt = WriteOptions()\n write_opt.sync = True\n batch_size = 0\n for doc in docs:\n write_batch.put(doc.id.encode(), doc.to_bytes(**self._serialize_config))\n batch_size += 1\n self._db.write(write_batch, write_opt=write_opt)\n self._size += batch_size\n\n def update(self, docs: 'DocumentArray'):\n write_batch = WriteBatch(raw_mode=True)\n write_opt = WriteOptions()\n write_opt.sync = True\n for doc in docs:\n key = doc.id.encode()\n if key not in self._db:\n raise ValueError(f'The Doc ({doc.id}) does not exist in database!')\n\n write_batch.put(key, doc.to_bytes(**self._serialize_config))\n self._db.write(write_batch, write_opt=write_opt)\n\n def delete(self, doc_ids: List[str]):\n write_batch = WriteBatch(raw_mode=True)\n write_opt = WriteOptions()\n write_opt.sync = True\n for doc_id in doc_ids:\n write_batch.delete(doc_id.encode())\n self._db.write(write_batch, write_opt=write_opt)\n self._size -= len(doc_ids)\n\n def get(self, doc_ids: Union[str, list]) -> DocumentArray:\n docs = DocumentArray()\n if isinstance(doc_ids, str):\n doc_ids = [doc_ids]\n\n for doc_bytes in self._db[[k.encode() for k in doc_ids]]:\n if doc_bytes:\n docs.append(Document.from_bytes(doc_bytes, **self._serialize_config))\n\n return docs\n\n def clear(self):\n if self._is_closed:\n warnings.warn(\n '`DocStorage` had been closed already, will skip this close operation.'\n )\n else:\n self._db.close()\n self._db.destroy(self._path)\n\n # re-initialize the database for the next usage\n self._init_db(create_if_missing=True, **self._kwargs)\n\n def close(self):\n if self._is_closed:\n warnings.warn(\n '`DocStorage` had been closed already, will skip this close operation.'\n )\n return\n try:\n self._db.flush(wait=True)\n self._db.close()\n except Exception as ex:\n if 'No such file or directory' not in str(ex):\n # this is a known bug, we can safely ignore it\n raise ex\n self._is_closed = True\n\n def __len__(self):\n return self._size\n\n @property\n def stat(self):\n return {'entries': len(self)}\n\n @property\n def size(self):\n return self.stat['entries']\n\n @property\n def last_transaction_id(self):\n return self._db.latest_sequence_number()\n\n def batched_iterator(self, batch_size: int = 1, **kwargs) -> 'DocumentArray':\n count = 0\n docs = DocumentArray()\n\n read_opt = ReadOptions()\n\n for value in self._db.values(read_opt=read_opt):\n doc = Document.from_bytes(value, **self._serialize_config)\n docs.append(doc)\n count += 1\n\n if count == batch_size:\n yield docs\n count = 0\n docs = DocumentArray()\n\n if count > 0:\n yield docs\n"
},
{
"path": "annlite/storage/table.py",
"content": "import datetime\nimport sqlite3\nimport threading\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union\n\nimport numpy as np\n\nif TYPE_CHECKING:\n from docarray import DocumentArray\n\nsqlite3.register_adapter(np.int64, lambda x: int(x))\nsqlite3.register_adapter(np.int32, lambda x: int(x))\n\nCOLUMN_TYPE_MAPPING = {\n float: 'FLOAT',\n int: 'INTEGER',\n bool: 'INTEGER',\n str: 'TEXT',\n bytes.__class__: 'BLOB',\n bytes: 'BLOB',\n memoryview: 'BLOB',\n # datetime.datetime: 'TEXT',\n # datetime.date: 'TEXT',\n # datetime.time: 'TEXT',\n None.__class__: 'TEXT',\n # SQLite explicit types\n 'TEXT': 'TEXT',\n 'INTEGER': 'INTEGER',\n 'FLOAT': 'FLOAT',\n 'BLOB': 'BLOB',\n 'text': 'TEXT',\n 'integer': 'INTEGER',\n 'float': 'FLOAT',\n 'blob': 'BLOB',\n}\n\n# If numpy is available, add more types\nif np:\n COLUMN_TYPE_MAPPING.update(\n {\n np.int8: 'INTEGER',\n np.int16: 'INTEGER',\n np.int32: 'INTEGER',\n np.int64: 'INTEGER',\n np.uint8: 'INTEGER',\n np.uint16: 'INTEGER',\n np.uint32: 'INTEGER',\n np.uint64: 'INTEGER',\n np.float16: 'FLOAT',\n np.float32: 'FLOAT',\n np.float64: 'FLOAT',\n }\n )\n\n\ndef _converting(value: Any) -> str:\n if isinstance(value, bool):\n if value:\n return 1\n else:\n return 0\n\n return str(value)\n\n\ndef time_now():\n return datetime.datetime.utcnow()\n\n\ndef _get_table_names(\n conn: 'sqlite3.Connection', fts4: bool = False, fts5: bool = False\n) -> List[str]:\n \"\"\"A list of string table names in this database.\"\"\"\n where = [\"type = 'table'\"]\n if fts4:\n where.append(\"sql like '%USING FTS4%'\")\n if fts5:\n where.append(\"sql like '%USING FTS5%'\")\n sql = 'select name from sqlite_master where {}'.format(' AND '.join(where))\n return [r[0] for r in conn.execute(sql).fetchall()]\n\n\nclass Table:\n def __init__(\n self,\n name: str,\n data_path: Optional[Union[Path, str]] = None,\n detect_types: int = 0,\n in_memory: bool = True,\n ):\n if in_memory:\n self._conn_name = ':memory:'\n else:\n if isinstance(data_path, str):\n data_path = Path(data_path)\n self._conn_name = data_path / f'{name}.db'\n self._name = name\n\n self.detect_types = detect_types\n\n self._conn = sqlite3.connect(\n self._conn_name, detect_types=detect_types, check_same_thread=False\n )\n self._conn_lock = threading.Lock()\n\n def execute(self, sql: str, commit: bool = True):\n self._conn.execute(sql)\n if commit:\n self.commit()\n\n def execute_many(self, sql: str, parameters: List[Tuple], commit: bool = True):\n self._conn.executemany(sql, parameters)\n if commit:\n self.commit()\n\n def commit(self):\n self._conn.commit()\n\n def create_table(self):\n ...\n\n def drop_table(self):\n self._conn.execute(f'DROP table {self.name}')\n self._conn.commit()\n\n def clear(self):\n \"\"\"Drop the table and create a new one\"\"\"\n self.drop_table()\n self.create_table()\n\n def load(self, data_file: Union[str, Path]):\n disk_db = sqlite3.connect(data_file, detect_types=self.detect_types)\n disk_db.backup(self._conn)\n disk_db.close()\n\n def dump(self, data_file: Union[str, Path]):\n backup_db = sqlite3.connect(data_file, detect_types=self.detect_types)\n self._conn.backup(backup_db)\n backup_db.close()\n\n def close(self):\n self._conn.close()\n\n @property\n def name(self):\n return self._name\n\n @property\n def schema(self):\n \"\"\"SQL schema for this database\"\"\"\n result = []\n for row in self._conn.execute(\n f'''PRAGMA table_info(\"{self.name}\")'''\n ).fetchall():\n result.append(', '.join([str(_) for _ in row]))\n return '\\n'.join(result)\n\n\nclass CellTable(Table):\n def __init__(\n self,\n name: str,\n columns: Optional[List[tuple]] = None,\n in_memory: bool = True,\n data_path: Optional[Path] = None,\n lazy_create: bool = False,\n ):\n super().__init__(name, data_path=data_path, in_memory=in_memory)\n\n self._columns = []\n self._indexed_keys = set()\n\n if columns is not None:\n for name, dtype in columns:\n self.add_column(name, dtype, True)\n if not lazy_create:\n self.create_table()\n\n @property\n def columns(self) -> List[str]:\n return ['_id', '_doc_id'] + [c.split()[0] for c in self._columns]\n\n def existed(self):\n return self.name in _get_table_names(self._conn)\n\n def add_column(self, name: str, dtype: str, create_index: bool = True):\n self._columns.append(f'{name} {COLUMN_TYPE_MAPPING[dtype]}')\n if create_index:\n self._indexed_keys.add(name)\n\n def create_index(self, column: str, commit: bool = True):\n sql_statement = f'''CREATE INDEX idx_{column}_\n ON {self.name}({column})'''\n self._conn.execute(sql_statement)\n\n if commit:\n self._conn.commit()\n\n def create_table(self):\n sql = f'''CREATE TABLE {self.name}\n (_id INTEGER PRIMARY KEY AUTOINCREMENT,\n _doc_id TEXT NOT NULL UNIQUE'''\n if len(self._columns) > 0:\n sql += ', ' + ', '.join(self._columns)\n sql += ')'\n self._conn.execute(sql)\n\n for name in self._indexed_keys:\n self.create_index(name, commit=False)\n self._conn.commit()\n\n def insert(\n self,\n docs: 'DocumentArray',\n commit: bool = True,\n ) -> List[int]:\n \"\"\"Add a single record into the table.\n\n :param docs: The list of dict docs\n :param commit: If set, commit is applied\n \"\"\"\n sql_template = 'INSERT INTO {table}({columns}) VALUES ({placeholders});'\n\n column_names = self.columns[1:]\n columns = ', '.join(column_names)\n placeholders = ', '.join('?' for c in column_names)\n sql = sql_template.format(\n table=self.name, columns=columns, placeholders=placeholders\n )\n\n values = []\n docs_size = 0\n for doc in docs:\n doc_value = tuple(\n [doc.id]\n + [\n _converting(doc.tags[c]) if c in doc.tags else None\n for c in self.columns[2:]\n ]\n )\n values.append(doc_value)\n docs_size += 1\n\n with self._conn_lock:\n cursor = self._conn.cursor()\n if docs_size > 1:\n cursor.executemany(sql, values[:-1])\n\n cursor.execute(sql, values[-1])\n last_row_id = cursor.lastrowid\n row_ids = list(range(last_row_id - len(docs), last_row_id))\n\n if commit:\n self._conn.commit()\n\n return row_ids\n\n def query(\n self,\n where_clause: str = '',\n where_params: Tuple = (),\n limit: int = -1,\n offset: int = 0,\n order_by: Optional[str] = None,\n ascending: bool = True,\n ) -> List[int]:\n \"\"\"Query the records which matches the given conditions\n\n :param where_clause: where clause for query\n :param where_params: where parameters for query\n :param limit: limit the number of results\n :param offset: offset the number of results\n :param order_by: order the results by the given column\n :param ascending: order the results in ascending or descending order\n :return: offsets list of matched docs\n \"\"\"\n\n where_conds = []\n where = None\n if where_clause:\n where_conds.append(where_clause)\n where = ' and '.join(where_conds)\n\n _order_by = f'{order_by or \"_id\"} {\"ASC\" if ascending else \"DESC\"}'\n _limit = f'LIMIT {limit}' if limit > 0 else ''\n _offset = f'OFFSET {offset}' if offset > 0 else ''\n\n if where:\n sql = f'SELECT _id from {self.name} WHERE {where} ORDER BY {_order_by} {_limit} {_offset}'\n else:\n sql = f'SELECT _id from {self.name} ORDER BY {_order_by} {_limit} {_offset}'\n\n params = tuple([_converting(p) for p in where_params])\n\n # # EXPLAIN SQL query\n # for row in self._conn.execute('EXPLAIN QUERY PLAN ' + sql, params):\n # print(row)\n\n # Use `row_factor`\n # https://docs.python.org/3.6/library/sqlite3.html#sqlite3.Connection.row_factory\n def _offset_factory(_, record):\n return record[0] - 1\n\n self._conn.row_factory = _offset_factory\n\n cursor = self._conn.cursor()\n\n try:\n if where:\n offsets = cursor.execute(sql, params).fetchall()\n else:\n offsets = cursor.execute(sql).fetchall()\n self._conn.row_factory = None\n return offsets if offsets else []\n except Exception as e:\n self._conn.row_factory = None\n raise e\n\n def delete(self, doc_ids: List[str]):\n \"\"\"Delete the docs\n\n :param doc_ids: The IDs of docs\n \"\"\"\n sql = f'DELETE from {self.name} WHERE _doc_id = ?'\n self._conn.executemany(sql, doc_ids)\n self._conn.commit()\n\n def get_docid_by_offset(self, offset: int):\n sql = f'SELECT _doc_id from {self.name} WHERE _id = ? LIMIT 1;'\n result = self._conn.execute(sql, (offset + 1,)).fetchone()\n if result:\n return result[0]\n return None\n\n def delete_by_offset(self, offset: int):\n \"\"\"Delete the doc with specific offset\n\n :param offset: The offset of the doc\n \"\"\"\n sql = f'DELETE FROM {self.name} WHERE _id = ?'\n self._conn.execute(sql, (offset + 1,))\n self._conn.commit()\n\n def exist(self, doc_id: str):\n sql = f'SELECT count(*) from {self.name} WHERE _doc_id = ?;'\n return self._conn.execute(sql, (doc_id,)).fetchone()[0] > 0\n\n def count(self, where_clause: str = '', where_params: Tuple = ()):\n \"\"\"Return the total number of records which match with the given conditions.\n :param where_clause: where clause for query\n :param where_params: where parameters for query\n :return: the total number of matched records\n \"\"\"\n\n if where_clause:\n sql = 'SELECT count(_id) from {table} WHERE {where} LIMIT 1;'\n where = where_clause\n sql = sql.format(table=self.name, where=where)\n\n params = tuple([_converting(p) for p in where_params])\n\n # # EXPLAIN SQL query\n # for row in self._conn.execute('EXPLAIN QUERY PLAN ' + sql, params):\n # print(row)\n return self._conn.execute(sql, params).fetchone()[0]\n else:\n sql = f'SELECT count(_id) from {self.name};'\n result = self._conn.execute(sql).fetchone()\n if result[0]:\n return result[0]\n return 0\n\n @property\n def size(self):\n return self.count()\n\n\nclass MetaTable(Table):\n def __init__(\n self,\n name: str = 'meta',\n data_path: Optional[Path] = None,\n in_memory: bool = False,\n ):\n super().__init__(\n name,\n data_path=data_path,\n detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES,\n in_memory=in_memory,\n )\n self.create_table()\n\n def create_table(self):\n sql = f'''CREATE TABLE if not exists {self.name}\n (_doc_id TEXT NOT NULL PRIMARY KEY,\n cell_id INTEGER NOT NULL,\n offset INTEGER NOT NULL,\n time_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)'''\n\n self._conn.execute(sql)\n\n self._conn.execute(\n f'CREATE INDEX if not exists idx_time_at_ ON {self.name}(time_at)'\n )\n\n self._conn.commit()\n\n def iter_addresses(\n self, time_since: 'datetime.datetime' = datetime.datetime(2020, 2, 2, 0, 0)\n ):\n sql = f'SELECT _doc_id, cell_id, offset from {self.name} WHERE time_at >= ? ORDER BY time_at ASC;'\n\n cursor = self._conn.cursor()\n for doc_id, cell_id, offset in cursor.execute(sql, (time_since,)):\n yield doc_id, cell_id, offset\n\n def get_latest_commit(self):\n sql = f'SELECT _doc_id, cell_id, offset, time_at from {self.name} ORDER BY time_at DESC LIMIT 1;'\n\n cursor = self._conn.execute(sql)\n row = cursor.fetchone()\n return row\n\n def get_address(self, doc_id: str):\n sql = f'SELECT cell_id, offset from {self.name} WHERE _doc_id = ? LIMIT 1;'\n cursor = self._conn.execute(sql, (doc_id,))\n row = cursor.fetchone()\n return (row[0], row[1]) if row else (None, None)\n\n def delete_address(self, doc_id: str, commit: bool = True):\n sql = f'DELETE from {self.name} WHERE _doc_id = ?'\n self._conn.execute(sql, (doc_id,))\n if commit:\n self._conn.commit()\n\n def add_address(self, doc_id: str, cell_id: int, offset: int, commit: bool = True):\n sql = f'INSERT OR REPLACE INTO {self.name}(_doc_id, cell_id, offset, time_at) VALUES (?, ?, ?, ?);'\n self._conn.execute(\n sql,\n (doc_id, cell_id, offset, time_now()),\n )\n if commit:\n self._conn.commit()\n\n def bulk_add_address(\n self,\n doc_ids: List[str],\n cell_ids: Union[List[int], np.ndarray],\n offsets: Union[List[int], np.ndarray],\n commit: bool = True,\n ):\n sql = f'INSERT OR REPLACE INTO {self.name}(_doc_id, cell_id, offset, time_at) VALUES (?, ?, ?, ?);'\n self._conn.executemany(\n sql,\n [\n (doc_id, cell_id, offset, time_now())\n for doc_id, cell_id, offset in zip(doc_ids, cell_ids, offsets)\n ],\n )\n if commit:\n self._conn.commit()\n"
},
{
"path": "annlite/utils.py",
"content": "import os\nimport shutil\n\nimport numpy as np\nfrom docarray import Document, DocumentArray\n\n\ndef clean_workspace():\n if os.path.exists('./data'):\n shutil.rmtree('./data')\n\n if os.path.exists('./workspace'):\n shutil.rmtree('./workspace')\n\n\ndef docs_with_tags(N, D, probs, categories):\n\n all_docs = []\n start_current = 0\n for k, prob in enumerate(probs):\n n_current = int(N * prob)\n X = np.random.random((n_current, D)).astype(np.float32)\n\n docs = [\n Document(\n embedding=X[i],\n id=f'{i+start_current}',\n tags={\n 'category': categories[k],\n },\n )\n for i in range(n_current)\n ]\n all_docs.extend(docs)\n start_current += n_current\n\n return DocumentArray(all_docs)\n\n\ndef _precision(predicted, relevant, eval_at):\n \"\"\"\n fraction of retrieved documents that are relevant to the query\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n\n return n_predicted_and_relevant / len(predicted)\n\n\ndef _recall(predicted, relevant, eval_at):\n \"\"\"\n fraction of the relevant documents that are successfully retrieved\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n return n_predicted_and_relevant / len(relevant)\n\n\ndef evaluate(predicts, relevants, top_k):\n recall = 0\n precision = 0\n for _predict, _relevant in zip(predicts, relevants):\n _predict = np.array([int(x) for x in _predict])\n recall += _recall(_predict, _relevant, top_k)\n precision += _precision(_predict, _relevant, top_k)\n\n return recall / len(predicts), precision / len(predicts)\n"
},
{
"path": "benchmarks/filtering_bench.py",
"content": "import os\nimport shutil\nimport tempfile\n\nimport numpy as np\nfrom jina import Document, DocumentArray\nfrom jina.logging.profile import TimeContext\n\nfrom annlite import AnnLite\n\nn_index = [10_000, 100_000, 500_000, 1_000_000]\n\nn_query = [1, 8, 64]\nD = 768\nR = 5\nB = 5000\nn_cells = 1\nprobs = [[0.20, 0.30, 0.50], [0.05, 0.15, 0.80]]\ncategories = ['comic', 'movie', 'audiobook']\n\n\ndef docs_with_tags(N, D, probs, categories):\n\n all_docs = []\n for k, prob in enumerate(probs):\n n_current = int(N * prob)\n X = np.random.random((n_current, D)).astype(np.float32)\n\n docs = [\n Document(\n embedding=X[i],\n tags={\n 'category': categories[k],\n },\n )\n for i in range(n_current)\n ]\n all_docs.extend(docs)\n\n return DocumentArray(all_docs)\n\n\nresults = []\nfor n_i in n_index:\n\n results_ni = []\n for current_probs in probs:\n with tempfile.TemporaryDirectory() as tmpdir:\n columns = [('category', str)]\n idxer = AnnLite(\n D,\n initial_size=n_i,\n n_cells=n_cells,\n columns=columns,\n data_path=tmpdir,\n )\n\n da = docs_with_tags(n_i, D, current_probs, categories)\n\n with TimeContext(f'indexing {n_i} docs') as t_i:\n for i, _batch in enumerate(da.batch(batch_size=B)):\n idxer.index(_batch)\n\n for cat, prob in zip(categories, current_probs):\n f = {'category': {'$eq': cat}}\n\n query_times = []\n for n_q in n_query:\n qa = DocumentArray.empty(n_q)\n q_embs = np.random.random([n_q, D]).astype(np.float32)\n qa.embeddings = q_embs\n t_qs = []\n\n for _ in range(R):\n with TimeContext(f'searching {n_q} docs') as t_q:\n idxer.search(qa, filter=f)\n t_qs.append(t_q.duration)\n query_times.append(np.mean(t_qs[1:]))\n\n print(f'\\n\\nprob={prob}, current_probs={current_probs}, n_i={n_i}\\n\\n')\n results_ni.append([n_i, int(100 * prob), t_i.duration] + query_times)\n\n results.append(results_ni)\n\n\ntitle = '| Stored data |% same filter| Indexing time | Query size=1 | Query size=8 | Query size=64|'\nprint(title)\nprint('|-----' * 6 + '|')\nfor block in results:\n sorted_elements_in_block = np.argsort([b[1] for b in block])\n for pos in sorted_elements_in_block:\n res = block[pos]\n print(\n ''.join(\n [f'| {x} ' for x in res[0:2]] + [f'| {x:.3f} ' for x in res[2:]] + ['|']\n )\n )\n"
},
{
"path": "benchmarks/hnsw_bench.py",
"content": "import tempfile\nimport time\nfrom datetime import date\n\nimport numpy as np\nimport pandas as pd\nfrom docarray import Document, DocumentArray\nfrom sklearn.datasets import make_blobs\nfrom sklearn.model_selection import train_test_split\n\nfrom annlite import AnnLite\nfrom annlite.math import cdist\nfrom annlite.math import top_k as _top_k\n\n\ndef _precision(predicted, relevant, eval_at):\n \"\"\"\n fraction of retrieved documents that are relevant to the query\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n\n return n_predicted_and_relevant / len(predicted)\n\n\ndef _recall(predicted, relevant, eval_at):\n \"\"\"\n fraction of the relevant documents that are successfully retrieved\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n return n_predicted_and_relevant / len(relevant)\n\n\ndef evaluate(predicts, relevants, eval_at):\n recall = 0\n precision = 0\n for _predict, _relevant in zip(predicts, relevants):\n _predict = np.array([int(x) for x in _predict])\n recall += _recall(_predict, _relevant, top_k)\n precision += _precision(_predict, _relevant, top_k)\n\n return recall / len(predicts), precision / len(predicts)\n\n\n# N = 100_000 # number of data points\nNt = 125_000\nNq = 1\nD = 128 # dimentionality / number of features\ntop_k = 10\nn_cells = 64\nn_subvectors = 64\nn_queries = 1000\n\n# 2,000 128-dim vectors for training\nnp.random.seed(123)\nXtr, Xte = train_test_split(\n make_blobs(n_samples=Nt, n_features=D)[0].astype(np.float32), test_size=20\n)\nprint(f'Xtr: {Xtr.shape} vs Xte: {Xte.shape}')\n\n\ndef get_documents(nr=10, index_start=0, embeddings=None):\n for i in range(index_start, nr + index_start):\n d = Document()\n d.id = f'{i}' # to test it supports non-int ids\n d.embedding = embeddings[i - index_start]\n yield d\n\n\nprecision_per_query = []\nrecall_per_query = []\nresults = []\n\nfor n_cells in [1, 8, 16, 32, 64, 128]:\n\n with tempfile.TemporaryDirectory() as tmpdir:\n\n pq = AnnLite(\n D,\n metric='euclidean',\n n_cells=n_cells,\n data_path=tmpdir,\n )\n\n t0 = time.time()\n pq.train(Xtr[:20480])\n train_time = abs(time.time() - t0)\n\n t0 = time.time()\n pq.index(DocumentArray(get_documents(len(Xtr), embeddings=Xtr)))\n index_time = abs(t0 - time.time())\n\n dists = cdist(Xte, Xtr, metric='euclidean')\n true_dists, true_ids = _top_k(dists, top_k, descending=False)\n\n t0 = time.time()\n docs = DocumentArray(get_documents(len(Xte), embeddings=Xte))\n pq.search(docs, limit=top_k)\n\n query_time = abs(t0 - time.time())\n pq_ids = []\n for doc in docs:\n pq_ids.append([m.id for m in doc.matches])\n\n recall, precision = evaluate(pq_ids, true_ids, top_k)\n\n results_dict = {\n 'precision': precision,\n 'recall': recall,\n 'train_time': train_time,\n 'index_time': index_time,\n 'query_time': query_time,\n 'query_qps': len(Xte) / query_time,\n 'index_qps': len(Xtr) / index_time,\n 'indexer_hyperparams': {'n_cells': n_cells},\n }\n print(results_dict)\n\n results.append(results_dict)\n pq.clear()\n pq.close()\n\ntoday = date.today()\nresults_df = pd.DataFrame(results)\nresults_df.sort_values('recall', ascending=False)\nresults_df.to_csv(f'bench-results-{today.strftime(\"%b-%d-%Y\")}.csv')\n"
},
{
"path": "bindings/hnsw_bindings.cpp",
"content": "#include \"hnswlib.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace py = pybind11;\nusing namespace pybind11::literals; // needed to bring in _a literal\n\n/*\n * replacement for the openmp '#pragma omp parallel for' directive\n * only handles a subset of functionality (no reductions etc)\n * Process ids from start (inclusive) to end (EXCLUSIVE)\n *\n * The method is borrowed from nmslib\n */\ntemplate \ninline void ParallelFor(size_t start, size_t end, size_t numThreads,\n Function fn) {\n if (numThreads <= 0) {\n numThreads = std::thread::hardware_concurrency();\n }\n\n if (numThreads == 1) {\n for (size_t id = start; id < end; id++) {\n fn(id, 0);\n }\n } else {\n std::vector threads;\n std::atomic current(start);\n\n // keep track of exceptions in threads\n // https://stackoverflow.com/a/32428427/1713196\n std::exception_ptr lastException = nullptr;\n std::mutex lastExceptMutex;\n\n for (size_t threadId = 0; threadId < numThreads; ++threadId) {\n threads.push_back(std::thread([&, threadId] {\n while (true) {\n size_t id = current.fetch_add(1);\n\n if ((id >= end)) {\n break;\n }\n\n try {\n fn(id, threadId);\n } catch (...) {\n std::unique_lock lastExcepLock(lastExceptMutex);\n lastException = std::current_exception();\n /*\n * This will work even when current is the largest value that\n * size_t can fit, because fetch_add returns the previous value\n * before the increment (what will result in overflow\n * and produce 0 instead of current + 1).\n */\n current = end;\n break;\n }\n }\n }));\n }\n for (auto &thread : threads) {\n thread.join();\n }\n if (lastException) {\n std::rethrow_exception(lastException);\n }\n }\n}\n\ninline void assert_true(bool expr, const std::string &msg) {\n if (expr == false)\n throw std::runtime_error(\"Unpickle Error: \" + msg);\n return;\n}\n\ntemplate class Index {\npublic:\n Index(const std::string &space_name, const int dim)\n : space_name(space_name), dim(dim) {\n normalize = false;\n if (space_name == \"l2\") {\n l2space = new hnswlib::L2Space(dim);\n } else if (space_name == \"ip\") {\n l2space = new hnswlib::InnerProductSpace(dim);\n } else if (space_name == \"cosine\") {\n l2space = new hnswlib::InnerProductSpace(dim);\n normalize = true;\n } else {\n throw new std::runtime_error(\n \"Space name must be one of l2, ip, or cosine.\");\n }\n appr_alg = NULL;\n ep_added = true;\n index_inited = false;\n num_threads_default = std::thread::hardware_concurrency();\n\n default_ef = 10;\n pq_enable = false;\n pq_n_clusters = -1;\n pq_n_subvectors = -1;\n pq_d_subvector = -1;\n pq_codec = py::none();\n }\n\n static const int ser_version = 1; // serialization version\n\n std::string space_name;\n int dim;\n size_t seed;\n size_t default_ef;\n\n bool index_inited;\n bool ep_added;\n bool normalize;\n size_t maxElements, M, efConstruction;\n int num_threads_default;\n hnswlib::labeltype cur_l;\n hnswlib::HierarchicalNSW *appr_alg;\n hnswlib::SpaceInterface *l2space;\n\n // quantization setting\n bool pq_enable;\n size_t pq_n_subvectors;\n size_t pq_n_clusters;\n size_t pq_d_subvector;\n py::object pq_codec;\n\n ~Index() {\n delete l2space;\n if (appr_alg)\n delete appr_alg;\n // WARNING: will python release the pq_codec?\n }\n\n void init_new_index(const size_t maxElements, const size_t M,\n const size_t efConstruction, const size_t random_seed,\n const py::object &pq_codec) {\n if (appr_alg) {\n throw new std::runtime_error(\"The index is already initiated.\");\n }\n if (!pq_codec.is_none()) {\n _loadPQ(pq_codec);\n }\n this->cur_l = 0;\n this->appr_alg = new hnswlib::HierarchicalNSW(\n l2space, maxElements, M, efConstruction, random_seed);\n this->index_inited = true;\n this->ep_added = false;\n this->appr_alg->ef_ = default_ef;\n this->maxElements = maxElements;\n this->M = M;\n this->efConstruction = efConstruction;\n this->seed = random_seed;\n }\n\n void loadPQ(const py::object &pq_codec) {\n if (this->appr_alg) {\n delete this->appr_alg;\n }\n if (pq_codec.is_none()) {\n throw new std::runtime_error(\"Passed PQ class is none\");\n }\n _loadPQ(pq_codec);\n this->cur_l = 0;\n this->ep_added = false;\n this->appr_alg = new hnswlib::HierarchicalNSW(\n l2space, maxElements, M, efConstruction, seed);\n this->appr_alg->ef_ = default_ef;\n }\n\n void set_ef(size_t ef) {\n default_ef = ef;\n if (appr_alg)\n appr_alg->ef_ = ef;\n }\n\n void set_num_threads(int num_threads) {\n this->num_threads_default = num_threads;\n }\n\n void saveIndex(const std::string &path_to_index) {\n appr_alg->saveIndex(path_to_index);\n }\n\n void loadIndex(const std::string &path_to_index, size_t max_elements) {\n if (appr_alg) {\n std::cerr << \"Warning: Calling load_index for an already inited index. \"\n \"Old index is being deallocated.\";\n delete appr_alg;\n }\n appr_alg = new hnswlib::HierarchicalNSW(l2space, path_to_index,\n false, max_elements);\n cur_l = appr_alg->cur_element_count;\n index_inited = true;\n }\n\n void normalize_vector(float *data, float *norm_array) {\n float norm = 0.0f;\n for (int i = 0; i < dim; i++)\n norm += data[i] * data[i];\n norm = 1.0f / (sqrtf(norm) + 1e-30f);\n for (int i = 0; i < dim; i++)\n norm_array[i] = data[i] * norm;\n }\n\n template \n void addRows_(int dim, int num_threads, const py::object &ids_,\n const py::object &input, const py::object dtables) {\n\n py::array_t items(input);\n auto buffer = items.request();\n size_t rows, features;\n\n rows = buffer.shape[0];\n features = buffer.shape[1];\n\n if (!dtables.is_none()) {\n hnswlib::pq_local_data_t pq_param;\n py::array_t\n dtable_items(dtables);\n auto dtable_buffer = dtable_items.request();\n pq_param.data = (float *)dtable_buffer.ptr;\n pq_param.batch_len = rows;\n\n this->l2space->attach_local_data(&pq_param);\n }\n\n if (features != dim)\n throw std::runtime_error(\"wrong dimensionality of the vectors\");\n\n if (num_threads <= 0)\n num_threads = num_threads_default;\n // avoid using threads when the number of searches is small:\n if (rows <= num_threads * 4) {\n num_threads = 1;\n }\n\n std::vector ids;\n if (!ids_.is_none()) {\n py::array_t items(\n ids_);\n auto ids_numpy = items.request();\n if (ids_numpy.ndim == 1 && ids_numpy.shape[0] == rows) {\n std::vector ids1(ids_numpy.shape[0]);\n for (size_t i = 0; i < ids1.size(); i++) {\n ids1[i] = items.data()[i];\n }\n ids.swap(ids1);\n } else if (ids_numpy.ndim == 0 && rows == 1) {\n ids.push_back(*items.data());\n } else\n throw std::runtime_error(\"wrong dimensionality of the labels\");\n }\n\n {\n\n int start = 0;\n if (!ep_added) {\n size_t id = ids.size() ? ids.at(0) : (cur_l);\n appr_alg->addPoint((void *)items.data(0), (size_t)id, 0);\n start = 1;\n ep_added = true;\n }\n\n py::gil_scoped_release l;\n ParallelFor(start, rows, num_threads, [&](size_t row, size_t threadId) {\n size_t id = ids.size() ? ids.at(row) : (cur_l + row);\n appr_alg->addPoint((void *)items.data(row), (size_t)id, row);\n });\n cur_l += rows;\n }\n\n if (!dtables.is_none()) {\n this->l2space->detach_local_data();\n }\n }\n void addItems(const py::object &input, py::object ids_ = py::none(),\n int num_threads = -1, py::object dtables = py::none()) {\n // move dim check and normalization to python\n if (!pq_enable) {\n addRows_(this->dim, num_threads, ids_, input, dtables);\n } else {\n if (pq_n_clusters <= (UINT8_MAX + 1)) {\n addRows_(pq_n_subvectors, num_threads, ids_, input, dtables);\n } else if (pq_n_clusters <= (UINT16_MAX + 1)) {\n addRows_(pq_n_subvectors, num_threads, ids_, input, dtables);\n } else if (pq_n_clusters <= (UINT32_MAX + 1)) {\n addRows_(pq_n_subvectors, num_threads, ids_, input, dtables);\n }\n }\n }\n\n template \n py::object knnQuery_return_numpy_(size_t k, int num_threads,\n const py::object &input,\n const py::object dtables) {\n py::array_t items(input);\n auto buffer = items.request();\n hnswlib::labeltype *data_numpy_l;\n dist_t *data_numpy_d;\n size_t rows, features;\n {\n py::gil_scoped_release l;\n\n rows = buffer.shape[0];\n features = buffer.shape[1];\n\n if (!dtables.is_none()) {\n hnswlib::pq_local_data_t pq_param;\n py::array_t\n dtable_items(dtables);\n auto dtable_buffer = dtable_items.request();\n pq_param.data = (float *)dtable_buffer.ptr;\n pq_param.batch_len = rows;\n\n this->l2space->attach_local_data(&pq_param);\n }\n\n if (num_threads <= 0)\n num_threads = num_threads_default;\n\n // avoid using threads when the number of searches is small:\n if (rows <= num_threads * 4) {\n num_threads = 1;\n }\n\n data_numpy_l = new hnswlib::labeltype[rows * k];\n data_numpy_d = new dist_t[rows * k];\n\n ParallelFor(0, rows, num_threads, [&](size_t row, size_t threadId) {\n std::priority_queue> result =\n appr_alg->searchKnn((void *)items.data(row), k, row);\n if (result.size() != k)\n throw std::runtime_error(\n \"Cannot return the results in a contigious 2D array. Probably \"\n \"ef or M is too small\");\n for (int i = k - 1; i >= 0; i--) {\n auto &result_tuple = result.top();\n data_numpy_d[row * k + i] = result_tuple.first;\n data_numpy_l[row * k + i] = result_tuple.second;\n result.pop();\n }\n });\n }\n\n if (!dtables.is_none()) {\n this->l2space->detach_local_data();\n }\n py::capsule free_when_done_l(data_numpy_l, [](void *f) { delete[] f; });\n py::capsule free_when_done_d(data_numpy_d, [](void *f) { delete[] f; });\n\n return py::make_tuple(\n py::array_t(\n {rows, k}, // shape\n {k * sizeof(hnswlib::labeltype),\n sizeof(\n hnswlib::labeltype)}, // C-style contiguous strides for double\n data_numpy_l, // the data pointer\n free_when_done_l),\n py::array_t(\n {rows, k}, // shape\n {k * sizeof(dist_t),\n sizeof(dist_t)}, // C-style contiguous strides for double\n data_numpy_d, // the data pointer\n free_when_done_d));\n }\n py::object knnQuery_return_numpy(const py::object &input, size_t k = 1,\n int num_threads = -1,\n py::object dtables = py::none()) {\n // move dim check and normalization to python\n if (!pq_enable) {\n return knnQuery_return_numpy_(k, num_threads, input, dtables);\n } else {\n if (pq_n_clusters <= (UINT8_MAX + 1)) {\n return knnQuery_return_numpy_(k, num_threads, input, dtables);\n } else if (pq_n_clusters <= (UINT16_MAX + 1)) {\n return knnQuery_return_numpy_(k, num_threads, input, dtables);\n } else if (pq_n_clusters <= (UINT32_MAX + 1)) {\n return knnQuery_return_numpy_(k, num_threads, input, dtables);\n }\n }\n }\n template \n py::object knnQuery_with_filter_(size_t k, int num_threads,\n const py::object &candidate_ids_,\n const py::object &input,\n const py::object dtables) {\n py::array_t items(input);\n auto buffer = items.request();\n hnswlib::labeltype *data_numpy_l;\n dist_t *data_numpy_d;\n\n if (num_threads <= 0)\n num_threads = num_threads_default;\n\n size_t rows;\n size_t features;\n\n rows = buffer.shape[0];\n features = buffer.shape[1];\n\n if (!dtables.is_none()) {\n hnswlib::pq_local_data_t pq_param;\n py::array_t\n dtable_items(dtables);\n auto dtable_buffer = dtable_items.request();\n pq_param.data = (float *)dtable_buffer.ptr;\n pq_param.batch_len = rows;\n\n this->l2space->attach_local_data(&pq_param);\n }\n // avoid using threads when the number of searches is small:\n\n if (rows <= num_threads * 4) {\n num_threads = 1;\n }\n\n // FuseFilter constructing\n binary_fuse16_t filter(appr_alg->max_elements_);\n\n if (!candidate_ids_.is_none()) {\n py::array_t items(\n candidate_ids_);\n auto ids_numpy = items.request();\n\n if (ids_numpy.ndim == 1) {\n const size_t size = ids_numpy.shape[0];\n std::vector big_set;\n big_set.reserve(size);\n for (size_t i = 0; i < size; i++) {\n big_set[i] = items.data()[i]; // we use contiguous values\n }\n\n if (!binary_fuse16_populate(big_set.data(), size, &filter)) {\n throw std::runtime_error(\"failure to populate the fuse filter\");\n }\n } else\n throw std::runtime_error(\"wrong dimensionality of the filter labels\");\n }\n\n {\n py::gil_scoped_release l;\n\n // would like to check the ownership of this data in more detail\n data_numpy_l = new hnswlib::labeltype[rows * k];\n data_numpy_d = new dist_t[rows * k];\n\n ParallelFor(0, rows, num_threads, [&](size_t row, size_t threadId) {\n std::priority_queue> result =\n appr_alg->searchKnnWithFilter((void *)items.data(row), &filter, k,\n row);\n if (result.size() != k)\n throw std::runtime_error(\n \"Cannot return the results in a contigious 2D array. Probably \"\n \"ef or M is too small\");\n for (int i = k - 1; i >= 0; i--) {\n auto &result_tuple = result.top();\n data_numpy_d[row * k + i] = result_tuple.first;\n data_numpy_l[row * k + i] = result_tuple.second;\n result.pop();\n }\n });\n }\n\n if (!dtables.is_none()) {\n this->l2space->detach_local_data();\n }\n\n py::capsule free_when_done_l(data_numpy_l, [](void *f) { delete[] f; });\n py::capsule free_when_done_d(data_numpy_d, [](void *f) { delete[] f; });\n\n return py::make_tuple(\n py::array_t(\n {rows, k}, // shape\n {k * sizeof(hnswlib::labeltype),\n sizeof(\n hnswlib::labeltype)}, // C-style contiguous strides for double\n data_numpy_l, // the data pointer\n free_when_done_l),\n py::array_t(\n {rows, k}, // shape\n {k * sizeof(dist_t),\n sizeof(dist_t)}, // C-style contiguous strides for double\n data_numpy_d, // the data pointer\n free_when_done_d));\n }\n py::object knnQuery_with_filter(py::object input,\n py::object candidate_ids_ = py::none(),\n size_t k = 1, int num_threads = -1,\n py::object dtables = py::none()) {\n // move dim check and normalization to python\n if (!pq_enable) {\n return knnQuery_with_filter_(k, num_threads, candidate_ids_, input,\n dtables);\n } else {\n if (pq_n_clusters <= (UINT8_MAX + 1)) {\n return knnQuery_with_filter_(k, num_threads, candidate_ids_,\n input, dtables);\n } else if (pq_n_clusters <= (UINT16_MAX + 1)) {\n return knnQuery_with_filter_(k, num_threads, candidate_ids_,\n input, dtables);\n } else if (pq_n_clusters <= (UINT32_MAX + 1)) {\n return knnQuery_with_filter_(k, num_threads, candidate_ids_,\n input, dtables);\n }\n }\n }\n\n std::vector>\n getDataReturnList(py::object ids_ = py::none()) {\n std::vector ids;\n if (!ids_.is_none()) {\n py::array_t items(\n ids_);\n auto ids_numpy = items.request();\n std::vector ids1(ids_numpy.shape[0]);\n for (size_t i = 0; i < ids1.size(); i++) {\n ids1[i] = items.data()[i];\n }\n ids.swap(ids1);\n }\n\n std::vector> data;\n for (auto id : ids) {\n data.push_back(appr_alg->template getDataByLabel(id));\n }\n return data;\n }\n\n std::vector getIdsList() {\n\n std::vector ids;\n\n for (auto kv : appr_alg->label_lookup_) {\n ids.push_back(kv.first);\n }\n return ids;\n }\n\n py::dict getAnnData() const { /* WARNING: Index::getAnnData is not thread-safe\n with Index::addItems */\n\n std::unique_lock templock(appr_alg->global);\n\n unsigned int level0_npy_size =\n appr_alg->cur_element_count * appr_alg->size_data_per_element_;\n unsigned int link_npy_size = 0;\n std::vector link_npy_offsets(appr_alg->cur_element_count);\n\n for (size_t i = 0; i < appr_alg->cur_element_count; i++) {\n unsigned int linkListSize =\n appr_alg->element_levels_[i] > 0\n ? appr_alg->size_links_per_element_ * appr_alg->element_levels_[i]\n : 0;\n link_npy_offsets[i] = link_npy_size;\n if (linkListSize)\n link_npy_size += linkListSize;\n }\n\n char *data_level0_npy = (char *)malloc(level0_npy_size);\n char *link_list_npy = (char *)malloc(link_npy_size);\n int *element_levels_npy =\n (int *)malloc(appr_alg->element_levels_.size() * sizeof(int));\n\n hnswlib::labeltype *label_lookup_key_npy = (hnswlib::labeltype *)malloc(\n appr_alg->label_lookup_.size() * sizeof(hnswlib::labeltype));\n hnswlib::tableint *label_lookup_val_npy = (hnswlib::tableint *)malloc(\n appr_alg->label_lookup_.size() * sizeof(hnswlib::tableint));\n\n memset(label_lookup_key_npy, -1,\n appr_alg->label_lookup_.size() * sizeof(hnswlib::labeltype));\n memset(label_lookup_val_npy, -1,\n appr_alg->label_lookup_.size() * sizeof(hnswlib::tableint));\n\n size_t idx = 0;\n for (auto it = appr_alg->label_lookup_.begin();\n it != appr_alg->label_lookup_.end(); ++it) {\n label_lookup_key_npy[idx] = it->first;\n label_lookup_val_npy[idx] = it->second;\n idx++;\n }\n\n memset(link_list_npy, 0, link_npy_size);\n\n memcpy(data_level0_npy, appr_alg->data_level0_memory_, level0_npy_size);\n memcpy(element_levels_npy, appr_alg->element_levels_.data(),\n appr_alg->element_levels_.size() * sizeof(int));\n\n for (size_t i = 0; i < appr_alg->cur_element_count; i++) {\n unsigned int linkListSize =\n appr_alg->element_levels_[i] > 0\n ? appr_alg->size_links_per_element_ * appr_alg->element_levels_[i]\n : 0;\n if (linkListSize) {\n memcpy(link_list_npy + link_npy_offsets[i], appr_alg->linkLists_[i],\n linkListSize);\n }\n }\n\n py::capsule free_when_done_l0(data_level0_npy, [](void *f) { delete[] f; });\n py::capsule free_when_done_lvl(element_levels_npy,\n [](void *f) { delete[] f; });\n py::capsule free_when_done_lb(label_lookup_key_npy,\n [](void *f) { delete[] f; });\n py::capsule free_when_done_id(label_lookup_val_npy,\n [](void *f) { delete[] f; });\n py::capsule free_when_done_ll(link_list_npy, [](void *f) { delete[] f; });\n\n /* TODO: serialize state of random generators appr_alg->level_generator_\n * and appr_alg->update_probability_generator_ */\n /* for full reproducibility / to avoid re-initializing generators\n * inside Index::createFromParams */\n\n return py::dict(\n \"offset_level0\"_a = appr_alg->offsetLevel0_,\n \"max_elements\"_a = appr_alg->max_elements_,\n \"cur_element_count\"_a = appr_alg->cur_element_count,\n \"size_data_per_element\"_a = appr_alg->size_data_per_element_,\n \"label_offset\"_a = appr_alg->label_offset_,\n \"offset_data\"_a = appr_alg->offsetData_,\n \"max_level\"_a = appr_alg->maxlevel_,\n \"enterpoint_node\"_a = appr_alg->enterpoint_node_,\n \"max_M\"_a = appr_alg->maxM_, \"max_M0\"_a = appr_alg->maxM0_,\n \"M\"_a = appr_alg->M_, \"mult\"_a = appr_alg->mult_,\n \"ef_construction\"_a = appr_alg->ef_construction_,\n \"ef\"_a = appr_alg->ef_, \"has_deletions\"_a = appr_alg->has_deletions_,\n \"size_links_per_element\"_a = appr_alg->size_links_per_element_,\n\n \"label_lookup_external\"_a = py::array_t(\n {appr_alg->label_lookup_.size()}, // shape\n {sizeof(\n hnswlib::labeltype)}, // C-style contiguous strides for double\n label_lookup_key_npy, // the data pointer\n free_when_done_lb),\n\n \"label_lookup_internal\"_a = py::array_t(\n {appr_alg->label_lookup_.size()}, // shape\n {sizeof(\n hnswlib::tableint)}, // C-style contiguous strides for double\n label_lookup_val_npy, // the data pointer\n free_when_done_id),\n\n \"element_levels\"_a = py::array_t(\n {appr_alg->element_levels_.size()}, // shape\n {sizeof(int)}, // C-style contiguous strides for double\n element_levels_npy, // the data pointer\n free_when_done_lvl),\n\n // linkLists_,element_levels_,data_level0_memory_\n \"data_level0\"_a = py::array_t(\n {level0_npy_size}, // shape\n {sizeof(char)}, // C-style contiguous strides for double\n data_level0_npy, // the data pointer\n free_when_done_l0),\n\n \"link_lists\"_a = py::array_t(\n {link_npy_size}, // shape\n {sizeof(char)}, // C-style contiguous strides for double\n link_list_npy, // the data pointer\n free_when_done_ll)\n\n );\n }\n\n py::dict getIndexParams() const { /* WARNING: Index::getAnnData is not\n thread-safe with Index::addItems */\n auto params = py::dict(\n \"ser_version\"_a =\n py::int_(Index::ser_version), // serialization version\n \"space\"_a = space_name, \"dim\"_a = dim, \"index_inited\"_a = index_inited,\n \"ep_added\"_a = ep_added, \"normalize\"_a = normalize,\n \"num_threads\"_a = num_threads_default, \"seed\"_a = seed);\n\n if (index_inited == false)\n return py::dict(**params, \"ef\"_a = default_ef);\n\n auto ann_params = getAnnData();\n\n return py::dict(**params, **ann_params);\n }\n\n static Index *createFromParams(const py::dict d) {\n\n // check serialization version\n assert_true(((int)py::int_(Index::ser_version)) >=\n d[\"ser_version\"].cast(),\n \"Invalid serialization version!\");\n\n auto space_name_ = d[\"space\"].cast();\n auto dim_ = d[\"dim\"].cast();\n auto index_inited_ = d[\"index_inited\"].cast();\n\n Index *new_index = new Index(space_name_, dim_);\n\n /* TODO: deserialize state of random generators into\n * new_index->level_generator_ and new_index->update_probability_generator_\n */\n /* for full reproducibility / state of generators is serialized\n * inside Index::getIndexParams */\n new_index->seed = d[\"seed\"].cast();\n\n if (index_inited_) {\n new_index->appr_alg = new hnswlib::HierarchicalNSW(\n new_index->l2space, d[\"max_elements\"].cast(),\n d[\"M\"].cast(), d[\"ef_construction\"].cast(),\n new_index->seed);\n new_index->cur_l = d[\"cur_element_count\"].cast();\n }\n\n new_index->index_inited = index_inited_;\n new_index->ep_added = d[\"ep_added\"].cast();\n new_index->num_threads_default = d[\"num_threads\"].cast();\n new_index->default_ef = d[\"ef\"].cast();\n\n if (index_inited_)\n new_index->setAnnData(d);\n\n return new_index;\n }\n\n static Index *createFromIndex(const Index &index) {\n return createFromParams(index.getIndexParams());\n }\n\n void setAnnData(const py::dict d) { /* WARNING: Index::setAnnData is not\n thread-safe with Index::addItems */\n\n std::unique_lock templock(appr_alg->global);\n\n assert_true(appr_alg->offsetLevel0_ == d[\"offset_level0\"].cast(),\n \"Invalid value of offsetLevel0_ \");\n assert_true(appr_alg->max_elements_ == d[\"max_elements\"].cast(),\n \"Invalid value of max_elements_ \");\n\n appr_alg->cur_element_count = d[\"cur_element_count\"].cast();\n\n assert_true(appr_alg->size_data_per_element_ ==\n d[\"size_data_per_element\"].cast(),\n \"Invalid value of size_data_per_element_ \");\n assert_true(appr_alg->label_offset_ == d[\"label_offset\"].cast(),\n \"Invalid value of label_offset_ \");\n assert_true(appr_alg->offsetData_ == d[\"offset_data\"].cast(),\n \"Invalid value of offsetData_ \");\n\n appr_alg->maxlevel_ = d[\"max_level\"].cast();\n appr_alg->enterpoint_node_ = d[\"enterpoint_node\"].cast();\n\n assert_true(appr_alg->maxM_ == d[\"max_M\"].cast(),\n \"Invalid value of maxM_ \");\n assert_true(appr_alg->maxM0_ == d[\"max_M0\"].cast(),\n \"Invalid value of maxM0_ \");\n assert_true(appr_alg->M_ == d[\"M\"].cast(), \"Invalid value of M_ \");\n assert_true(appr_alg->mult_ == d[\"mult\"].cast(),\n \"Invalid value of mult_ \");\n assert_true(appr_alg->ef_construction_ ==\n d[\"ef_construction\"].cast(),\n \"Invalid value of ef_construction_ \");\n\n appr_alg->ef_ = d[\"ef\"].cast();\n appr_alg->has_deletions_ = d[\"has_deletions\"].cast();\n\n assert_true(appr_alg->size_links_per_element_ ==\n d[\"size_links_per_element\"].cast(),\n \"Invalid value of size_links_per_element_ \");\n\n auto label_lookup_key_npy =\n d[\"label_lookup_external\"]\n .cast>();\n auto label_lookup_val_npy =\n d[\"label_lookup_internal\"]\n .cast>();\n auto element_levels_npy =\n d[\"element_levels\"]\n .cast<\n py::array_t>();\n auto data_level0_npy =\n d[\"data_level0\"]\n .cast<\n py::array_t>();\n auto link_list_npy =\n d[\"link_lists\"]\n .cast<\n py::array_t>();\n\n for (size_t i = 0; i < appr_alg->cur_element_count; i++) {\n if (label_lookup_val_npy.data()[i] < 0) {\n throw std::runtime_error(\"internal id cannot be negative!\");\n } else {\n appr_alg->label_lookup_.insert(std::make_pair(\n label_lookup_key_npy.data()[i], label_lookup_val_npy.data()[i]));\n }\n }\n\n memcpy(appr_alg->element_levels_.data(), element_levels_npy.data(),\n element_levels_npy.nbytes());\n\n unsigned int link_npy_size = 0;\n std::vector link_npy_offsets(appr_alg->cur_element_count);\n\n for (size_t i = 0; i < appr_alg->cur_element_count; i++) {\n unsigned int linkListSize =\n appr_alg->element_levels_[i] > 0\n ? appr_alg->size_links_per_element_ * appr_alg->element_levels_[i]\n : 0;\n link_npy_offsets[i] = link_npy_size;\n if (linkListSize)\n link_npy_size += linkListSize;\n }\n\n memcpy(appr_alg->data_level0_memory_, data_level0_npy.data(),\n data_level0_npy.nbytes());\n\n for (size_t i = 0; i < appr_alg->max_elements_; i++) {\n unsigned int linkListSize =\n appr_alg->element_levels_[i] > 0\n ? appr_alg->size_links_per_element_ * appr_alg->element_levels_[i]\n : 0;\n if (linkListSize == 0) {\n appr_alg->linkLists_[i] = nullptr;\n } else {\n appr_alg->linkLists_[i] = (char *)malloc(linkListSize);\n if (appr_alg->linkLists_[i] == nullptr)\n throw std::runtime_error(\n \"Not enough memory: loadIndex failed to allocate linklist\");\n\n memcpy(appr_alg->linkLists_[i],\n link_list_npy.data() + link_npy_offsets[i], linkListSize);\n }\n }\n }\n\n void markDeleted(size_t label) { appr_alg->markDelete(label); }\n\n void resizeIndex(size_t new_size) { appr_alg->resizeIndex(new_size); }\n\n size_t getMaxElements() const { return appr_alg->max_elements_; }\n\n size_t getCurrentCount() const { return appr_alg->cur_element_count; }\n\n void _loadPQ(const py::object &pq_abstract) {\n int exist_attr = 1, attr_correct = 1;\n PyObject *pq_raw_ptr = pq_abstract.ptr();\n exist_attr *= PyObject_HasAttrString(pq_raw_ptr, \"encode\");\n exist_attr *= PyObject_HasAttrString(pq_raw_ptr, \"get_codebook\");\n exist_attr *= PyObject_HasAttrString(pq_raw_ptr, \"get_subspace_splitting\");\n if (exist_attr <= 0) {\n throw py::index_error(\n \"PQ class should at least have the following attributes:\\n\"\n \"(encode, get_codebook, get_subspace_splitting)\");\n }\n attr_correct *=\n PyMethod_Check(PyObject_GetAttrString(pq_raw_ptr, \"encode\"));\n attr_correct *=\n PyMethod_Check(PyObject_GetAttrString(pq_raw_ptr, \"get_codebook\"));\n attr_correct *= PyMethod_Check(\n PyObject_GetAttrString(pq_raw_ptr, \"get_subspace_splitting\"));\n if (attr_correct <= 0) {\n throw py::attribute_error(\n \"PQ class have at least one of the following attributes' type \"\n \"INCORRECT:\\n(encode: ,\\n codebook: \"\n \",\\n get_subspace_splitting: )\");\n }\n\n this->pq_enable = true;\n this->pq_codec = pq_abstract;\n\n py::tuple subspaces_param = pq_abstract.attr(\"get_subspace_splitting\")();\n this->pq_n_subvectors = py::cast(subspaces_param[0]);\n this->pq_n_clusters = py::cast(subspaces_param[1]);\n this->pq_d_subvector = py::cast(subspaces_param[2]);\n\n size_t pq_total_dims = (this->pq_n_subvectors * this->pq_d_subvector);\n if (this->dim != pq_total_dims) {\n throw py::value_error(\n \"Initialization Error, expect HNSW.dim == \"\n \"PQ.n_subvector*PQ.d_subvector, but got:\\n\"\n \"HNSW.dim =\" +\n std::to_string(this->dim) +\n \", PQ.n_subvector*PQ.d_subvector=\" + std::to_string(pq_total_dims));\n }\n // reading codebook into float buffer\n py::array_t items(\n pq_abstract.attr(\"get_codebook\")());\n auto buffer = items.request();\n if (buffer.ndim != 3 || buffer.shape[0] != pq_n_subvectors ||\n buffer.shape[1] != pq_n_clusters || buffer.shape[2] != pq_d_subvector) {\n py::print(\"Expect the codebook with shape (\", pq_n_subvectors,\n pq_n_clusters, pq_d_subvector, \"seq\"_a = \",\");\n py::print(\" but got shape \", buffer.shape);\n throw py::attribute_error(\n \"PQ class returning the codebook with wrong dimension\");\n }\n // std::shared_ptr codebook_buffer((float *)buffer.ptr);\n float *codebook_buffer = (float *)buffer.ptr;\n if (l2space) {\n delete l2space;\n }\n if (pq_n_clusters <= (UINT8_MAX + 1)) {\n l2space = new hnswlib::PQ_Space(space_name, pq_n_subvectors,\n pq_n_clusters, pq_d_subvector,\n codebook_buffer);\n } else if (pq_n_clusters <= (UINT16_MAX + 1)) {\n l2space = new hnswlib::PQ_Space(space_name, pq_n_subvectors,\n pq_n_clusters, pq_d_subvector,\n codebook_buffer);\n } else if (pq_n_clusters <= (UINT32_MAX + 1)) {\n l2space = new hnswlib::PQ_Space(space_name, pq_n_subvectors,\n pq_n_clusters, pq_d_subvector,\n codebook_buffer);\n } else {\n throw py::value_error(\n \"PQ clustering exceed the maximum, annlite set the maximum of \"\n \"clusters = \" +\n std::to_string(UINT16_MAX + 1) +\n \", but got PQ.n_clusters=\" + std::to_string(pq_n_clusters));\n }\n }\n};\n\nPYBIND11_PLUGIN(hnsw_bind) {\n py::module m(\"hnsw_bind\");\n\n py::class_>(m, \"Index\")\n .def(py::init(&Index::createFromParams), py::arg(\"params\"))\n /* WARNING: Index::createFromIndex is not thread-safe with Index::addItems\n */\n .def(py::init(&Index::createFromIndex), py::arg(\"index\"))\n .def(py::init(), py::arg(\"space\"),\n py::arg(\"dim\"))\n .def(\"init_index\", &Index::init_new_index, py::arg(\"max_elements\"),\n py::arg(\"M\") = 16, py::arg(\"ef_construction\") = 200,\n py::arg(\"random_seed\") = 100, py::arg(\"pq_codec\") = py::none())\n .def(\"knn_query\", &Index::knnQuery_return_numpy, py::arg(\"data\"),\n py::arg(\"k\") = 1, py::arg(\"num_threads\") = -1,\n py::arg(\"dtables\") = py::none())\n .def(\"knn_query_with_filter\", &Index::knnQuery_with_filter,\n py::arg(\"data\"), py::arg(\"filters\") = py::none(), py::arg(\"k\") = 1,\n py::arg(\"num_threads\") = -1, py::arg(\"dtables\") = py::none())\n .def(\"add_items\", &Index::addItems, py::arg(\"data\"),\n py::arg(\"ids\") = py::none(), py::arg(\"num_threads\") = -1,\n py::arg(\"dtables\") = py::none())\n .def(\"get_items\", &Index::getDataReturnList,\n py::arg(\"ids\") = py::none())\n .def(\"get_ids_list\", &Index::getIdsList)\n .def(\"set_ef\", &Index::set_ef, py::arg(\"ef\"))\n .def(\"set_num_threads\", &Index::set_num_threads,\n py::arg(\"num_threads\"))\n .def(\"save_index\", &Index::saveIndex, py::arg(\"path_to_index\"))\n .def(\"load_index\", &Index::loadIndex, py::arg(\"path_to_index\"),\n py::arg(\"max_elements\") = 0)\n .def(\"mark_deleted\", &Index::markDeleted, py::arg(\"label\"))\n .def(\"resize_index\", &Index::resizeIndex, py::arg(\"new_size\"))\n .def(\"get_max_elements\", &Index::getMaxElements)\n .def(\"get_current_count\", &Index::getCurrentCount)\n .def(\"loadPQ\", &Index::loadPQ)\n .def_readonly(\"space\", &Index::space_name)\n .def_readonly(\"dim\", &Index::dim)\n .def_readonly(\"pq_enable\", &Index::pq_enable)\n .def_readwrite(\"num_threads\", &Index::num_threads_default)\n .def_property(\n \"ef\",\n [](const Index &index) {\n return index.index_inited ? index.appr_alg->ef_ : index.default_ef;\n },\n [](Index &index, const size_t ef_) {\n index.default_ef = ef_;\n if (index.appr_alg)\n index.appr_alg->ef_ = ef_;\n })\n .def_property_readonly(\"max_elements\",\n [](const Index &index) {\n return index.index_inited\n ? index.appr_alg->max_elements_\n : 0;\n })\n .def_property_readonly(\"element_count\",\n [](const Index &index) {\n return index.index_inited\n ? index.appr_alg->cur_element_count\n : 0;\n })\n .def_property_readonly(\"ef_construction\",\n [](const Index &index) {\n return index.index_inited\n ? index.appr_alg->ef_construction_\n : 0;\n })\n .def_property_readonly(\"M\",\n [](const Index &index) {\n return index.index_inited ? index.appr_alg->M_\n : 0;\n })\n\n .def(py::pickle(\n [](const Index &ind) { // __getstate__\n return py::make_tuple(\n ind.getIndexParams()); /* Return dict (wrapped in a tuple) that\n fully encodes state of the Index\n object */\n },\n [](py::tuple t) { // __setstate__\n if (t.size() != 1)\n throw std::runtime_error(\"Invalid state!\");\n\n return Index::createFromParams(t[0].cast());\n }))\n\n .def(\"__repr__\", [](const Index &a) {\n return \"\";\n });\n\n return m.ptr();\n}\n"
},
{
"path": "bindings/pq_bindings.pyx",
"content": "# distutils: language = c++\n\nimport numpy as np\n\ncimport cython\nfrom libc.stdint cimport (\n int8_t,\n int16_t,\n int32_t,\n int64_t,\n uint8_t,\n uint16_t,\n uint32_t,\n uint64_t,\n)\nfrom libcpp.vector cimport vector\n\nctypedef fused any_int:\n uint8_t\n uint16_t\n uint32_t\n uint64_t\n int8_t\n int16_t\n int32_t\n int64_t\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ncdef inline dist_pqcode_to_codebook(long M,const float[:,:] adtable, any_int[:] pq_code):\n \"\"\"Compute the distance between each codevector and the pq_code of a query.\n\n :param M: Number of sub-vectors in the original feature space.\n :param adtable: 2D Memoryview[float] containing precomputed Asymmetric Distances.\n :param pq_code: 1D Memoriview[any_int] containing a pq code.\n\n :return: Distance between pq code and query according to the Asymmetric Distance table.\n\n \"\"\"\n cdef:\n float dist = 0\n int m\n\n for m in range(M):\n dist += adtable[m, pq_code[m]]\n\n return dist\n\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ncpdef dist_pqcodes_to_codebooks(const float[:,:] adtable, any_int[:,:] pq_codes):\n \"\"\"\n Compute the distance between each row in pq_codes and each codevector using a adtable.\n\n :param adtable: 2D Memoryview of precomputed Asymmetric Distances.\n :param pq_codes: 2D Memoryview of pq_codes.\n\n :return: List of Asymmetric Distances distances between pq_codes and the query.\n\n This function is equivalent to:\n '''\n dists = np.zeros((N, )).astype(np.float32)\n for n in range(N):\n for m in range(M):\n dists[n] += self.adtable[m][codes[n][m]]\n '''\n\n \"\"\"\n\n cdef:\n int m\n int N = pq_codes.shape[0]\n int M = pq_codes.shape[1]\n vector[float] dists\n\n for n in range(N):\n dists.push_back(dist_pqcode_to_codebook(M, adtable, pq_codes[n,:]))\n\n return dists\n\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ncpdef precompute_adc_table(const float[:] query,\n long d_subvector,\n long n_clusters,\n const float[:,:,:] codebooks):\n \"\"\"\n Compute the Asymmetric Distance Table between a query and a PQ space.\n\n :param query: Memoryview a query in the original feature space (not a pqcode).\n :param d_subvector: Number of dimensions in a subvector.\n :param n_clusters: Number of clusters per sub-space (number of prototypes per sub-space).\n :param codebooks: Memoryview containing the learned codevectors for each slice.\n This is a 3D view with (slice index, prototype index, vector values).\n\n :return: Memoryview with a 2D matrix containing the Asymmetric Distance Computation.\n\n This function is equivalent to\n '''\n def numpy_adc_table(query, n_subvectors, n_clusters, d_subvector, codebooks):\n adtable = np.empty((n_subvectors, n_clusters), dtype=np.float32)\n for m in range(n_subvectors):\n query_sub = query[m * d_subvector: (m + 1) * d_subvector]\n adtable[m, :] = np.linalg.norm(codebooks[m] - query_sub, axis=1) ** 2\n\n return adtable\n '''\n But avoids generating views and calling numpy functions.\n \"\"\"\n\n cdef:\n int D = len(query)\n int M = int(D/d_subvector)\n int n_subvectors = int(D/d_subvector)\n int m, i, k, ind_prototype, j\n float[:, ::1] adtable = np.empty((M, n_clusters), dtype=np.float32)\n float[:] query_subvec = np.empty(d_subvector, dtype=np.float32)\n float[:] query_subcodeword = np.empty(d_subvector, dtype=np.float32)\n float dist_subprototype_to_subquery, coord_j\n\n for m in range(n_subvectors):\n\n # load m'th subquery\n i = 0\n for k in range(m * d_subvector, (m + 1) * d_subvector):\n query_subvec[i] = query[k]\n i += 1\n\n for ind_prototype in range(n_clusters):\n\n # load prototype ind_prototype for the m'th subspace\n for i in range(d_subvector):\n query_subcodeword[i] = codebooks[m, ind_prototype, i]\n\n # compute the distance between subprototype and subquery\n dist_subprototype_to_subquery = 0.\n for j in range(d_subvector):\n coord_j = query_subcodeword[j] - query_subvec[j]\n dist_subprototype_to_subquery += coord_j * coord_j\n\n adtable[m, ind_prototype] = dist_subprototype_to_subquery\n\n return adtable\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ncpdef batch_precompute_adc_table(const float[:, :] queries,\n long d_subvector,\n long n_clusters,\n const float[:,:,:] codebooks):\n \"\"\"\n Compute the Asymmetric Distance Table between a query and a PQ space.\n\n :param query: Memoryview a query in the original feature space (not a pqcode).\n :param d_subvector: Number of dimensions in a subvector.\n :param n_clusters: Number of clusters per sub-space (number of prototypes per sub-space).\n :param codebooks: Memoryview containing the learned codevectors for each slice.\n This is a 3D view with (slice index, prototype index, vector values).\n\n :return: Memoryview with a 2D matrix containing the Asymmetric Distance Computation.\n\n This function is equivalent to\n '''\n def numpy_adc_table(query, n_subvectors, n_clusters, d_subvector, codebooks):\n adtable = np.empty((n_subvectors, n_clusters), dtype=np.float32)\n for m in range(n_subvectors):\n query_sub = query[m * d_subvector: (m + 1) * d_subvector]\n adtable[m, :] = np.linalg.norm(codebooks[m] - query_sub, axis=1) ** 2\n\n return adtable\n '''\n But avoids generating views and calling numpy functions.\n \"\"\"\n\n cdef:\n int N = queries.shape[0]\n int D = queries.shape[1]\n int n_subvectors = int(D/d_subvector)\n int m, i, k, ind_prototype, j\n float[:, :, :] adtable = np.empty((N, n_subvectors, n_clusters), dtype=np.float32)\n float[:] query_subvec = np.empty(d_subvector, dtype=np.float32)\n float[:] query_subcodeword = np.empty(d_subvector, dtype=np.float32)\n float dist_subprototype_to_subquery, coord_j\n\n for index in range(N):\n for m in range(n_subvectors):\n\n # load m'th subquery\n i = 0\n for k in range(m * d_subvector, (m + 1) * d_subvector):\n query_subvec[i] = queries[index, k]\n i += 1\n\n for ind_prototype in range(n_clusters):\n\n # load prototype ind_prototype for the m'th subspace\n for i in range(d_subvector):\n query_subcodeword[i] = codebooks[m, ind_prototype, i]\n\n # compute the distance between subprototype and subquery\n dist_subprototype_to_subquery = 0.\n for j in range(d_subvector):\n coord_j = query_subcodeword[j] - query_subvec[j]\n dist_subprototype_to_subquery += coord_j * coord_j\n\n adtable[index, m, ind_prototype] = dist_subprototype_to_subquery\n\n return adtable\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ncpdef batch_precompute_adc_table_ip(const float[:, :] queries,\n long d_subvector,\n long n_clusters,\n const float[:,:,:] codebooks):\n \"\"\"\n Compute the Asymmetric Distance Table between a query and a PQ space.\n\n :param query: Memoryview a query in the original feature space (not a pqcode).\n :param d_subvector: Number of dimensions in a subvector.\n :param n_clusters: Number of clusters per sub-space (number of prototypes per sub-space).\n :param codebooks: Memoryview containing the learned codevectors for each slice.\n This is a 3D view with (slice index, prototype index, vector values).\n\n :return: Memoryview with a 2D matrix containing the Asymmetric Distance Computation.\n\n This function is equivalent to\n '''\n def numpy_adc_table(query, n_subvectors, n_clusters, d_subvector, codebooks):\n adtable = np.empty((n_subvectors, n_clusters), dtype=np.float32)\n for m in range(n_subvectors):\n query_sub = query[m * d_subvector: (m + 1) * d_subvector]\n adtable[m, :] = np.linalg.norm(codebooks[m] - query_sub, axis=1) ** 2\n\n return adtable\n '''\n But avoids generating views and calling numpy functions.\n \"\"\"\n\n cdef:\n int N = queries.shape[0]\n int D = queries.shape[1]\n int n_subvectors = int(D/d_subvector)\n int m, i, k, ind_prototype, j\n float[:, :, :] adtable = np.empty((N, n_subvectors, n_clusters), dtype=np.float32)\n float[:] query_subvec = np.empty(d_subvector, dtype=np.float32)\n float[:] query_subcodeword = np.empty(d_subvector, dtype=np.float32)\n float dist_subprototype_to_subquery\n\n for index in range(N):\n for m in range(n_subvectors):\n\n # load m'th subquery\n i = 0\n for k in range(m * d_subvector, (m + 1) * d_subvector):\n query_subvec[i] = queries[index, k]\n i += 1\n\n for ind_prototype in range(n_clusters):\n\n # load prototype ind_prototype for the m'th subspace\n for i in range(d_subvector):\n query_subcodeword[i] = codebooks[m, ind_prototype, i]\n\n # compute the distance between subprototype and subquery\n dist_subprototype_to_subquery = 0.\n for j in range(d_subvector):\n dist_subprototype_to_subquery += (query_subcodeword[j] * query_subvec[j])\n\n adtable[index, m, ind_prototype] = dist_subprototype_to_subquery\n\n return adtable\n"
},
{
"path": "examples/annlite_vs_simpleindexer.py",
"content": "import os\nimport shutil\nimport tempfile\nimport time\n\nimport numpy as np\nimport pandas as pd\nfrom jina import Document, DocumentArray, Flow\nfrom sklearn.datasets import make_blobs\nfrom sklearn.model_selection import train_test_split\n\nfrom executor.executor import AnnLiteIndexer\n\nNq = 1\nD = 128\ntop_k = 10\nR = 5\nn_cells = 64\nn_subvectors = 64\nn_queries = 1\n\nBENCHMARK_SIMPLEINDEXER = False\nBENCHMARK_ANNLITE = True\n\n\ndef _precision(predicted, relevant, eval_at):\n \"\"\"\n fraction of retrieved documents that are relevant to the query\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n\n return n_predicted_and_relevant / len(predicted)\n\n\ndef _recall(predicted, relevant, eval_at):\n \"\"\"\n fraction of the relevant documents that are successfully retrieved\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n return n_predicted_and_relevant / len(relevant)\n\n\ndef evaluate(predicts, relevants, eval_at):\n recall = 0\n precision = 0\n for _predict, _relevant in zip(predicts, relevants):\n _predict = np.array([int(x) for x in _predict])\n recall += _recall(_predict, _relevant, top_k)\n precision += _precision(_predict, _relevant, top_k)\n\n return recall / len(predicts), precision / len(predicts)\n\n\ndef create_data(n_examples, D):\n np.random.seed(123)\n Xtr, Xte = train_test_split(\n make_blobs(n_samples=n_examples, n_features=D)[0].astype(np.float32),\n test_size=1,\n )\n return Xtr, Xte\n\n\ndef create_data_online(n_examples, D, batch_size):\n np.random.seed(123)\n num = 0\n while True:\n Xtr_batch = make_blobs(n_samples=batch_size, n_features=D)[0].astype(np.float32)\n yield DocumentArray([Document(embedding=x) for x in Xtr_batch])\n num += batch_size\n\n if num + batch_size >= n_examples:\n break\n\n if num < n_examples:\n Xtr_batch = make_blobs(n_samples=n_examples - num, n_features=D)[0].astype(\n np.float32\n )\n yield DocumentArray([Document(embedding=x) for x in Xtr_batch])\n\n\ndef create_test_data(D, Nq):\n np.random.seed(123)\n Xte = make_blobs(n_samples=Nq, n_features=D)[0].astype(np.float32)\n return DocumentArray([Document(embedding=x) for x in Xte])\n\n\nif BENCHMARK_SIMPLEINDEXER:\n\n ################ SimpleIndexer Benchmark BEGIN #################\n n_datasets = [10001, 50001, 200001, 400001]\n times = []\n\n for n_examples in n_datasets:\n time_taken = 0\n\n Xtr, Xte = create_data(n_examples, D)\n\n with tempfile.TemporaryDirectory() as tmpdir:\n\n f = Flow().add(\n uses='jinahub://SimpleIndexer',\n uses_with={'match_args': {'metric': 'euclidean', 'limit': 10}},\n workspace=tmpdir,\n )\n docs = [Document(id=f'{i}', embedding=Xtr[i]) for i in range(len(Xtr))]\n\n with f:\n resp = f.post(\n on='/index',\n inputs=docs,\n )\n\n with f:\n t0 = time.time()\n resp = f.post(\n on='/search',\n inputs=DocumentArray([Document(embedding=Xte[0])]),\n return_results=True,\n )\n time_taken = time.time() - t0\n\n times.append(time_taken)\n\n df = pd.DataFrame({'n_examples': n_datasets, 'times': times})\n df.to_csv('simpleindexer.csv')\n print(df)\n ################ SimpleIndexer Benchmark END #################\n\n\nif BENCHMARK_ANNLITE:\n\n ################ AnnLite Benchmark BEGIN ######################\n n_datasets = [10_000, 100_000, 500_000, 1_000_000, 10_000_000]\n # n_datasets = [10_000, 100_000]\n n_queries = [1, 8, 64]\n\n batch_size = 4096\n times = []\n\n results = {}\n for n_examples in n_datasets:\n print(f'\\n\\nWorking with n_examples={n_examples}\\n\\n')\n time_taken = 0\n\n with tempfile.TemporaryDirectory() as tmpdir:\n\n f = Flow().add(\n uses=AnnLiteIndexer,\n uses_with={\n 'n_dim': D,\n 'limit': 10,\n },\n workspace=tmpdir,\n )\n\n docs = create_data_online(n_examples, D, batch_size)\n\n results_current = {}\n with f:\n time_taken = 0\n for batch in docs:\n t0 = time.time()\n resp = f.post(on='/index', inputs=batch, request_size=10240)\n # This is done to avoid data creation time loaded in index time\n time_taken += time.time() - t0\n results_current['index_time'] = time_taken\n\n times_per_n_query = []\n with f:\n for n_query in n_queries:\n da_queries = create_test_data(D, n_query)\n t_qs = []\n for _ in range(R):\n t0 = time.time()\n resp = f.post(\n on='/search',\n inputs=da_queries,\n return_results=True,\n )\n time_taken = time.time() - t0\n t_qs.append(time_taken)\n # remove warm-up\n times_per_n_query.append(np.mean(t_qs[1:]))\n\n results_current['query_times'] = times_per_n_query\n print(f'==> query_times: {times_per_n_query}')\n df = pd.DataFrame({'results': results_current})\n df.to_csv(f'annlite_{n_examples}.csv')\n results[n_examples] = results_current\n\n df = pd.DataFrame(results)\n df.to_csv('annlite.csv')\n clean_workspace()\n ################ AnnLite Benchmark END #########################\n"
},
{
"path": "examples/filter_example.py",
"content": "import os\nimport random\nimport shutil\n\nimport numpy as np\nfrom jina import Document, DocumentArray\nfrom jina.logging.profile import TimeContext\n\nfrom annlite import AnnLite\n\nn_index = [10_000, 100_000, 500_000, 1_000_000]\nn_index = [100_000]\nn_query = [1, 8, 64]\nn_query = [1]\nD = 768\nR = 5\nB = 5000\nn_cells = 1\n# probs =[[0.20, 0.30, 0.50],\n# [0.05, 0.15, 0.80]]\ncategories = ['comic', 'movie', 'audiobook']\n\n\ndef clean_workspace():\n if os.path.exists('./data'):\n shutil.rmtree('./data')\n\n if os.path.exists('./workspace'):\n shutil.rmtree('./workspace')\n\n\ndef docs_with_tags(N, D, probs, categories):\n\n all_docs = []\n for k, prob in enumerate(probs):\n n_current = int(N * prob)\n X = np.random.random((n_current, D)).astype(np.float32)\n\n docs = [\n Document(\n embedding=X[i],\n tags={'category': categories[k], 'x': random.randint(0, 5)},\n )\n for i in range(n_current)\n ]\n all_docs.extend(docs)\n\n return DocumentArray(all_docs)\n\n\nresults = []\nfor n_i in n_index:\n clean_workspace()\n results_ni = []\n current_probs = [0.05, 0.15, 0.80]\n\n columns = [('category', str)]\n idxer = AnnLite(\n D,\n initial_size=n_i,\n n_cells=n_cells,\n data_path='./workspace',\n columns=columns,\n )\n\n da = docs_with_tags(n_i, D, current_probs, categories)\n\n with TimeContext(f'indexing {n_i} docs') as t_i:\n for i, _batch in enumerate(da.batch(batch_size=B)):\n idxer.index(_batch)\n\n for cat, prob in zip(categories, current_probs):\n f = {'category': {'$eq': cat}}\n\n query_times = []\n for n_q in n_query:\n qa = DocumentArray.empty(n_q)\n q_embs = np.random.random([n_q, D]).astype(np.float32)\n qa.embeddings = q_embs\n t_qs = []\n\n for _ in range(R):\n with TimeContext(f'searching {n_q} docs') as t_q:\n idxer.search(qa, filter=f)\n t_qs.append(t_q.duration)\n query_times.append(np.mean(t_qs[1:]))\n\n print(f'\\n\\nprob={prob}, current_probs={current_probs}, n_i={n_i}\\n\\n')\n results_ni.append([n_i, prob, t_i.duration] + query_times)\n\n results.append(results_ni)\n\n\ntitle = '| Stored data |% same filter| Indexing time | Query size=1 | Query size=8 | Query size=64|'\nprint(title)\nprint('|-----' * 6 + '|')\nfor block in results:\n sorted_elements_in_block = np.argsort([b[1] for b in block])\n for pos in sorted_elements_in_block:\n res = block[pos]\n print(''.join([f'| {x:.3f} ' for x in res] + ['|']))\n"
},
{
"path": "examples/hnsw_example.py",
"content": "import random\nimport tempfile\n\nimport numpy as np\nfrom docarray import Document, DocumentArray\n\nfrom annlite import AnnLite\n\nN = 1000 # number of data points\nNq = 5\nNt = 2000\nD = 128 # dimensionality / number of features\n\ndirpath = tempfile.mkdtemp()\n\nwith tempfile.TemporaryDirectory() as tmpdirname:\n\n index = AnnLite(\n D, columns=[('x', float)], data_path=tmpdirname, include_metadata=True\n )\n\n X = np.random.random((N, D)).astype(\n np.float32\n ) # 10,000 128-dim vectors to be indexed\n\n docs = DocumentArray(\n [\n Document(id=f'{i}', embedding=X[i], tags={'x': random.random()})\n for i in range(N)\n ]\n )\n index.index(docs)\n\n X = np.random.random((Nq, D)).astype(np.float32) # a 128-dim query vector\n query = DocumentArray([Document(embedding=X[i]) for i in range(5)])\n\n index.search(query, filter={'x': {'$lt': 0.2}}, limit=10, include_metadata=True)\n\n for m in query[0].matches:\n print(f'{m.scores[\"euclidean\"].value} -> x={m.tags[\"x\"]}')\n assert m.tags['x'] < 0.2\n\n print(f'====')\n\n index.search(query, filter={'x': {'$gte': 0.9}}, limit=10, include_metadata=True)\n\n for m in query[0].matches:\n print(f'{m.scores[\"euclidean\"].value} -> x={m.tags[\"x\"]}')\n assert m.tags['x'] >= 0.9\n\n#\n# print(f'{[m.scores[\"euclidean\"].value for m in query[0].matches]}')\n# for i in range(len(query[0].matches) - 1):\n# assert (\n# query[0].matches[i].scores['euclidean'].value\n# <= query[0].matches[i + 1].scores['euclidean'].value\n# )\n"
},
{
"path": "examples/pq_benchmark.py",
"content": "import time\nfrom datetime import date\n\nimport numpy as np\nimport pandas as pd\nfrom docarray import Document, DocumentArray\nfrom sklearn.datasets import make_blobs\nfrom sklearn.model_selection import train_test_split\nfrom utils import evaluate\n\nfrom annlite import AnnLite\nfrom annlite.math import cdist\nfrom annlite.math import top_k as _top_k\n\n# N = 100_000 # number of data points\nNt = 100_020\nNq = 1\nD = 128 # dimentionality / number of features\ntop_k = 10\nn_cells = 64\nn_subvectors = 64\nn_queries = 1000\n\n# 2,000 128-dim vectors for training\nnp.random.seed(123)\nXtr, Xte = train_test_split(\n make_blobs(n_samples=Nt, n_features=D)[0].astype(np.float32), test_size=20\n)\nprint(f'Xtr: {Xtr.shape} vs Xte: {Xte.shape}')\n\n\ndef get_documents(nr=10, index_start=0, embeddings=None):\n for i in range(index_start, nr + index_start):\n d = Document()\n d.id = f'{i}' # to test it supports non-int ids\n d.embedding = embeddings[i - index_start]\n yield d\n\n\nprecision_per_query = []\nrecall_per_query = []\nresults = []\n\nfor n_cells in [1, 4, 8]:\n for n_subvectors in [64, 128]:\n\n pq = AnnLite(D, metric='euclidean', n_cells=n_cells, n_subvectors=n_subvectors)\n\n t0 = time.time()\n pq.train(Xtr[:20480])\n train_time = abs(time.time() - t0)\n\n t0 = time.time()\n pq.index(DocumentArray(get_documents(len(Xtr), embeddings=Xtr)))\n index_time = abs(t0 - time.time())\n\n dists = cdist(Xte, Xtr, metric='euclidean')\n true_dists, true_ids = _top_k(dists, top_k, descending=False)\n\n docs = DocumentArray(get_documents(len(Xte), embeddings=Xte))\n t0 = time.time()\n pq.search(docs, limit=top_k)\n query_time = abs(t0 - time.time())\n\n pq_ids = []\n for doc in docs:\n pq_ids.append([m.id for m in doc.matches])\n\n recall, precision = evaluate(pq_ids, true_ids, top_k)\n\n results_dict = {\n 'precision': precision,\n 'recall': recall,\n 'train_time': train_time,\n 'index_time': index_time,\n 'query_time': query_time,\n 'query_qps': len(Xte) / query_time,\n 'index_qps': len(Xtr) / index_time,\n 'indexer_hyperparams': {'n_cells': n_cells, 'n_subvectors': n_subvectors},\n }\n print(results_dict)\n\n results.append(results_dict)\n pq.clear()\n pq.close()\n\ntoday = date.today()\nresults_df = pd.DataFrame(results)\nresults_df.sort_values('recall', ascending=False)\nresults_df.to_csv(f'bench-results-{today.strftime(\"%b-%d-%Y\")}.csv')\n"
},
{
"path": "examples/pqlinearscann_benchmark_with_filtering.py",
"content": "import numpy as np\nfrom docarray import Document, DocumentArray\nfrom docarray.math.distance import cdist\nfrom docarray.math.helper import top_k as _top_k\nfrom jina.logging.profile import TimeContext\nfrom utils import clean_workspace, docs_with_tags, evaluate\n\nfrom annlite import AnnLite\n\nn_index = [10_000, 100_000, 500_000, 1_000_000]\nn_query = [1, 8, 64]\nD = 768\nR = 5\nB = 100_000\nn_cells = 1\nprobs = [[0.20, 0.30, 0.50], [0.05, 0.15, 0.80]]\n\ncategories = ['comic', 'movie', 'audiobook']\n\ntop_k = 20\nn_cells = 1\nn_subvectors = D\n\nresults = []\nfor n_i in n_index:\n\n results_ni = []\n for current_probs in probs:\n\n clean_workspace()\n columns = [('category', str)]\n indexer = AnnLite(\n D,\n initial_size=n_i,\n n_subvectors=n_subvectors,\n n_cells=n_cells,\n metas={'workspace': './workspace'},\n columns=columns,\n )\n\n da = docs_with_tags(n_i, D, current_probs, categories)\n da_embeddings = da.embeddings\n\n with TimeContext(f'indexing {n_i} docs') as t_i:\n n_train_quantizer = min(n_i, 20_000)\n row_ids = np.random.choice(range(n_i), n_train_quantizer, replace=False)\n indexer.partial_train(da_embeddings[row_ids, :])\n\n indexer.build_codebook()\n for i, _batch in enumerate(da.batch(batch_size=B)):\n indexer.index(_batch)\n\n for cat, prob in zip(categories, current_probs):\n f = {'category': {'$eq': cat}}\n\n indices_cat = (\n np.array([t['category'] for t in da.get_attributes('tags')]) == cat\n )\n ids_indices_cat = np.array(\n [d.id for d in da if d.tags['category'] == cat], dtype='int'\n )\n\n da_embeddings_cat = da_embeddings[indices_cat, :]\n\n query_times = []\n for n_q in n_query:\n qa = DocumentArray.empty(n_q)\n q_embs = np.random.random([n_q, D]).astype(np.float32)\n qa.embeddings = q_embs\n t_qs = []\n\n for _ in range(R):\n with TimeContext(f'searching {n_q} docs') as t_q:\n indexer.search(qa, filter=f, limit=top_k)\n t_qs.append(t_q.duration)\n\n query_times.append(np.mean(t_qs[1:]))\n\n if n_q == 8:\n dists = cdist(q_embs, da_embeddings_cat, metric='euclidean')\n true_dists, true_local_ids = _top_k(dists, top_k, descending=False)\n\n # Note:x\n # `true_ids` are not really positions within the original data\n # but positions within the subset `da_embeddings_cat`\n # We need to go from these positions to the original positions\n true_ids = ids_indices_cat[true_local_ids]\n ids = []\n for doc in qa:\n ids.append([m.id for m in doc.matches])\n\n recall, precision = evaluate(ids, true_ids, top_k)\n\n print(\n f'\\n\\nprob={prob}, current_probs={current_probs}, n_i={n_i}, recall={recall}\\n\\n'\n )\n results_ni.append([n_i, prob, t_i.duration] + query_times + [recall])\n\n results.append(results_ni)\n\n\ntitle = '| Stored data |% same filter| Indexing time | Query size=1 | Query size=8 | Query size=64| Recall |'\nprint(title)\nprint('|-----' * 6 + '|')\nfor block in results:\n sorted_elements_in_block = np.argsort([b[1] for b in block])\n for pos in sorted_elements_in_block:\n res = block[pos]\n print(''.join([f'| {x:.3f} ' for x in res] + ['|']))\n"
},
{
"path": "examples/utils.py",
"content": "import os\nimport shutil\n\nimport numpy as np\nfrom docarray import Document, DocumentArray\n\n\ndef clean_workspace():\n if os.path.exists('./data'):\n shutil.rmtree('./data')\n\n if os.path.exists('./workspace'):\n shutil.rmtree('./workspace')\n\n\ndef docs_with_tags(N, D, probs, categories):\n\n all_docs = []\n start_current = 0\n for k, prob in enumerate(probs):\n n_current = int(N * prob)\n X = np.random.random((n_current, D)).astype(np.float32)\n\n docs = [\n Document(\n embedding=X[i],\n id=f'{i+start_current}',\n tags={\n 'category': categories[k],\n },\n )\n for i in range(n_current)\n ]\n all_docs.extend(docs)\n start_current += n_current\n\n return DocumentArray(all_docs)\n\n\ndef _precision(predicted, relevant, eval_at):\n \"\"\"\n fraction of retrieved documents that are relevant to the query\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n\n return n_predicted_and_relevant / len(predicted)\n\n\ndef _recall(predicted, relevant, eval_at):\n \"\"\"\n fraction of the relevant documents that are successfully retrieved\n \"\"\"\n if eval_at == 0:\n return 0.0\n predicted_at_k = predicted[:eval_at]\n n_predicted_and_relevant = len(set(predicted_at_k).intersection(set(relevant)))\n return n_predicted_and_relevant / len(relevant)\n\n\ndef evaluate(predicts, relevants, top_k):\n recall = 0\n precision = 0\n for _predict, _relevant in zip(predicts, relevants):\n _predict = np.array([int(x) for x in _predict])\n recall += _recall(_predict, _relevant, top_k)\n precision += _precision(_predict, _relevant, top_k)\n\n return recall / len(predicts), precision / len(predicts)\n"
},
{
"path": "executor/Dockerfile",
"content": "FROM jinaai/jina:3-py38-perf\n\nRUN apt-get update && apt-get install --no-install-recommends -y gcc g++ git \\\n && rm -rf /var/lib/apt/lists/*\n\nCOPY . /workspace\nWORKDIR /workspace\n\nRUN pip install -r requirements.txt --no-cache-dir\n\nENTRYPOINT [\"jina\", \"executor\", \"--uses\", \"config.yml\"]\n"
},
{
"path": "executor/README.md",
"content": "# AnnLiteIndexer\n\n`AnnLiteIndexer` uses the [AnnLite](https://github.com/jina-ai/annlite) class for indexing Jina `Document` objects.\nThe `AnnLite` class partitions the data into cells at index time, and instantiates a \"sub-indexer\" in each cell. Search is performed aggregating results retrieved from cells.\n\nThis indexer is recommended to be used when an application requires **search with filters** applied on `Document` tags.\nThe `filtering query language` is based on [MongoDB's query and projection operators](https://docs.mongodb.com/manual/reference/operator/query/). We currently support a subset of those selectors.\nThe tags filters can be combined with `$and` and `$or`:\n\n- `$eq` - Equal to (number, string)\n- `$ne` - Not equal to (number, string)\n- `$gt` - Greater than (number)\n- `$gte` - Greater than or equal to (number)\n- `$lt` - Less than (number)\n- `$lte` - Less than or equal to (number)\n- `$in` - Included in an array\n- `$nin` - Not included in an array\n\nFor example, we want to search for a product with a price no more than `50$`.\n```python\nindex.search(query, filter={\"price\": {\"$lte\": 50}})\n```\n\nMore example filter expresses\n\n- A Nike shoes with white color\n\n```JSON\n{\n \"brand\": {\"$eq\": \"Nike\"},\n \"category\": {\"$eq\": \"Shoes\"},\n \"color\": {\"$eq\": \"White\"}\n}\n```\n\nOr\n\n```JSON\n{\n \"$and\":\n {\n \"brand\": {\"$eq\": \"Nike\"},\n \"category\": {\"$eq\": \"Shoes\"},\n \"color\": {\"$eq\": \"White\"}\n }\n}\n```\n\n\n- A Nike shoes or price less than `100$`\n\n```JSON\n{\n \"$or\":\n {\n \"brand\": {\"$eq\": \"Nike\"},\n \"price\": {\"$lt\": 100}\n }\n}\n```\n\n## Performance\n\nOne can run `benchmark.py` to get a quick performance overview.\n\n|Stored data| Indexing time | Query size=1 | Query size=8 | Query size=64|\n|---|---|---|---|---|\n|10000 | 2.970 | 0.002 | 0.013 | 0.100|\n|100000 | 76.474 | 0.011 | 0.078 | 0.649|\n|500000 | 467.936 | 0.046 | 0.356 | 2.823|\n|1000000 | 1025.506 | 0.091 | 0.695 | 5.778|\n\n## Getting Started\n\nFor an in-depth overview of the features of AnnLite\nyou can follow along with one of the examples below:\n\n\n| Name | Link |\n|----------------------------------------------|---|\n| E-commerce product image search with AnnLite | [](https://colab.research.google.com/github/jina-ai/pqlite/blob/main/notebooks/fashion_product_search.ipynb)|\n\n## Quick Start\n\n`AnnLiteIndexer` stores `Document` objects at the `workspace` directory, specified under the [`metas`](https://docs.jina.ai/fundamentals/executor/executor-built-in-features/#meta-attributes) attribute.\n\n#### Example: Selecting items whose 'price' is less than 50\n\nIf documents have a tag `'price'` that stores floating point values this indexer allows searching documents with a filter, such as `price <= 50`.\n\n```python\ncolumns = [('price', 'float')]\n\nf = Flow().add(\n uses='jinahub://AnnLiteIndexer/latest',\n uses_with={\n 'dim': 256,\n 'columns': columns,\n 'metric': 'cosine'\n },\n uses_metas={'workspace': '/my/tmp_folder'},\n install_requirements=True\n)\n\nsearch_filter = {\"price\": {\"$lte\": 50}}\nwith f:\n f.post(on='/index', inputs=docs)\n query_res = f.post(on='/search',\n inputs=query_docs,\n return_results=True,\n parameters={'filter': search_filter})\n```\n\n## CRUD operations\n\nYou can perform all the usual operations on the respective endpoints\n\n- `/index` Add documents\n- `/search` Search with query documents.\n- `/update` Update documents\n- `/delete` Delete documents\n- `/clear` Clear the index\n- `/status` Return the status of index\n - `total_docs`: the total number of indexed documents\n - `dim`: the dimension of the embeddings\n - `metric`: the distance metric type\n - `is_trained`: whether the index is already trained\n"
},
{
"path": "executor/benchmark.py",
"content": "import tempfile\nimport time\n\nimport numpy as np\nfrom jina import DocumentArray\nfrom jina.logging.profile import TimeContext\n\nn_index = [10_000, 100_000, 500_000, 1_000_000]\n\nn_query = [1, 8, 64]\nD = 768\nR = 5\nB = 4096\nn_cells = 1\n\nfrom annlite.executor import AnnLiteIndexer\n\ntimes = {}\n\nfor n_i in n_index:\n with tempfile.TemporaryDirectory() as tempdir:\n idxer = AnnLiteIndexer(\n n_dim=D,\n initial_size=n_i,\n n_cells=n_cells,\n data_path=str(tempdir),\n )\n\n # build index docs\n i_embs = np.random.random([n_i, D]).astype(np.float32)\n\n if n_cells > 1:\n idxer._index.vq_codec.fit(i_embs)\n\n da = DocumentArray.empty(n_i)\n da.embeddings = i_embs\n\n with TimeContext(f'indexing {n_i} docs') as t_i:\n for _batch in da.batch(batch_size=B):\n idxer.index(_batch)\n\n times[n_i] = {}\n times[n_i]['index'] = t_i.duration\n\n # waiting for the index to be ready\n time.sleep(5)\n\n for n_q in n_query:\n q_embs = np.random.random([n_q, D]).astype(np.float32)\n qa = DocumentArray.empty(n_q)\n qa.embeddings = q_embs\n\n t_qs = []\n for _ in range(R):\n with TimeContext(f'searching {n_q} docs') as t_q:\n idxer.search(qa)\n t_qs.append(t_q.duration)\n # # check if it return the full doc\n # assert qa[0].matches\n # assert qa[0].matches.embeddings.shape\n times[n_i][f'query_{n_q}'] = np.mean(t_qs[1:]) # remove warm-up\n\n idxer.clear()\n idxer.close()\n\nprint('|Stored data| Indexing time | Query size=1 | Query size=8 | Query size=64|')\nprint('|---' * (len(list(times.values())[0]) + 1) + '|')\nfor k, v in times.items():\n s = ' | '.join(f'{v[vv]:.3f}' for vv in ['index', 'query_1', 'query_8', 'query_64'])\n print(f'|{k} | {s}|')\n"
},
{
"path": "executor/config.yml",
"content": "jtype: AnnLiteIndexer\npy_modules:\n - ./executor.py\nmetas:\n name: AnnLiteIndexer_v3\n description: A similarity search based on Annlite\n url: https://github.com/jina-ai/annlite\n keywords: [ann, similarity_search, indexer, pq, hnsw, pre-filtering]\n"
},
{
"path": "executor/executor.py",
"content": "import threading\nimport time\nimport traceback\nimport warnings\nfrom threading import Thread\nfrom typing import Dict, List, Optional, Tuple, Union\n\nfrom docarray import Document, DocumentArray\nfrom jina import Executor, requests\nfrom jina.logging.logger import JinaLogger\n\nINDEX_BATCH_SIZE = 1024\n\n\nclass AnnLiteIndexer(Executor):\n \"\"\"A simple indexer that wraps the AnnLite indexer and adds a simple interface for indexing and searching.\n\n :param n_dim: Dimensionality of vectors to index\n :param metric: Distance metric type. Can be 'euclidean', 'inner_product', or 'cosine'\n :param limit: Number of results to get for each query document in search\n :param n_components: Number of components to use for dimensionality reduction\n :param match_args: the arguments to `DocumentArray`'s match function\n :param data_path: the workspace of the AnnLiteIndexer but not support when shards > 1.\n :param ef_construction: The construction time/accuracy trade-off\n :param ef_search: The query time accuracy/speed trade-off\n :param max_connection: The maximum number of outgoing connections in the\n graph (the \"M\" parameter)\n :param include_metadata: If True, return the document metadata in response\n :param index_access_paths: Default traversal paths on docs\n (used for indexing, delete and update), e.g. '@r', '@c', '@r,c'\n :param search_access_paths: Default traversal paths on docs\n (used for search), e.g. '@r', '@c', '@r,c'\n :param columns: A list or dict of column names to index.\n :param dim: Deprecated, use n_dim instead\n \"\"\"\n\n def __init__(\n self,\n n_dim: int = 0,\n metric: str = 'cosine',\n limit: int = 10,\n n_components: Optional[int] = None,\n match_args: Optional[Dict] = None,\n data_path: Optional[str] = None,\n ef_construction: Optional[int] = None,\n ef_search: Optional[int] = None,\n max_connection: Optional[int] = None,\n include_metadata: bool = True,\n index_access_paths: str = '@r',\n search_access_paths: str = '@r',\n columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None,\n dim: int = None,\n *args,\n **kwargs,\n ):\n\n super().__init__(*args, **kwargs)\n self.logger = JinaLogger(self.__class__.__name__)\n\n n_dim = n_dim or dim\n if not n_dim:\n raise ValueError('Please specify the dimension of the vectors to index!')\n\n self.n_components = n_components\n self.metric = metric\n self.match_args = match_args or {}\n self.include_metadata = include_metadata\n if limit:\n self.match_args.update({'limit': limit})\n\n self.index_access_paths = index_access_paths\n if 'index_traversal_paths' in kwargs:\n warnings.warn(\n f'`index_traversal_paths` is deprecated. Use `index_access_paths` instead.'\n )\n self.index_access_paths = kwargs['index_traversal_paths']\n\n self.search_access_paths = search_access_paths\n if 'search_traversal_paths' in kwargs:\n warnings.warn(\n f'`search_traversal_paths` is deprecated. Use `search_access_paths` instead.'\n )\n self.search_access_paths = kwargs['search_traversal_paths']\n\n self._data_buffer = DocumentArray()\n self._index_batch_size = INDEX_BATCH_SIZE\n self._max_length_queue = 2 * self._index_batch_size\n self._index_lock = threading.Lock()\n\n self.logger = JinaLogger(getattr(self.metas, 'name', self.__class__.__name__))\n\n if getattr(self.runtime_args, 'shards', 1) > 1 and data_path:\n raise ValueError(\n '`data_path` is not supported when shards > 1, please use `workspace` instead'\n )\n\n config = {\n 'n_dim': n_dim,\n 'n_components': n_components,\n 'metric': metric,\n 'ef_construction': ef_construction,\n 'ef_search': ef_search,\n 'max_connection': max_connection,\n 'data_path': data_path or self.workspace or './workspace',\n 'columns': columns,\n }\n self._index = DocumentArray(storage='annlite', config=config)\n\n # start indexing thread in background to group indexing requests\n # together and perform batch indexing at once\n self._start_index_loop()\n\n @requests(on='/index')\n def index(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Index new documents\n\n :param docs: the Documents to index\n :param parameters: dictionary with options for indexing\n Keys accepted:\n - 'access_paths': traversal paths on docs, e.g. '@r', '@c', '@r,c'\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.index_access_paths)\n flat_docs = docs[access_paths]\n if len(flat_docs) == 0:\n return\n\n while len(self._data_buffer) >= self._max_length_queue:\n time.sleep(0.001)\n\n with self._index_lock:\n self._data_buffer.extend(flat_docs)\n\n def _start_index_loop(self):\n \"\"\"Start the indexing loop in background.\n\n This loop is responsible for batch indexing the documents in the buffer.\n \"\"\"\n\n def _index_loop():\n try:\n while True:\n # if the buffer is none, will break the loop\n if self._data_buffer is None:\n break\n\n # if the buffer is empty, will wait for new documents to be added\n if len(self._data_buffer) == 0:\n time.sleep(0.1) # sleep for 100ms\n continue\n\n # acquire the lock to prevent threading issues\n with self._index_lock:\n batch_docs = self._data_buffer.pop(\n range(\n self._index_batch_size\n if len(self._data_buffer) > self._index_batch_size\n else len(self._data_buffer)\n )\n )\n self._index.extend(batch_docs)\n self.logger.debug(f'indexing {len(batch_docs)} docs done...')\n except Exception as e:\n self.logger.error(traceback.format_exc())\n raise e\n\n self._index_thread = Thread(target=_index_loop, daemon=False)\n self._index_thread.start()\n\n @requests(on='/update')\n def update(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Update existing documents\n\n :param docs: the Documents to update\n :param parameters: dictionary with options for updating\n Keys accepted:\n - 'access_paths': traversal paths on docs, e.g. '@r', '@c', '@r,c'\n - 'raise_errors_on_not_found': if True, raise an error if a document is not found. Default is False.\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.index_access_paths)\n raise_errors_on_not_found = parameters.get('raise_errors_on_not_found', False)\n flat_docs = docs[access_paths]\n if len(flat_docs) == 0:\n return\n\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot update documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n for doc in flat_docs:\n try:\n self._index[doc.id] = doc\n except IndexError:\n if raise_errors_on_not_found:\n raise Exception(\n f'The document (id={doc.id}) cannot be updated as'\n f'it is not found in the index'\n )\n else:\n self.logger.warning(\n f'cannot update doc {doc.id} as it does not exist in storage'\n )\n\n @requests(on='/delete')\n def delete(self, parameters: dict = {}, **kwargs):\n \"\"\"Delete existing documents\n\n Delete entries from the index by id\n :param parameters: parameters to the request\n \"\"\"\n\n delete_ids = parameters.get('ids', [])\n if len(delete_ids) == 0:\n return\n\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot delete documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n\n del self._index[delete_ids]\n\n @requests(on='/search')\n def search(\n self, docs: Optional[DocumentArray] = None, parameters: dict = {}, **kwargs\n ):\n \"\"\"Perform a vector similarity search and retrieve Document matches\n\n Search can be performed with candidate filtering. Filters are a triplet (column,operator,value).\n More than a filter can be applied during search. Therefore, conditions for a filter are specified as a list triplets.\n Each triplet contains:\n\n - column: Column used to filter.\n - operator: Binary operation between two values. Some supported operators include `['>','<','=','<=','>=']`.\n - value: value used to compare a candidate.\n\n :param docs: the Documents to search with\n :param parameters: dictionary for parameters for the search operation\n Keys accepted:\n - 'access_paths' (str): traversal paths on docs, e.g. '@r', '@c', '@r,c'\n - 'filter' (dict): the filtering conditions on document tags\n - 'limit' (int): nr of matches to get per Document\n \"\"\"\n\n if not docs:\n return\n\n access_paths = parameters.get('access_paths', self.search_access_paths)\n flat_docs = docs[access_paths]\n match_args = (\n {**self.match_args, **parameters}\n if parameters is not None\n else self.match_args\n )\n\n with self._index_lock:\n # if len(self._data_buffer) > 0:\n # raise RuntimeError(\n # f'Cannot search documents while the pending documents in the buffer are not indexed yet. '\n # 'Please wait for the pending documents to be indexed.'\n # )\n\n flat_docs.match(self._index, **match_args)\n\n @requests(on='/backup')\n def backup(self, parameters: Optional[Dict] = {}, **kwargs):\n \"\"\"\n Backup data to local or remote.\n Use api of \n\n Keys accepted:\n - 'target' (str): the name of indexer you want to backup as\n \"\"\"\n\n target_name = parameters.get('target_name', None)\n token = parameters.get('token', None)\n if target_name:\n target_name = f'{target_name}_{self.runtime_args.shard_id}'\n with self._index_lock:\n if len(self._data_buffer) > 0:\n raise RuntimeError(\n f'Cannot backup documents while the pending documents in the buffer are not indexed yet. '\n 'Please wait for the pending documents to be indexed.'\n )\n self._index._annlite.backup(target_name, token)\n\n @requests(on='/restore')\n def restore(self, parameters: Optional[Dict] = {}, **kwargs):\n \"\"\"\n Restore data from local or remote.\n Use api of \n \"\"\"\n source_name = parameters.get('source_name', None)\n token = parameters.get('token', None)\n if source_name:\n source_name = f'{source_name}_{self.runtime_args.shard_id}'\n self._index._annlite.restore(source_name, token)\n\n @requests(on='/filter')\n def filter(self, parameters: Dict, **kwargs):\n \"\"\"\n Query documents from the indexer by the filter `query` object in parameters. The `query` object must follow the\n specifications in the `find` method of `DocumentArray` using annlite: https://docarray.jina.ai/fundamentals/documentarray/find/#filter-with-query-operators\n :param parameters: Dictionary to define the `filter` that you want to use.\n \"\"\"\n\n return self._index.find(parameters.get('filter', None))\n\n @requests(on='/fill_embedding')\n def fill_embedding(self, docs: DocumentArray, **kwargs):\n \"\"\"\n retrieve embedding of Documents by id\n :param docs: DocumentArray to search with\n \"\"\"\n\n for doc in docs:\n doc.embedding = self._index[doc.id].embedding\n\n @requests(on='/status')\n def status(self, **kwargs) -> DocumentArray:\n \"\"\"Return the document containing status information about the indexer.\n\n The status will contain information on the total number of indexed and deleted\n documents, and on the number of (searchable) documents currently in the index.\n \"\"\"\n\n status = Document(\n tags={\n 'appending_size': len(self._data_buffer),\n 'total_docs': len(self._index),\n 'index_size': len(self._index),\n }\n )\n return DocumentArray([status])\n\n def flush(self):\n \"\"\"Flush all the data in the buffer to the index\"\"\"\n while len(self._data_buffer) > 0:\n time.sleep(0.1)\n\n @requests(on='/clear')\n def clear(self, **kwargs):\n \"\"\"Clear the index of all entries.\"\"\"\n self.flush()\n\n with self._index_lock:\n self._data_buffer = None\n self._index_thread.join()\n\n self._data_buffer = DocumentArray()\n self._index.clear()\n\n self._start_index_loop()\n\n def close(self, **kwargs):\n \"\"\"Close the index.\"\"\"\n super().close()\n\n self.flush()\n\n # wait for the index thread to finish\n with self._index_lock:\n self._data_buffer = None\n self._index_thread.join()\n\n # WARNING: the commented code below hangs the close in pytest `pytest tests/test_*.py`\n # But don't know why. It works fine in `pytest tests/test_executor.py` and normal python execution\n del self._index\n"
},
{
"path": "executor/requirements.txt",
"content": "annlite\ncertifi\ndocarray\n"
},
{
"path": "include/hnswlib/bruteforce.h",
"content": "#pragma once\n#include \n#include \n#include \n#include \n\nnamespace hnswlib {\n template\n class BruteforceSearch : public AlgorithmInterface {\n public:\n BruteforceSearch(SpaceInterface *s) {\n\n }\n BruteforceSearch(SpaceInterface *s, const std::string &location) {\n loadIndex(location, s);\n }\n\n BruteforceSearch(SpaceInterface *s, size_t maxElements) {\n maxelements_ = maxElements;\n data_size_ = s->get_data_size();\n fstdistfunc_ = s->get_dist_func();\n dist_func_param_ = s->get_dist_func_param();\n size_per_element_ = data_size_ + sizeof(labeltype);\n data_ = (char *) malloc(maxElements * size_per_element_);\n if (data_ == nullptr)\n std::runtime_error(\"Not enough memory: BruteforceSearch failed to allocate data\");\n cur_element_count = 0;\n }\n\n ~BruteforceSearch() {\n free(data_);\n }\n\n char *data_;\n size_t maxelements_;\n size_t cur_element_count;\n size_t size_per_element_;\n\n size_t data_size_;\n DISTFUNC fstdistfunc_;\n void *dist_func_param_;\n std::mutex index_lock;\n\n std::unordered_map dict_external_to_internal;\n\n void addPoint(const void *datapoint, labeltype label) {\n\n int idx;\n {\n std::unique_lock lock(index_lock);\n\n\n\n auto search=dict_external_to_internal.find(label);\n if (search != dict_external_to_internal.end()) {\n idx=search->second;\n }\n else{\n if (cur_element_count >= maxelements_) {\n throw std::runtime_error(\"The number of elements exceeds the specified limit\\n\");\n }\n idx=cur_element_count;\n dict_external_to_internal[label] = idx;\n cur_element_count++;\n }\n }\n memcpy(data_ + size_per_element_ * idx + data_size_, &label, sizeof(labeltype));\n memcpy(data_ + size_per_element_ * idx, datapoint, data_size_);\n\n\n\n\n };\n\n void removePoint(labeltype cur_external) {\n size_t cur_c=dict_external_to_internal[cur_external];\n\n dict_external_to_internal.erase(cur_external);\n\n labeltype label=*((labeltype*)(data_ + size_per_element_ * (cur_element_count-1) + data_size_));\n dict_external_to_internal[label]=cur_c;\n memcpy(data_ + size_per_element_ * cur_c,\n data_ + size_per_element_ * (cur_element_count-1),\n data_size_+sizeof(labeltype));\n cur_element_count--;\n\n }\n\n\n std::priority_queue>\n searchKnn(const void *query_data, size_t k, size_t batch_index) const {\n std::priority_queue> topResults;\n if (cur_element_count == 0) return topResults;\n for (int i = 0; i < k; i++) {\n dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_, nullptr);\n topResults.push(std::pair(dist, *((labeltype *)(data_ + size_per_element_ * i +\n data_size_))));\n }\n dist_t lastdist = topResults.top().first;\n for (int i = k; i < cur_element_count; i++) {\n dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_, nullptr);\n if (dist <= lastdist) {\n topResults.push(std::pair(dist, *((labeltype *) (data_ + size_per_element_ * i +\n data_size_))));\n if (topResults.size() > k)\n topResults.pop();\n lastdist = topResults.top().first;\n }\n\n }\n return topResults;\n };\n\n void saveIndex(const std::string &location) {\n std::ofstream output(location, std::ios::binary);\n std::streampos position;\n\n writeBinaryPOD(output, maxelements_);\n writeBinaryPOD(output, size_per_element_);\n writeBinaryPOD(output, cur_element_count);\n\n output.write(data_, maxelements_ * size_per_element_);\n\n output.close();\n }\n\n void loadIndex(const std::string &location, SpaceInterface *s) {\n\n\n std::ifstream input(location, std::ios::binary);\n std::streampos position;\n\n readBinaryPOD(input, maxelements_);\n readBinaryPOD(input, size_per_element_);\n readBinaryPOD(input, cur_element_count);\n\n data_size_ = s->get_data_size();\n fstdistfunc_ = s->get_dist_func();\n dist_func_param_ = s->get_dist_func_param();\n size_per_element_ = data_size_ + sizeof(labeltype);\n data_ = (char *) malloc(maxelements_ * size_per_element_);\n if (data_ == nullptr)\n std::runtime_error(\"Not enough memory: loadIndex failed to allocate data\");\n\n input.read(data_, maxelements_ * size_per_element_);\n\n input.close();\n\n }\n\n };\n}\n"
},
{
"path": "include/hnswlib/fusefilter.h",
"content": "#ifndef BINARYFUSEFILTER_H\n#define BINARYFUSEFILTER_H\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#ifndef XOR_MAX_ITERATIONS\n#define XOR_MAX_ITERATIONS \\\n 100 // probabillity of success should always be > 0.5 so 100 iterations is\n // highly unlikely\n#endif\n#ifdef _MSC_VER\n#include \n#include \n#endif\n\n/**\n * We start with a few utilities.\n ***/\nstatic inline uint64_t binary_fuse_murmur64(uint64_t h) {\n h ^= h >> 33;\n h *= UINT64_C(0xff51afd7ed558ccd);\n h ^= h >> 33;\n h *= UINT64_C(0xc4ceb9fe1a85ec53);\n h ^= h >> 33;\n return h;\n}\nstatic inline uint64_t binary_fuse_mix_split(uint64_t key, uint64_t seed) {\n return binary_fuse_murmur64(key + seed);\n}\nstatic inline uint64_t binary_fuse_rotl64(uint64_t n, unsigned int c) {\n return (n << (c & 63)) | (n >> ((-c) & 63));\n}\nstatic inline uint32_t binary_fuse_reduce(uint32_t hash, uint32_t n) {\n // http://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/\n return (uint32_t)(((uint64_t)hash * n) >> 32);\n}\nstatic inline uint64_t binary_fuse8_fingerprint(uint64_t hash) {\n return hash ^ (hash >> 32);\n}\n\n/**\n * We need a decent random number generator.\n **/\n\n// returns random number, modifies the seed\nstatic inline uint64_t binary_fuse_rng_splitmix64(uint64_t *seed) {\n uint64_t z = (*seed += UINT64_C(0x9E3779B97F4A7C15));\n z = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9);\n z = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB);\n return z ^ (z >> 31);\n}\n\ntypedef struct binary_fuse8_s {\n uint64_t Seed;\n uint32_t SegmentLength;\n uint32_t SegmentLengthMask;\n uint32_t SegmentCount;\n uint32_t SegmentCountLength;\n uint32_t ArrayLength;\n uint8_t *Fingerprints;\n} binary_fuse8_t;\n\n#ifdef _MSC_VER\n// Windows programmers who target 32-bit platform may need help:\nstatic inline uint64_t binary_fuse_mulhi(uint64_t a, uint64_t b) { return __umulh(a, b); }\n#else\nstatic inline uint64_t binary_fuse_mulhi(uint64_t a, uint64_t b) {\n return ((__uint128_t)a * b) >> 64;\n}\n#endif\n\ntypedef struct binary_hashes_s {\n uint32_t h0;\n uint32_t h1;\n uint32_t h2;\n} binary_hashes_t;\n\nstatic inline binary_hashes_t binary_fuse8_hash_batch(uint64_t hash,\n const binary_fuse8_t *filter) {\n uint64_t hi = binary_fuse_mulhi(hash, filter->SegmentCountLength);\n binary_hashes_t ans;\n ans.h0 = (uint32_t)hi;\n ans.h1 = ans.h0 + filter->SegmentLength;\n ans.h2 = ans.h1 + filter->SegmentLength;\n ans.h1 ^= (uint32_t)(hash >> 18) & filter->SegmentLengthMask;\n ans.h2 ^= (uint32_t)(hash)&filter->SegmentLengthMask;\n return ans;\n}\nstatic inline uint32_t binary_fuse8_hash(int index, uint64_t hash,\n const binary_fuse8_t *filter) {\n uint64_t h = binary_fuse_mulhi(hash, filter->SegmentCountLength);\n h += index * filter->SegmentLength;\n // keep the lower 36 bits\n uint64_t hh = hash & ((1UL << 36) - 1);\n // index 0: right shift by 36; index 1: right shift by 18; index 2: no shift\n h ^= (size_t)((hh >> (36 - 18 * index)) & filter->SegmentLengthMask);\n return h;\n}\n\n// Report if the key is in the set, with false positive rate.\nstatic inline bool binary_fuse8_contain(uint64_t key,\n const binary_fuse8_t *filter) {\n uint64_t hash = binary_fuse_mix_split(key, filter->Seed);\n uint8_t f = binary_fuse8_fingerprint(hash);\n binary_hashes_t hashes = binary_fuse8_hash_batch(hash, filter);\n f ^= filter->Fingerprints[hashes.h0] ^ filter->Fingerprints[hashes.h1] ^\n filter->Fingerprints[hashes.h2];\n return f == 0;\n}\n\nstatic inline uint32_t binary_fuse_calculate_segment_length(uint32_t arity,\n uint32_t size) {\n // These parameters are very sensitive. Replacing 'floor' by 'round' can\n // substantially affect the construction time.\n if (arity == 3) {\n return ((uint32_t)1) << (int)(floor(log((double)(size)) / log(3.33) + 2.25));\n } else if (arity == 4) {\n return ((uint32_t)1) << (int)(floor(log((double)(size)) / log(2.91) - 0.5));\n } else {\n return 65536;\n }\n}\n\nstatic inline double binary_fuse8_max(double a, double b) {\n if (a < b) {\n return b;\n }\n return a;\n}\n\nstatic inline double binary_fuse_calculate_size_factor(uint32_t arity,\n uint32_t size) {\n if (arity == 3) {\n return binary_fuse8_max(1.125, 0.875 + 0.25 * log(1000000.0) / log((double)size));\n } else if (arity == 4) {\n return binary_fuse8_max(1.075, 0.77 + 0.305 * log(600000.0) / log((double)size));\n } else {\n return 2.0;\n }\n}\n\n// allocate enough capacity for a set containing up to 'size' elements\n// caller is responsible to call binary_fuse8_free(filter)\n// size should be at least 2.\nstatic inline bool binary_fuse8_allocate(uint32_t size,\n binary_fuse8_t *filter) {\n uint32_t arity = 3;\n filter->SegmentLength = size == 0 ? 4 : binary_fuse_calculate_segment_length(arity, size);\n if (filter->SegmentLength > 262144) {\n filter->SegmentLength = 262144;\n }\n filter->SegmentLengthMask = filter->SegmentLength - 1;\n double sizeFactor = binary_fuse_calculate_size_factor(arity, size);\n uint32_t capacity = size <= 1 ? 0 : (uint32_t)(round((double)size * sizeFactor));\n uint32_t initSegmentCount =\n (capacity + filter->SegmentLength - 1) / filter->SegmentLength -\n (arity - 1);\n filter->ArrayLength = (initSegmentCount + arity - 1) * filter->SegmentLength;\n filter->SegmentCount =\n (filter->ArrayLength + filter->SegmentLength - 1) / filter->SegmentLength;\n if (filter->SegmentCount <= arity - 1) {\n filter->SegmentCount = 1;\n } else {\n filter->SegmentCount = filter->SegmentCount - (arity - 1);\n }\n filter->ArrayLength =\n (filter->SegmentCount + arity - 1) * filter->SegmentLength;\n filter->SegmentCountLength = filter->SegmentCount * filter->SegmentLength;\n filter->Fingerprints = (uint8_t *)malloc(filter->ArrayLength);\n return filter->Fingerprints != NULL;\n}\n\n// report memory usage\nstatic inline size_t binary_fuse8_size_in_bytes(const binary_fuse8_t *filter) {\n return filter->ArrayLength * sizeof(uint8_t) + sizeof(binary_fuse8_t);\n}\n\n// release memory\nstatic inline void binary_fuse8_free(binary_fuse8_t *filter) {\n free(filter->Fingerprints);\n filter->Fingerprints = NULL;\n filter->Seed = 0;\n filter->SegmentLength = 0;\n filter->SegmentLengthMask = 0;\n filter->SegmentCount = 0;\n filter->SegmentCountLength = 0;\n filter->ArrayLength = 0;\n}\n\nstatic inline uint8_t binary_fuse_mod3(uint8_t x) {\n return x > 2 ? x - 3 : x;\n}\n\n// construct the filter, returns true on success, false on failure.\n// most likely, a failure is due to too high a memory usage\n// size is the number of keys\n// The caller is responsable for calling binary_fuse8_allocate(size,filter)\n// before. The caller is responsible to ensure that there are not too many duplicated\n// keys. The inner loop will run up to XOR_MAX_ITERATIONS times (default on\n// 100), it should never fail, except if there are many duplicated keys. If it fails,\n// a return value of false is provided.\n//\n//\n// If there are many duplicated keys and you do not want to remove them, you can first\n// sort your input, the algorithm will then work adequately.\nbool binary_fuse8_populate(const uint64_t *keys, uint32_t size,\n binary_fuse8_t *filter) {\n uint64_t rng_counter = 0x726b2b9d438b9d4d;\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n uint64_t *reverseOrder = (uint64_t *)calloc((size + 1), sizeof(uint64_t));\n uint32_t capacity = filter->ArrayLength;\n uint32_t *alone = (uint32_t *)malloc(capacity * sizeof(uint32_t));\n uint8_t *t2count = (uint8_t *)calloc(capacity, sizeof(uint8_t));\n uint8_t *reverseH = (uint8_t *)malloc(size * sizeof(uint8_t));\n uint64_t *t2hash = (uint64_t *)calloc(capacity, sizeof(uint64_t));\n\n uint32_t blockBits = 1;\n while (((uint32_t)1 << blockBits) < filter->SegmentCount) {\n blockBits += 1;\n }\n uint32_t block = ((uint32_t)1 << blockBits);\n uint32_t *startPos = (uint32_t *)malloc((1 << blockBits) * sizeof(uint32_t));\n uint32_t h012[5];\n\n if ((alone == NULL) || (t2count == NULL) || (reverseH == NULL) ||\n (t2hash == NULL) || (reverseOrder == NULL) || (startPos == NULL)) {\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return false;\n }\n reverseOrder[size] = 1;\n for (int loop = 0; true; ++loop) {\n if (loop + 1 > XOR_MAX_ITERATIONS) {\n fprintf(stderr, \"Too many iterations. Are all your keys unique?\");\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return false;\n }\n\n for (uint32_t i = 0; i < block; i++) {\n // important : i * size would overflow as a 32-bit number in some\n // cases.\n startPos[i] = ((uint64_t)i * size) >> blockBits;\n }\n\n uint64_t maskblock = block - 1;\n for (uint32_t i = 0; i < size; i++) {\n uint64_t hash = binary_fuse_murmur64(keys[i] + filter->Seed);\n uint64_t segment_index = hash >> (64 - blockBits);\n while (reverseOrder[startPos[segment_index]] != 0) {\n segment_index++;\n segment_index &= maskblock;\n }\n reverseOrder[startPos[segment_index]] = hash;\n startPos[segment_index]++;\n }\n int error = 0;\n uint32_t duplicates = 0;\n for (uint32_t i = 0; i < size; i++) {\n uint64_t hash = reverseOrder[i];\n uint32_t h0 = binary_fuse8_hash(0, hash, filter);\n t2count[h0] += 4;\n t2hash[h0] ^= hash;\n uint32_t h1 = binary_fuse8_hash(1, hash, filter);\n t2count[h1] += 4;\n t2count[h1] ^= 1;\n t2hash[h1] ^= hash;\n uint32_t h2 = binary_fuse8_hash(2, hash, filter);\n t2count[h2] += 4;\n t2hash[h2] ^= hash;\n t2count[h2] ^= 2;\n if ((t2hash[h0] & t2hash[h1] & t2hash[h2]) == 0) {\n if (((t2hash[h0] == 0) && (t2count[h0] == 8)) ||\n ((t2hash[h1] == 0) && (t2count[h1] == 8)) ||\n ((t2hash[h2] == 0) && (t2count[h2] == 8))) {\n duplicates += 1;\n t2count[h0] -= 4;\n t2hash[h0] ^= hash;\n t2count[h1] -= 4;\n t2count[h1] ^= 1;\n t2hash[h1] ^= hash;\n t2count[h2] -= 4;\n t2count[h2] ^= 2;\n t2hash[h2] ^= hash;\n }\n }\n error = (t2count[h0] < 4) ? 1 : error;\n error = (t2count[h1] < 4) ? 1 : error;\n error = (t2count[h2] < 4) ? 1 : error;\n }\n if (error) {\n memset(reverseOrder, 0, sizeof(uint64_t)*size);\n memset(t2count, 0, sizeof(uint8_t)*capacity);\n memset(t2hash, 0, sizeof(uint64_t)*capacity);\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n continue;\n }\n\n // End of key addition\n uint32_t Qsize = 0;\n // Add sets with one key to the queue.\n for (uint32_t i = 0; i < capacity; i++) {\n alone[Qsize] = i;\n Qsize += ((t2count[i] >> 2) == 1) ? 1 : 0;\n }\n uint32_t stacksize = 0;\n while (Qsize > 0) {\n Qsize--;\n uint32_t index = alone[Qsize];\n if ((t2count[index] >> 2) == 1) {\n uint64_t hash = t2hash[index];\n\n // h012[0] = binary_fuse8_hash(0, hash, filter);\n h012[1] = binary_fuse8_hash(1, hash, filter);\n h012[2] = binary_fuse8_hash(2, hash, filter);\n h012[3] = binary_fuse8_hash(0, hash, filter); // == h012[0];\n h012[4] = h012[1];\n uint8_t found = t2count[index] & 3;\n reverseH[stacksize] = found;\n reverseOrder[stacksize] = hash;\n stacksize++;\n uint32_t other_index1 = h012[found + 1];\n alone[Qsize] = other_index1;\n Qsize += ((t2count[other_index1] >> 2) == 2 ? 1 : 0);\n\n t2count[other_index1] -= 4;\n t2count[other_index1] ^= binary_fuse_mod3(found + 1);\n t2hash[other_index1] ^= hash;\n\n uint32_t other_index2 = h012[found + 2];\n alone[Qsize] = other_index2;\n Qsize += ((t2count[other_index2] >> 2) == 2 ? 1 : 0);\n t2count[other_index2] -= 4;\n t2count[other_index2] ^= binary_fuse_mod3(found + 2);\n t2hash[other_index2] ^= hash;\n }\n }\n if (stacksize + duplicates == size) {\n // success\n size = stacksize;\n break;\n }\n memset(reverseOrder, 0, sizeof(uint64_t)*size);\n memset(t2count, 0, sizeof(uint8_t)*capacity);\n memset(t2hash, 0, sizeof(uint64_t)*capacity);\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n }\n\n for (uint32_t i = size - 1; i < size; i--) {\n // the hash of the key we insert next\n uint64_t hash = reverseOrder[i];\n uint8_t xor2 = binary_fuse8_fingerprint(hash);\n uint8_t found = reverseH[i];\n h012[0] = binary_fuse8_hash(0, hash, filter);\n h012[1] = binary_fuse8_hash(1, hash, filter);\n h012[2] = binary_fuse8_hash(2, hash, filter);\n h012[3] = h012[0];\n h012[4] = h012[1];\n filter->Fingerprints[h012[found]] = xor2 ^\n filter->Fingerprints[h012[found + 1]] ^\n filter->Fingerprints[h012[found + 2]];\n }\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return true;\n}\n\n//////////////////\n// fuse16\n//////////////////\n\ntypedef struct binary_fuse16_s {\n uint64_t Seed;\n uint32_t SegmentLength;\n uint32_t SegmentLengthMask;\n uint32_t SegmentCount;\n uint32_t SegmentCountLength;\n uint32_t ArrayLength;\n uint16_t *Fingerprints;\n\n binary_fuse16_s(uint32_t size) {\n uint32_t arity = 3;\n SegmentLength =\n size == 0 ? 4 : binary_fuse_calculate_segment_length(arity, size);\n if (SegmentLength > 262144) {\n SegmentLength = 262144;\n }\n SegmentLengthMask = SegmentLength - 1;\n double sizeFactor =\n size <= 1 ? 0 : binary_fuse_calculate_size_factor(arity, size);\n uint32_t capacity = (uint32_t)(round((double)size * sizeFactor));\n uint32_t initSegmentCount =\n (capacity + SegmentLength - 1) / SegmentLength - (arity - 1);\n ArrayLength = (initSegmentCount + arity - 1) * SegmentLength;\n SegmentCount = (ArrayLength + SegmentLength - 1) / SegmentLength;\n if (SegmentCount <= arity - 1) {\n SegmentCount = 1;\n } else {\n SegmentCount = SegmentCount - (arity - 1);\n }\n ArrayLength = (SegmentCount + arity - 1) * SegmentLength;\n SegmentCountLength = SegmentCount * SegmentLength;\n Fingerprints = (uint16_t *)malloc(ArrayLength * sizeof(uint16_t));\n memset(Fingerprints, 0, ArrayLength * sizeof(uint16_t));\n if (Fingerprints == NULL) {\n throw std::runtime_error(\"not enough memory to hold the fuse filter\");\n }\n }\n\n ~binary_fuse16_s() {\n free(Fingerprints);\n Fingerprints = NULL;\n Seed = 0;\n SegmentLength = 0;\n SegmentLengthMask = 0;\n SegmentCount = 0;\n SegmentCountLength = 0;\n ArrayLength = 0;\n }\n} binary_fuse16_t;\n\nstatic inline uint64_t binary_fuse16_fingerprint(uint64_t hash) {\n return hash ^ (hash >> 32);\n}\n\nstatic inline binary_hashes_t binary_fuse16_hash_batch(uint64_t hash,\n const binary_fuse16_t *filter) {\n uint64_t hi = binary_fuse_mulhi(hash, filter->SegmentCountLength);\n binary_hashes_t ans;\n ans.h0 = (uint32_t)hi;\n ans.h1 = ans.h0 + filter->SegmentLength;\n ans.h2 = ans.h1 + filter->SegmentLength;\n ans.h1 ^= (uint32_t)(hash >> 18) & filter->SegmentLengthMask;\n ans.h2 ^= (uint32_t)(hash)&filter->SegmentLengthMask;\n return ans;\n}\nstatic inline uint32_t binary_fuse16_hash(int index, uint64_t hash,\n const binary_fuse16_t *filter) {\n uint64_t h = binary_fuse_mulhi(hash, filter->SegmentCountLength);\n h += index * filter->SegmentLength;\n // keep the lower 36 bits\n uint64_t hh = hash & ((1UL << 36) - 1);\n // index 0: right shift by 36; index 1: right shift by 18; index 2: no shift\n h ^= (size_t)((hh >> (36 - 18 * index)) & filter->SegmentLengthMask);\n return h;\n}\n\n// Report if the key is in the set, with false positive rate.\nstatic inline bool binary_fuse16_contain(uint64_t key,\n const binary_fuse16_t *filter) {\n uint64_t hash = binary_fuse_mix_split(key, filter->Seed);\n uint16_t f = binary_fuse16_fingerprint(hash);\n binary_hashes_t hashes = binary_fuse16_hash_batch(hash, filter);\n f ^= filter->Fingerprints[hashes.h0] ^ filter->Fingerprints[hashes.h1] ^\n filter->Fingerprints[hashes.h2];\n return f == 0;\n}\n\n// allocate enough capacity for a set containing up to 'size' elements\n// caller is responsible to call binary_fuse16_free(filter)\n// size should be at least 2.\nstatic inline bool binary_fuse16_allocate(uint32_t size,\n binary_fuse16_t *filter) {\n uint32_t arity = 3;\n filter->SegmentLength = size == 0 ? 4 : binary_fuse_calculate_segment_length(arity, size);\n if (filter->SegmentLength > 262144) {\n filter->SegmentLength = 262144;\n }\n filter->SegmentLengthMask = filter->SegmentLength - 1;\n double sizeFactor = size <= 1 ? 0 : binary_fuse_calculate_size_factor(arity, size);\n uint32_t capacity = (uint32_t)(round((double)size * sizeFactor));\n uint32_t initSegmentCount =\n (capacity + filter->SegmentLength - 1) / filter->SegmentLength -\n (arity - 1);\n filter->ArrayLength = (initSegmentCount + arity - 1) * filter->SegmentLength;\n filter->SegmentCount =\n (filter->ArrayLength + filter->SegmentLength - 1) / filter->SegmentLength;\n if (filter->SegmentCount <= arity - 1) {\n filter->SegmentCount = 1;\n } else {\n filter->SegmentCount = filter->SegmentCount - (arity - 1);\n }\n filter->ArrayLength =\n (filter->SegmentCount + arity - 1) * filter->SegmentLength;\n filter->SegmentCountLength = filter->SegmentCount * filter->SegmentLength;\n filter->Fingerprints = (uint16_t*)malloc(filter->ArrayLength * sizeof(uint16_t));\n return filter->Fingerprints != NULL;\n}\n\n// report memory usage\nstatic inline size_t\nbinary_fuse16_size_in_bytes(const binary_fuse16_t *filter) {\n return filter->ArrayLength * sizeof(uint16_t) + sizeof(binary_fuse16_t);\n}\n\n// release memory\nstatic inline void binary_fuse16_free(binary_fuse16_t *filter) {\n free(filter->Fingerprints);\n filter->Fingerprints = NULL;\n filter->Seed = 0;\n filter->SegmentLength = 0;\n filter->SegmentLengthMask = 0;\n filter->SegmentCount = 0;\n filter->SegmentCountLength = 0;\n filter->ArrayLength = 0;\n}\n\n// construct the filter, returns true on success, false on failure.\n// most likely, a failure is due to too high a memory usage\n// size is the number of keys\n// The caller is responsable for calling binary_fuse8_allocate(size,filter)\n// before. The caller is responsible to ensure that there are not too many duplicated\n// keys. The inner loop will run up to XOR_MAX_ITERATIONS times (default on\n// 100), it should never fail, except if there are many duplicated keys. If it fails,\n// a return value of false is provided.\n//\n// If there are many duplicated keys and you do not want to remove them, you can first\n// sort your input, the algorithm will then work adequately.\ninline bool binary_fuse16_populate(const uint64_t *keys, uint32_t size,\n binary_fuse16_t *filter) {\n uint64_t rng_counter = 0x726b2b9d438b9d4d;\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n uint64_t *reverseOrder = (uint64_t *)calloc((size + 1), sizeof(uint64_t));\n uint32_t capacity = filter->ArrayLength;\n uint32_t *alone = (uint32_t *)malloc(capacity * sizeof(uint32_t));\n uint8_t *t2count = (uint8_t *)calloc(capacity, sizeof(uint8_t));\n uint8_t *reverseH = (uint8_t *)malloc(size * sizeof(uint8_t));\n uint64_t *t2hash = (uint64_t *)calloc(capacity, sizeof(uint64_t));\n\n uint32_t blockBits = 1;\n while (((uint32_t)1 << blockBits) < filter->SegmentCount) {\n blockBits += 1;\n }\n uint32_t block = ((uint32_t)1 << blockBits);\n uint32_t *startPos = (uint32_t *)malloc((1 << blockBits) * sizeof(uint32_t));\n uint32_t h012[5];\n\n if ((alone == NULL) || (t2count == NULL) || (reverseH == NULL) ||\n (t2hash == NULL) || (reverseOrder == NULL) || (startPos == NULL)) {\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return false;\n }\n reverseOrder[size] = 1;\n for (int loop = 0; true; ++loop) {\n if (loop + 1 > XOR_MAX_ITERATIONS) {\n fprintf(stderr, \"Too many iterations. Are all your keys unique?\");\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return false;\n }\n\n for (uint32_t i = 0; i < block; i++) {\n // important : i * size would overflow as a 32-bit number in some\n // cases.\n startPos[i] = ((uint64_t)i * size) >> blockBits;\n }\n\n uint64_t maskblock = block - 1;\n for (uint32_t i = 0; i < size; i++) {\n uint64_t hash = binary_fuse_murmur64(keys[i] + filter->Seed);\n uint64_t segment_index = hash >> (64 - blockBits);\n while (reverseOrder[startPos[segment_index]] != 0) {\n segment_index++;\n segment_index &= maskblock;\n }\n reverseOrder[startPos[segment_index]] = hash;\n startPos[segment_index]++;\n }\n int error = 0;\n uint32_t duplicates = 0;\n for (uint32_t i = 0; i < size; i++) {\n uint64_t hash = reverseOrder[i];\n uint32_t h0 = binary_fuse16_hash(0, hash, filter);\n t2count[h0] += 4;\n t2hash[h0] ^= hash;\n uint32_t h1 = binary_fuse16_hash(1, hash, filter);\n t2count[h1] += 4;\n t2count[h1] ^= 1;\n t2hash[h1] ^= hash;\n uint32_t h2 = binary_fuse16_hash(2, hash, filter);\n t2count[h2] += 4;\n t2hash[h2] ^= hash;\n t2count[h2] ^= 2;\n if ((t2hash[h0] & t2hash[h1] & t2hash[h2]) == 0) {\n if (((t2hash[h0] == 0) && (t2count[h0] == 8)) ||\n ((t2hash[h1] == 0) && (t2count[h1] == 8)) ||\n ((t2hash[h2] == 0) && (t2count[h2] == 8))) {\n duplicates += 1;\n t2count[h0] -= 4;\n t2hash[h0] ^= hash;\n t2count[h1] -= 4;\n t2count[h1] ^= 1;\n t2hash[h1] ^= hash;\n t2count[h2] -= 4;\n t2count[h2] ^= 2;\n t2hash[h2] ^= hash;\n }\n }\n error = (t2count[h0] < 4) ? 1 : error;\n error = (t2count[h1] < 4) ? 1 : error;\n error = (t2count[h2] < 4) ? 1 : error;\n }\n if (error) {\n memset(reverseOrder, 0, sizeof(uint64_t)*size);\n memset(t2count, 0, sizeof(uint8_t)*capacity);\n memset(t2hash, 0, sizeof(uint64_t)*capacity);\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n continue;\n }\n\n // End of key addition\n uint32_t Qsize = 0;\n // Add sets with one key to the queue.\n for (uint32_t i = 0; i < capacity; i++) {\n alone[Qsize] = i;\n Qsize += ((t2count[i] >> 2) == 1) ? 1 : 0;\n }\n uint32_t stacksize = 0;\n while (Qsize > 0) {\n Qsize--;\n uint32_t index = alone[Qsize];\n if ((t2count[index] >> 2) == 1) {\n uint64_t hash = t2hash[index];\n\n // h012[0] = binary_fuse16_hash(0, hash, filter);\n h012[1] = binary_fuse16_hash(1, hash, filter);\n h012[2] = binary_fuse16_hash(2, hash, filter);\n h012[3] = binary_fuse16_hash(0, hash, filter); // == h012[0];\n h012[4] = h012[1];\n uint8_t found = t2count[index] & 3;\n reverseH[stacksize] = found;\n reverseOrder[stacksize] = hash;\n stacksize++;\n uint32_t other_index1 = h012[found + 1];\n alone[Qsize] = other_index1;\n Qsize += ((t2count[other_index1] >> 2) == 2 ? 1 : 0);\n\n t2count[other_index1] -= 4;\n t2count[other_index1] ^= binary_fuse_mod3(found + 1);\n t2hash[other_index1] ^= hash;\n\n uint32_t other_index2 = h012[found + 2];\n alone[Qsize] = other_index2;\n Qsize += ((t2count[other_index2] >> 2) == 2 ? 1 : 0);\n t2count[other_index2] -= 4;\n t2count[other_index2] ^= binary_fuse_mod3(found + 2);\n t2hash[other_index2] ^= hash;\n }\n }\n if (stacksize + duplicates == size) {\n // success\n size = stacksize;\n break;\n }\n memset(reverseOrder, 0, sizeof(uint64_t)*size);\n memset(t2count, 0, sizeof(uint8_t)*capacity);\n memset(t2hash, 0, sizeof(uint64_t)*capacity);\n filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);\n }\n\n for (uint32_t i = size - 1; i < size; i--) {\n // the hash of the key we insert next\n uint64_t hash = reverseOrder[i];\n uint16_t xor2 = binary_fuse16_fingerprint(hash);\n uint8_t found = reverseH[i];\n h012[0] = binary_fuse16_hash(0, hash, filter);\n h012[1] = binary_fuse16_hash(1, hash, filter);\n h012[2] = binary_fuse16_hash(2, hash, filter);\n h012[3] = h012[0];\n h012[4] = h012[1];\n filter->Fingerprints[h012[found]] = xor2 ^\n filter->Fingerprints[h012[found + 1]] ^\n filter->Fingerprints[h012[found + 2]];\n }\n free(alone);\n free(t2count);\n free(reverseH);\n free(t2hash);\n free(reverseOrder);\n free(startPos);\n return true;\n}\n\n#endif"
},
{
"path": "include/hnswlib/hnswalg.h",
"content": "#pragma once\n\n#include \"hnswlib.h\"\n#include \"visited_list_pool.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace hnswlib {\n typedef unsigned int tableint;\n typedef unsigned int linklistsizeint;\n\n template\n class HierarchicalNSW : public AlgorithmInterface {\n public:\n static const tableint max_update_element_locks = 65536;\n HierarchicalNSW(SpaceInterface *s) {\n }\n\n HierarchicalNSW(SpaceInterface *s, const std::string &location, bool nmslib = false, size_t max_elements=0) {\n loadIndex(location, s, max_elements);\n }\n\n HierarchicalNSW(SpaceInterface *s, size_t max_elements, size_t M = 16, size_t ef_construction = 200, size_t random_seed = 100) :\n link_list_locks_(max_elements), link_list_update_locks_(max_update_element_locks), element_levels_(max_elements) {\n max_elements_ = max_elements;\n\n has_deletions_ = false;\n num_deleted_ = 0;\n data_size_ = s->get_data_size();\n fstdistfunc_ = s->get_dist_func();\n dist_func_param_ = s->get_dist_func_param();\n M_ = M;\n maxM_ = M_;\n maxM0_ = M_ * 2;\n ef_construction_ = std::max(ef_construction,M_);\n ef_ = 10;\n\n level_generator_.seed(random_seed);\n update_probability_generator_.seed(random_seed + 1);\n\n size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint);\n size_data_per_element_ = size_links_level0_ + data_size_ + sizeof(labeltype);\n offsetData_ = size_links_level0_;\n label_offset_ = size_links_level0_ + data_size_;\n offsetLevel0_ = 0;\n\n data_level0_memory_ = (char *) malloc(max_elements_ * size_data_per_element_);\n if (data_level0_memory_ == nullptr)\n throw std::runtime_error(\"Not enough memory\");\n\n cur_element_count = 0;\n\n visited_list_pool_ = new VisitedListPool(1, max_elements);\n\n //initializations for special treatment of the first node\n enterpoint_node_ = -1;\n maxlevel_ = -1;\n\n linkLists_ = (char **) malloc(sizeof(void *) * max_elements_);\n if (linkLists_ == nullptr)\n throw std::runtime_error(\"Not enough memory: HierarchicalNSW failed to allocate linklists\");\n size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint);\n mult_ = 1 / log(1.0 * M_);\n revSize_ = 1.0 / mult_;\n }\n\n struct CompareByFirst {\n constexpr bool operator()(std::pair const &a,\n std::pair const &b) const noexcept {\n return a.first < b.first;\n }\n };\n\n ~HierarchicalNSW() {\n\n free(data_level0_memory_);\n for (tableint i = 0; i < cur_element_count; i++) {\n if (element_levels_[i] > 0)\n free(linkLists_[i]);\n }\n free(linkLists_);\n delete visited_list_pool_;\n }\n\n size_t max_elements_;\n size_t cur_element_count;\n size_t size_data_per_element_;\n size_t size_links_per_element_;\n size_t num_deleted_;\n\n size_t M_;\n size_t maxM_;\n size_t maxM0_;\n size_t ef_construction_;\n\n double mult_, revSize_;\n int maxlevel_;\n\n\n VisitedListPool *visited_list_pool_;\n std::mutex cur_element_count_guard_;\n\n std::vector link_list_locks_;\n\n // Locks to prevent race condition during update/insert of an element at same time.\n // Note: Locks for additions can also be used to prevent this race condition if the querying of KNN is not exposed along with update/inserts i.e multithread insert/update/query in parallel.\n std::vector link_list_update_locks_;\n tableint enterpoint_node_;\n\n size_t size_links_level0_;\n size_t offsetData_, offsetLevel0_;\n\n char *data_level0_memory_;\n char **linkLists_;\n std::vector element_levels_;\n\n size_t data_size_;\n\n bool has_deletions_;\n\n size_t label_offset_;\n DISTFUNC fstdistfunc_;\n void *dist_func_param_;\n std::unordered_map label_lookup_;\n\n std::default_random_engine level_generator_;\n std::default_random_engine update_probability_generator_;\n\n inline labeltype getExternalLabel(tableint internal_id) const {\n labeltype return_label;\n memcpy(&return_label,(data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), sizeof(labeltype));\n return return_label;\n }\n\n inline void setExternalLabel(tableint internal_id, labeltype label) const {\n memcpy((data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), &label, sizeof(labeltype));\n }\n\n inline labeltype *getExternalLabeLp(tableint internal_id) const {\n return (labeltype *) (data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_);\n }\n\n inline char *getDataByInternalId(tableint internal_id) const {\n return (data_level0_memory_ + internal_id * size_data_per_element_ + offsetData_);\n }\n\n int getRandomLevel(double reverse_size) {\n std::uniform_real_distribution distribution(0.0, 1.0);\n double r = -log(distribution(level_generator_)) * reverse_size;\n return (int) r;\n }\n\n // TODO:\n std::priority_queue, std::vector>, CompareByFirst>\n searchBaseLayer(tableint ep_id, const void *data_point, int layer, const local_state_t *local_state_ptr) {\n VisitedList *vl = visited_list_pool_->getFreeVisitedList();\n vl_type *visited_array = vl->mass;\n vl_type visited_array_tag = vl->curV;\n\n std::priority_queue, std::vector>, CompareByFirst> top_candidates;\n std::priority_queue, std::vector>, CompareByFirst> candidateSet;\n\n dist_t lowerBound;\n if (!isMarkedDeleted(ep_id)) {\n dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_, local_state_ptr);\n top_candidates.emplace(dist, ep_id);\n lowerBound = dist;\n candidateSet.emplace(-dist, ep_id);\n } else {\n lowerBound = std::numeric_limits::max();\n candidateSet.emplace(-lowerBound, ep_id);\n }\n visited_array[ep_id] = visited_array_tag;\n\n while (!candidateSet.empty()) {\n std::pair curr_el_pair = candidateSet.top();\n if ((-curr_el_pair.first) > lowerBound) {\n break;\n }\n candidateSet.pop();\n\n tableint curNodeNum = curr_el_pair.second;\n\n std::unique_lock lock(link_list_locks_[curNodeNum]);\n\n int *data;// = (int *)(linkList0_ + curNodeNum * size_links_per_element0_);\n if (layer == 0) {\n data = (int*)get_linklist0(curNodeNum);\n } else {\n data = (int*)get_linklist(curNodeNum, layer);\n// data = (int *) (linkLists_[curNodeNum] + (layer - 1) * size_links_per_element_);\n }\n size_t size = getListCount((linklistsizeint*)data);\n tableint *datal = (tableint *) (data + 1);\n#ifdef USE_SSE\n _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0);\n _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0);\n _mm_prefetch(getDataByInternalId(*datal), _MM_HINT_T0);\n _mm_prefetch(getDataByInternalId(*(datal + 1)), _MM_HINT_T0);\n#endif\n\n for (size_t j = 0; j < size; j++) {\n tableint candidate_id = *(datal + j);\n// if (candidate_id == 0) continue;\n#ifdef USE_SSE\n _mm_prefetch((char *) (visited_array + *(datal + j + 1)), _MM_HINT_T0);\n _mm_prefetch(getDataByInternalId(*(datal + j + 1)), _MM_HINT_T0);\n#endif\n if (visited_array[candidate_id] == visited_array_tag) continue;\n visited_array[candidate_id] = visited_array_tag;\n char *currObj1 = (getDataByInternalId(candidate_id));\n\n dist_t dist1 = fstdistfunc_(data_point, currObj1, dist_func_param_, local_state_ptr);\n if (top_candidates.size() < ef_construction_ || lowerBound > dist1) {\n candidateSet.emplace(-dist1, candidate_id);\n#ifdef USE_SSE\n _mm_prefetch(getDataByInternalId(candidateSet.top().second), _MM_HINT_T0);\n#endif\n\n if (!isMarkedDeleted(candidate_id))\n top_candidates.emplace(dist1, candidate_id);\n\n if (top_candidates.size() > ef_construction_)\n top_candidates.pop();\n\n if (!top_candidates.empty())\n lowerBound = top_candidates.top().first;\n }\n }\n }\n visited_list_pool_->releaseVisitedList(vl);\n\n return top_candidates;\n }\n\n mutable std::atomic metric_distance_computations;\n mutable std::atomic metric_hops;\n\n template \n std::priority_queue, std::vector>, CompareByFirst>\n searchBaseLayerST(tableint ep_id, const void *data_point, size_t ef, const local_state_t *local_state_ptr) const {\n VisitedList *vl = visited_list_pool_->getFreeVisitedList();\n vl_type *visited_array = vl->mass;\n vl_type visited_array_tag = vl->curV;\n\n std::priority_queue, std::vector>, CompareByFirst> top_candidates;\n std::priority_queue, std::vector>, CompareByFirst> candidate_set;\n\n dist_t lowerBound;\n if (!has_deletions || !isMarkedDeleted(ep_id)) {\n dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_, local_state_ptr);\n lowerBound = dist;\n top_candidates.emplace(dist, ep_id);\n candidate_set.emplace(-dist, ep_id);\n } else {\n lowerBound = std::numeric_limits::max();\n candidate_set.emplace(-lowerBound, ep_id);\n }\n\n visited_array[ep_id] = visited_array_tag;\n\n while (!candidate_set.empty()) {\n\n std::pair current_node_pair = candidate_set.top();\n\n if ((-current_node_pair.first) > lowerBound && (top_candidates.size() == ef || has_deletions == false)) {\n break;\n }\n candidate_set.pop();\n\n tableint current_node_id = current_node_pair.second;\n int *data = (int *) get_linklist0(current_node_id);\n size_t size = getListCount((linklistsizeint*)data);\n// bool cur_node_deleted = isMarkedDeleted(current_node_id);\n if(collect_metrics){\n metric_hops++;\n metric_distance_computations+=size;\n }\n\n#ifdef USE_SSE\n _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0);\n _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0);\n _mm_prefetch(data_level0_memory_ + (*(data + 1)) * size_data_per_element_ + offsetData_, _MM_HINT_T0);\n _mm_prefetch((char *) (data + 2), _MM_HINT_T0);\n#endif\n\n for (size_t j = 1; j <= size; j++) {\n int candidate_id = *(data + j);\n// if (candidate_id == 0) continue;\n#ifdef USE_SSE\n _mm_prefetch((char *) (visited_array + *(data + j + 1)), _MM_HINT_T0);\n _mm_prefetch(data_level0_memory_ + (*(data + j + 1)) * size_data_per_element_ + offsetData_,\n _MM_HINT_T0);////////////\n#endif\n if (!(visited_array[candidate_id] == visited_array_tag)) {\n\n visited_array[candidate_id] = visited_array_tag;\n\n char *currObj1 = (getDataByInternalId(candidate_id));\n dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_, local_state_ptr);\n\n if (top_candidates.size() < ef || lowerBound > dist) {\n candidate_set.emplace(-dist, candidate_id);\n#ifdef USE_SSE\n _mm_prefetch(data_level0_memory_ + candidate_set.top().second * size_data_per_element_ +\n offsetLevel0_,///////////\n _MM_HINT_T0);////////////////////////\n#endif\n\n if (!has_deletions || !isMarkedDeleted(candidate_id))\n top_candidates.emplace(dist, candidate_id);\n\n if (top_candidates.size() > ef)\n top_candidates.pop();\n\n if (!top_candidates.empty())\n lowerBound = top_candidates.top().first;\n }\n }\n }\n }\n\n visited_list_pool_->releaseVisitedList(vl);\n return top_candidates;\n }\n\n // TODO\n template \n std::priority_queue,\n std::vector>, CompareByFirst>\n searchBaseLayerSTWithFilter(tableint ep_id, const void *data_point,\n const binary_fuse16_t *filter, size_t ef,\n const local_state_t *local_state_ptr) const {\n VisitedList *vl = visited_list_pool_->getFreeVisitedList();\n vl_type *visited_array = vl->mass;\n vl_type visited_array_tag = vl->curV;\n\n std::priority_queue,\n std::vector>,\n CompareByFirst>\n top_candidates;\n std::priority_queue,\n std::vector>,\n CompareByFirst>\n candidate_set;\n\n dist_t lowerBound;\n\n uint64_t label = getExternalLabel(ep_id);\n if (binary_fuse16_contain(label, filter)) {\n dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id),\n dist_func_param_, local_state_ptr);\n lowerBound = dist;\n top_candidates.emplace(dist, ep_id);\n candidate_set.emplace(-dist, ep_id);\n } else {\n lowerBound = std::numeric_limits::max();\n candidate_set.emplace(-lowerBound, ep_id);\n }\n\n visited_array[ep_id] = visited_array_tag;\n\n while (!candidate_set.empty()) {\n\n std::pair current_node_pair = candidate_set.top();\n\n if ((-current_node_pair.first) > lowerBound) {\n break;\n }\n candidate_set.pop();\n\n tableint current_node_id = current_node_pair.second;\n int *data = (int *)get_linklist0(current_node_id);\n size_t size = getListCount((linklistsizeint *)data);\n // bool cur_node_deleted =\n // isMarkedDeleted(current_node_id);\n if (collect_metrics) {\n metric_hops++;\n metric_distance_computations += size;\n }\n\n#ifdef USE_SSE\n _mm_prefetch((char *)(visited_array + *(data + 1)), _MM_HINT_T0);\n _mm_prefetch((char *)(visited_array + *(data + 1) + 64), _MM_HINT_T0);\n _mm_prefetch(data_level0_memory_ +\n (*(data + 1)) * size_data_per_element_ + offsetData_,\n _MM_HINT_T0);\n _mm_prefetch((char *)(data + 2), _MM_HINT_T0);\n#endif\n\n for (size_t j = 1; j <= size; j++) {\n int candidate_id = *(data + j);\n// if (candidate_id == 0) continue;\n#ifdef USE_SSE\n _mm_prefetch((char *)(visited_array + *(data + j + 1)), _MM_HINT_T0);\n _mm_prefetch(data_level0_memory_ +\n (*(data + j + 1)) * size_data_per_element_ +\n offsetData_,\n _MM_HINT_T0); ////////////\n#endif\n if (!(visited_array[candidate_id] == visited_array_tag)) {\n\n visited_array[candidate_id] = visited_array_tag;\n\n char *currObj1 = (getDataByInternalId(candidate_id));\n dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_,\n local_state_ptr);\n\n if (top_candidates.size() < ef || lowerBound > dist) {\n candidate_set.emplace(-dist, candidate_id);\n#ifdef USE_SSE\n _mm_prefetch(data_level0_memory_ +\n candidate_set.top().second *\n size_data_per_element_ +\n offsetLevel0_, ///////////\n _MM_HINT_T0); ////////////////////////\n#endif\n\n uint64_t et_label = getExternalLabel(candidate_id);\n\n if (binary_fuse16_contain(et_label, filter))\n top_candidates.emplace(dist, candidate_id);\n\n if (top_candidates.size() > ef)\n top_candidates.pop();\n\n if (!top_candidates.empty())\n lowerBound = top_candidates.top().first;\n }\n }\n }\n }\n\n visited_list_pool_->releaseVisitedList(vl);\n return top_candidates;\n }\n\n // TODO\n void getNeighborsByHeuristic2(\n std::priority_queue, std::vector>, CompareByFirst> &top_candidates,\n const size_t M, const local_state_t *local_state_ptr) {\n if (top_candidates.size() < M) {\n return;\n }\n\n std::priority_queue> queue_closest;\n std::vector> return_list;\n while (top_candidates.size() > 0) {\n queue_closest.emplace(-top_candidates.top().first, top_candidates.top().second);\n top_candidates.pop();\n }\n\n while (queue_closest.size()) {\n if (return_list.size() >= M)\n break;\n std::pair curent_pair = queue_closest.top();\n dist_t dist_to_query = -curent_pair.first;\n queue_closest.pop();\n bool good = true;\n\n for (std::pair second_pair : return_list) {\n dist_t curdist =\n fstdistfunc_(getDataByInternalId(second_pair.second),\n getDataByInternalId(curent_pair.second),\n dist_func_param_, local_state_ptr);;\n if (curdist < dist_to_query) {\n good = false;\n break;\n }\n }\n if (good) {\n return_list.push_back(curent_pair);\n }\n }\n\n for (std::pair curent_pair : return_list) {\n top_candidates.emplace(-curent_pair.first, curent_pair.second);\n }\n }\n\n\n linklistsizeint *get_linklist0(tableint internal_id) const {\n return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_);\n };\n\n linklistsizeint *get_linklist0(tableint internal_id, char *data_level0_memory_) const {\n return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_);\n };\n\n linklistsizeint *get_linklist(tableint internal_id, int level) const {\n return (linklistsizeint *) (linkLists_[internal_id] + (level - 1) * size_links_per_element_);\n };\n\n linklistsizeint *get_linklist_at_level(tableint internal_id, int level) const {\n return level == 0 ? get_linklist0(internal_id) : get_linklist(internal_id, level);\n };\n\n tableint mutuallyConnectNewElement(const void *data_point, tableint cur_c,\n std::priority_queue, std::vector>, CompareByFirst> &top_candidates,\n int level, bool isUpdate, const local_state_t *local_state_ptr) {\n size_t Mcurmax = level ? maxM_ : maxM0_;\n getNeighborsByHeuristic2(top_candidates, M_, local_state_ptr);\n if (top_candidates.size() > M_)\n throw std::runtime_error(\"Should be not be more than M_ candidates returned by the heuristic\");\n\n std::vector selectedNeighbors;\n selectedNeighbors.reserve(M_);\n while (top_candidates.size() > 0) {\n selectedNeighbors.push_back(top_candidates.top().second);\n top_candidates.pop();\n }\n\n tableint next_closest_entry_point = selectedNeighbors.back();\n\n {\n linklistsizeint *ll_cur;\n if (level == 0)\n ll_cur = get_linklist0(cur_c);\n else\n ll_cur = get_linklist(cur_c, level);\n\n if (*ll_cur && !isUpdate) {\n throw std::runtime_error(\"The newly inserted element should have blank link list\");\n }\n setListCount(ll_cur,selectedNeighbors.size());\n tableint *data = (tableint *) (ll_cur + 1);\n for (size_t idx = 0; idx < selectedNeighbors.size(); idx++) {\n if (data[idx] && !isUpdate)\n throw std::runtime_error(\"Possible memory corruption\");\n if (level > element_levels_[selectedNeighbors[idx]])\n throw std::runtime_error(\"Trying to make a link on a non-existent level\");\n\n data[idx] = selectedNeighbors[idx];\n\n }\n }\n\n for (size_t idx = 0; idx < selectedNeighbors.size(); idx++) {\n\n std::unique_lock lock(link_list_locks_[selectedNeighbors[idx]]);\n\n linklistsizeint *ll_other;\n if (level == 0)\n ll_other = get_linklist0(selectedNeighbors[idx]);\n else\n ll_other = get_linklist(selectedNeighbors[idx], level);\n\n size_t sz_link_list_other = getListCount(ll_other);\n\n if (sz_link_list_other > Mcurmax)\n throw std::runtime_error(\"Bad value of sz_link_list_other\");\n if (selectedNeighbors[idx] == cur_c)\n throw std::runtime_error(\"Trying to connect an element to itself\");\n if (level > element_levels_[selectedNeighbors[idx]])\n throw std::runtime_error(\"Trying to make a link on a non-existent level\");\n\n tableint *data = (tableint *) (ll_other + 1);\n\n bool is_cur_c_present = false;\n if (isUpdate) {\n for (size_t j = 0; j < sz_link_list_other; j++) {\n if (data[j] == cur_c) {\n is_cur_c_present = true;\n break;\n }\n }\n }\n\n // If cur_c is already present in the neighboring connections of `selectedNeighbors[idx]` then no need to modify any connections or run the heuristics.\n if (!is_cur_c_present) {\n if (sz_link_list_other < Mcurmax) {\n data[sz_link_list_other] = cur_c;\n setListCount(ll_other, sz_link_list_other + 1);\n } else {\n // finding the \"weakest\" element to replace it with the new one\n dist_t d_max = fstdistfunc_(getDataByInternalId(cur_c), getDataByInternalId(selectedNeighbors[idx]),\n dist_func_param_, local_state_ptr);\n // Heuristic:\n std::priority_queue, std::vector>, CompareByFirst> candidates;\n candidates.emplace(d_max, cur_c);\n\n for (size_t j = 0; j < sz_link_list_other; j++) {\n candidates.emplace(\n fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(selectedNeighbors[idx]),\n dist_func_param_, local_state_ptr), data[j]);\n }\n\n getNeighborsByHeuristic2(candidates, Mcurmax, local_state_ptr);\n\n int indx = 0;\n while (candidates.size() > 0) {\n data[indx] = candidates.top().second;\n candidates.pop();\n indx++;\n }\n\n setListCount(ll_other, indx);\n // Nearest K:\n /*int indx = -1;\n for (int j = 0; j < sz_link_list_other; j++) {\n dist_t d = fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(rez[idx]), dist_func_param_, nullptr);\n if (d > d_max) {\n indx = j;\n d_max = d;\n }\n }\n if (indx >= 0) {\n data[indx] = cur_c;\n } */\n }\n }\n }\n\n return next_closest_entry_point;\n }\n\n std::mutex global;\n size_t ef_;\n\n void setEf(size_t ef) {\n ef_ = ef;\n }\n\n\n std::priority_queue> searchKnnInternal(void *query_data, int k) {\n std::priority_queue> top_candidates;\n if (cur_element_count == 0) return top_candidates;\n\n // FIXME: Never used, so fix index to 0.\n local_state_t local_state;\n local_state.batch_index = 0;\n\n tableint currObj = enterpoint_node_;\n dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_, &local_state);\n\n for (size_t level = maxlevel_; level > 0; level--) {\n bool changed = true;\n while (changed) {\n changed = false;\n int *data;\n data = (int *) get_linklist(currObj,level);\n int size = getListCount(data);\n tableint *datal = (tableint *) (data + 1);\n for (int i = 0; i < size; i++) {\n tableint cand = datal[i];\n if (cand < 0 || cand > max_elements_)\n throw std::runtime_error(\"cand error\");\n dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_, nullptr);\n\n if (d < curdist) {\n curdist = d;\n currObj = cand;\n changed = true;\n }\n }\n }\n }\n\n if (num_deleted_) {\n std::priority_queue> top_candidates1=searchBaseLayerST(currObj, query_data,\n ef_, &local_state);\n top_candidates.swap(top_candidates1);\n }\n else{\n std::priority_queue