diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 2c0e88993..b72c1da95 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,2 +1,2 @@
# This should match the owning team set up in https://github.com/orgs/opensearch-project/teams
-* @opensearch-project/k-nn
\ No newline at end of file
+* @heemin32 @navneet1v @VijayanB @vamshin @jmazanec15 @naveentatikonda @junqiu-lei @martin-gaievski @ryanbogan @luyuncheng
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index f9e914099..c37596a1f 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -6,79 +6,190 @@ on:
branches:
- "*"
- "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - 'micro-benchmarks/**'
+ - '.github/workflows/CI.yml'
pull_request:
branches:
- "*"
- "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - 'micro-benchmarks/**'
+ - '.github/workflows/CI.yml'
+env:
+ ACTIONS_ALLOW_USE_UNSECURE_NODE_VERSION: true
jobs:
- Build-k-NN:
+ Get-CI-Image-Tag:
+ uses: opensearch-project/opensearch-build/.github/workflows/get-ci-image-tag.yml@main
+ with:
+ product: opensearch
+
+ Build-k-NN-Linux:
strategy:
matrix:
- java: [11, 17]
+ java: [11, 17, 21]
- name: Build and Test k-NN Plugin
+ name: Build and Test k-NN Plugin on Linux
runs-on: ubuntu-latest
+ needs: Get-CI-Image-Tag
+ container:
+ # using the same image which is used by opensearch-build team to build the OpenSearch Distribution
+ # this image tag is subject to change as more dependencies and updates will arrive over time
+ image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }}
+ # need to switch to root so that github actions can install runner binary on container without permission issues.
+ options: --user root
steps:
- name: Checkout k-NN
uses: actions/checkout@v1
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ su `id -un 1000` -c 'git config --global user.name "github-actions[bot]"'
+ su `id -un 1000` -c 'git config --global user.email "github-actions[bot]@users.noreply.github.com"'
+
- name: Setup Java ${{ matrix.java }}
uses: actions/setup-java@v1
with:
java-version: ${{ matrix.java }}
- - name: Install dependencies
- run: |
- sudo apt-get install libopenblas-dev gfortran -y
-
- name: Run build
+ # switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
run: |
- ./gradlew build
+ chown -R 1000:1000 `pwd`
+ if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
+ then
+ echo "avx512 available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
+ elif lscpu | grep -i avx2
+ then
+ echo "avx2 available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false"
+ else
+ echo "avx512 and avx2 not available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`"
+ fi
+
- name: Upload Coverage Report
uses: codecov/codecov-action@v1
with:
token: ${{ secrets.CODECOV_TOKEN }}
-# - name: Pull and Run Docker for security tests
-# run: |
-# plugin=`ls build/distributions/*.zip`
-# version=`echo $plugin|awk -F- '{print $4}'| cut -d. -f 1-3`
-# plugin_version=`echo $plugin|awk -F- '{print $4}'| cut -d. -f 1-4`
-# echo $version
-# cd ..
-# if docker pull opendistroforelasticsearch/opendistroforelasticsearch:$version
-# then
-# echo "FROM opendistroforelasticsearch/opendistroforelasticsearch:$version" >> Dockerfile
-# echo "RUN if [ -d /usr/share/elasticsearch/plugins/opendistro-knn ]; then /usr/share/elasticsearch/bin/elasticsearch-plugin remove opendistro-knn; fi" >> Dockerfile
-# echo "RUN yum -y update \ && yum -y groupinstall "Development Tools" \ && yum install -y unzip glibc.x86_64 cmake \ && yum clean all" >> Dockerfile
-# echo "RUN git clone --recursive --branch ${GITHUB_REF##*/} https://github.com/opendistro-for-elasticsearch/k-NN.git /usr/share/elasticsearch/k-NN \ " >> Dockerfile
-# echo "&& cd /usr/share/elasticsearch/k-NN/jni \ && sed -i 's/-march=native/-march=x86-64/g' external/nmslib/similarity_search/CMakeLists.txt \ && cmake . \ && make \ " >> Dockerfile
-# echo "&& mkdir /tmp/jni/ && cp release/*.so /tmp/jni/ && ls -ltr /tmp/jni/ \ && cp /tmp/jni/libKNNIndex*.so /usr/lib \ && rm -rf /usr/share/elasticsearch/k-NN" >> Dockerfile
-# echo "RUN cd /usr/share/elasticsearch/" >> Dockerfile
-# echo "ADD k-NN/build/distributions/opendistro-knn-$plugin_version.zip /tmp/" >> Dockerfile
-# echo "RUN /usr/share/elasticsearch/bin/elasticsearch-plugin install --batch file:/tmp/opendistro-knn-$plugin_version.zip" >> Dockerfile
-# docker build -t odfe-knn:test .
-# echo "imagePresent=true" >> $GITHUB_ENV
-# else
-# echo "imagePresent=false" >> $GITHUB_ENV
-# fi
-# - name: Run Docker Image
-# if: env.imagePresent == 'true'
-# run: |
-# cd ..
-# docker run -p 9200:9200 -d -p 9600:9600 -e "discovery.type=single-node" odfe-knn:test
-# sleep 90
-# - name: Run k-NN Test
-# if: env.imagePresent == 'true'
-# run: |
-# security=`curl -XGET https://localhost:9200/_cat/plugins?v -u admin:admin --insecure |grep opendistro_security|wc -l`
-# if [ $security -gt 0 ]
-# then
-# echo "Security plugin is available. Running tests in security mode"
-# ./gradlew :integTest -Dtests.rest.cluster=localhost:9200 -Dtests.cluster=localhost:9200 -Dtests.clustername="docker-cluster" -Dhttps=true -Duser=admin -Dpassword=admin
-# else
-# echo "Security plugin is NOT available. Skipping tests as they are already ran part of ./gradlew build"
-# fi
+ Build-k-NN-MacOS:
+ strategy:
+ matrix:
+ java: [ 11, 17, 21 ]
+
+ name: Build and Test k-NN Plugin on MacOS
+ needs: Get-CI-Image-Tag
+ runs-on: macos-13
+
+ steps:
+ - name: Checkout k-NN
+ uses: actions/checkout@v1
+
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ git config --global user.name "github-actions[bot]"
+ git config --global user.email "github-actions[bot]@users.noreply.github.com"
+
+ - name: Setup Java ${{ matrix.java }}
+ uses: actions/setup-java@v1
+ with:
+ java-version: ${{ matrix.java }}
+
+ - name: Install dependencies on macos
+ run: |
+ brew install libomp
+ brew reinstall gcc
+ export FC=/usr/local/Cellar/gcc/12.2.0/bin/gfortran
+
+ # TODO: Detect processor count and set the value of nproc.count
+ - name: Run build
+ run: |
+ if sysctl -n machdep.cpu.features machdep.cpu.leaf7_features | grep -i AVX2
+ then
+ echo "avx2 available on system"
+ ./gradlew build -Dnproc.count=3
+ else
+ echo "avx2 not available on system"
+ ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=3
+ fi
+
+ Build-k-NN-Windows:
+ strategy:
+ matrix:
+ java: [ 11, 17, 21 ]
+
+ name: Build and Test k-NN Plugin on Windows
+ needs: Get-CI-Image-Tag
+ runs-on: windows-latest
+
+ steps:
+ - name: Checkout k-NN
+ uses: actions/checkout@v1
+
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ git config --global user.name "github-actions[bot]"
+ git config --global user.email "github-actions[bot]@users.noreply.github.com"
+
+ - name: Setup Java ${{ matrix.java }}
+ uses: actions/setup-java@v1
+ with:
+ java-version: ${{ matrix.java }}
+
+ - name: Install MinGW Using Scoop
+ run: |
+ iex "& {$(irm get.scoop.sh)} -RunAsAdmin"
+ scoop bucket add main
+ scoop install mingw
+
+ - name: Add MinGW to PATH
+ run: |
+ echo "C:/Users/runneradmin/scoop/apps/mingw/current/bin" >> $env:GITHUB_PATH
+ Import-Module "$env:ChocolateyInstall\helpers\chocolateyProfile.psm1"
+ refreshenv
+
+ - name: Install Zlib Using Scoop
+ run: |
+ echo "C:/Users/runneradmin/scoop/shims" >> $env:GITHUB_PATH
+ Import-Module "$env:ChocolateyInstall\helpers\chocolateyProfile.psm1"
+ refreshenv
+ scoop bucket add extras
+ scoop install zlib
+ regedit /s "C:\\Users\\runneradmin\\scoop\\apps\\zlib\\current\\register.reg"
+
+ - name: Download OpenBLAS
+ run: |
+ curl -L -O https://github.com/xianyi/OpenBLAS/releases/download/v0.3.21/OpenBLAS-0.3.21-x64.zip
+ mkdir OpenBLAS
+ Expand-Archive -Path .\OpenBLAS-0.3.21-x64.zip -DestinationPath .\OpenBLAS\
+ mkdir ./src/main/resources/windowsDependencies
+ cp ./OpenBLAS/bin/libopenblas.dll ./src/main/resources/windowsDependencies/
+ rm .\OpenBLAS-0.3.21-x64.zip
+ rm -r .\OpenBLAS\
+
+ # TODO: Detect processor count and set the value of nproc.count
+ - name: Run build
+ run: |
+ ./gradlew.bat build -D'avx2.enabled=false' -D'avx512.enabled=false' -D'nproc.count=4'
diff --git a/.github/workflows/auto-release.yml b/.github/workflows/auto-release.yml
new file mode 100644
index 000000000..c67e12bad
--- /dev/null
+++ b/.github/workflows/auto-release.yml
@@ -0,0 +1,28 @@
+name: Releases
+
+on:
+ push:
+ tags:
+ - '*'
+
+jobs:
+ build:
+ runs-on: ubuntu-latest
+ permissions:
+ contents: write
+ steps:
+ - name: GitHub App token
+ id: github_app_token
+ uses: tibdex/github-app-token@v1.5.0
+ with:
+ app_id: ${{ secrets.APP_ID }}
+ private_key: ${{ secrets.APP_PRIVATE_KEY }}
+ installation_id: 22958780
+ - name: Get tag
+ id: tag
+ uses: dawidd6/action-get-tag@v1
+ - uses: actions/checkout@v2
+ - uses: ncipollo/release-action@v1
+ with:
+ github_token: ${{ steps.github_app_token.outputs.token }}
+ bodyFile: release-notes/opensearch-knn.release-notes-${{steps.tag.outputs.tag}}.md
diff --git a/.github/workflows/backwards_compatibility_tests_workflow.yml b/.github/workflows/backwards_compatibility_tests_workflow.yml
index ce654eee9..3229b5d69 100644
--- a/.github/workflows/backwards_compatibility_tests_workflow.yml
+++ b/.github/workflows/backwards_compatibility_tests_workflow.yml
@@ -4,18 +4,38 @@ on:
branches:
- "*"
- "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - 'qa/**'
+ - '.github/workflows/backwards_compatibility_tests_workflow.yml'
pull_request:
branches:
- "*"
- "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - 'qa/**'
+ - '.github/workflows/backwards_compatibility_tests_workflow.yml'
jobs:
Restart-Upgrade-BWCTests-k-NN:
strategy:
matrix:
java: [ 11, 17 ]
- bwc_version : [ "1.0.0", "1.1.0", "1.2.4", "1.3.2", "2.0.0" ]
- opensearch_version : [ "2.1.0-SNAPSHOT" ]
+ bwc_version : [ "1.1.0", "1.2.4", "1.3.8", "2.0.1", "2.1.0", "2.2.1", "2.3.0", "2.4.1", "2.5.0", "2.6.0", "2.7.0", "2.8.0", "2.9.0", "2.10.0", "2.11.0", "2.12.0", "2.13.0", "2.14.0", "2.15.0", "2.16.0", "2.17.0", "2.18.0" ]
+ opensearch_version : [ "2.19.0-SNAPSHOT" ]
name: k-NN Restart-Upgrade BWC Tests
runs-on: ubuntu-latest
@@ -26,6 +46,12 @@ jobs:
- name: Checkout k-NN
uses: actions/checkout@v1
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ git config --global user.name "github-actions[bot]"
+ git config --global user.email "github-actions[bot]@users.noreply.github.com"
+
- name: Setup Java ${{ matrix.java }}
uses: actions/setup-java@v1
with:
@@ -38,26 +64,42 @@ jobs:
- name: Run k-NN Restart-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }}
run: |
- echo "Running restart-upgrade backwards compatibility tests ..."
- ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE
-
+ echo "Running restart-upgrade backwards compatibility tests ..."
+ if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
+ then
+ echo "avx512 available on system"
+ ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc`
+ elif lscpu | grep -i avx2
+ then
+ echo "avx2 available on system"
+ ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false
+ else
+ echo "avx512 and avx2 not available on system"
+ ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc`
+ fi
Rolling-Upgrade-BWCTests-k-NN:
strategy:
matrix:
java: [ 11, 17 ]
- bwc_version: [ "1.3.2", "2.0.0" ]
- opensearch_version: [ "2.1.0-SNAPSHOT" ]
+ bwc_version: [ "1.3.8", "2.0.1", "2.1.0", "2.2.1", "2.3.0", "2.4.1", "2.5.0", "2.6.0", "2.7.0", "2.8.0", "2.9.0", "2.10.0", "2.11.0", "2.12.0", "2.13.0", "2.14.0", "2.15.0", "2.16.0", "2.17.0", "2.18.0"]
+ opensearch_version: [ "2.19.0-SNAPSHOT" ]
name: k-NN Rolling-Upgrade BWC Tests
runs-on: ubuntu-latest
env:
- OPENSEARCH_VERSION_ROLLING_UPGRADE: ${{ matrix.opensearch_version }}
+ BWC_VERSION_ROLLING_UPGRADE: ${{ matrix.bwc_version }}
steps:
- name: Checkout k-NN
uses: actions/checkout@v1
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ git config --global user.name "github-actions[bot]"
+ git config --global user.email "github-actions[bot]@users.noreply.github.com"
+
- name: Setup Java ${{ matrix.java }}
uses: actions/setup-java@v1
with:
@@ -70,4 +112,15 @@ jobs:
- name: Run k-NN Rolling-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }}
run: |
echo "Running rolling-upgrade backwards compatibility tests ..."
- ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Drolling.bwctests.opensearch.version=$OPENSEARCH_VERSION_ROLLING_UPGRADE
+ if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
+ then
+ echo "avx512 available on system"
+ ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc`
+ elif lscpu | grep -i avx2
+ then
+ echo "avx2 available on system"
+ ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false
+ else
+ echo "avx512 and avx2 not available on system"
+ ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc`
+ fi
diff --git a/.github/workflows/changelog_verifier.yml b/.github/workflows/changelog_verifier.yml
new file mode 100644
index 000000000..992a38b62
--- /dev/null
+++ b/.github/workflows/changelog_verifier.yml
@@ -0,0 +1,18 @@
+name: "Changelog Verifier"
+on:
+ pull_request:
+ types: [opened, edited, review_requested, synchronize, reopened, ready_for_review, labeled, unlabeled]
+
+jobs:
+ # Enforces the update of a changelog file on every pull request
+ verify-changelog:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v3
+ with:
+ token: ${{ secrets.GITHUB_TOKEN }}
+ ref: ${{ github.event.pull_request.head.sha }}
+
+ - uses: dangoslen/changelog-enforcer@v3
+ with:
+ skipLabels: "autocut, skip-changelog"
diff --git a/.github/workflows/dco.yml b/.github/workflows/dco.yml
deleted file mode 100644
index cf30ea89d..000000000
--- a/.github/workflows/dco.yml
+++ /dev/null
@@ -1,18 +0,0 @@
-name: Developer Certificate of Origin Check
-
-on: [pull_request]
-
-jobs:
- check:
- runs-on: ubuntu-latest
-
- steps:
- - name: Get PR Commits
- id: 'get-pr-commits'
- uses: tim-actions/get-pr-commits@v1.1.0
- with:
- token: ${{ secrets.GITHUB_TOKEN }}
- - name: DCO Check
- uses: tim-actions/dco@v1.1.0
- with:
- commits: ${{ steps.get-pr-commits.outputs.commits }}
diff --git a/.github/workflows/delete_backport_branch.yml b/.github/workflows/delete_backport_branch.yml
index d654df6b4..a4d186ca7 100644
--- a/.github/workflows/delete_backport_branch.yml
+++ b/.github/workflows/delete_backport_branch.yml
@@ -7,9 +7,16 @@ on:
jobs:
delete-branch:
runs-on: ubuntu-latest
- if: startsWith(github.event.pull_request.head.ref,'backport/')
+ permissions:
+ contents: write
+ if: github.repository == 'opensearch-project/k-NN' && startsWith(github.event.pull_request.head.ref,'backport/')
steps:
- - name: Delete merged branch
- uses: SvanBoxel/delete-merged-branch@main
- env:
- GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
\ No newline at end of file
+ - name: Delete merged branch
+ uses: actions/github-script@v7
+ with:
+ script: |
+ github.rest.git.deleteRef({
+ owner: context.repo.owner,
+ repo: context.repo.repo,
+ ref: `heads/${context.payload.pull_request.head.ref}`,
+ })
diff --git a/.github/workflows/maven-publish.yml b/.github/workflows/maven-publish.yml
new file mode 100644
index 000000000..724e3a213
--- /dev/null
+++ b/.github/workflows/maven-publish.yml
@@ -0,0 +1,35 @@
+name: Publish snapshots to maven
+
+on:
+ workflow_dispatch:
+ push:
+ branches:
+ - 'main'
+ - '1.*'
+ - '2.*'
+
+jobs:
+ build-and-publish-snapshots:
+ runs-on: ubuntu-latest
+
+ permissions:
+ id-token: write
+ contents: write
+
+ steps:
+ - uses: actions/setup-java@v3
+ with:
+ distribution: temurin # Temurin is a distribution of adoptium
+ java-version: 11
+ - uses: actions/checkout@v3
+ - uses: aws-actions/configure-aws-credentials@v1
+ with:
+ role-to-assume: ${{ secrets.PUBLISH_SNAPSHOTS_ROLE }}
+ aws-region: us-east-1
+ - name: publish snapshots to maven
+ run: |
+ export SONATYPE_USERNAME=$(aws secretsmanager get-secret-value --secret-id maven-snapshots-username --query SecretString --output text)
+ export SONATYPE_PASSWORD=$(aws secretsmanager get-secret-value --secret-id maven-snapshots-password --query SecretString --output text)
+ echo "::add-mask::$SONATYPE_USERNAME"
+ echo "::add-mask::$SONATYPE_PASSWORD"
+ ./gradlew publishPluginZipPublicationToSnapshotsRepository
diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml
new file mode 100644
index 000000000..f3deb096e
--- /dev/null
+++ b/.github/workflows/test_security.yml
@@ -0,0 +1,84 @@
+name: Test k-NN on Secure Cluster
+on:
+ schedule:
+ - cron: '0 0 * * *' # every night
+ push:
+ branches:
+ - "*"
+ - "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - '.github/workflows/test_security.yml'
+ pull_request:
+ branches:
+ - "*"
+ - "feature/**"
+ paths:
+ - 'build.gradle'
+ - 'settings.gradle'
+ - 'src/**'
+ - 'build-tools/**'
+ - 'buildSrc/**'
+ - 'gradle/**'
+ - 'jni/**'
+ - '.github/workflows/test_security.yml'
+env:
+ ACTIONS_ALLOW_USE_UNSECURE_NODE_VERSION: true
+
+jobs:
+ Get-CI-Image-Tag:
+ uses: opensearch-project/opensearch-build/.github/workflows/get-ci-image-tag.yml@main
+ with:
+ product: opensearch
+
+ integ-test-with-security-linux:
+ strategy:
+ matrix:
+ java: [21]
+
+ name: Run Integration Tests on Linux
+ runs-on: ubuntu-latest
+ needs: Get-CI-Image-Tag
+ container:
+ # using the same image which is used by opensearch-build team to build the OpenSearch Distribution
+ # this image tag is subject to change as more dependencies and updates will arrive over time
+ image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }}
+ # need to switch to root so that github actions can install runner binary on container without permission issues.
+ options: --user root
+
+ steps:
+ - name: Checkout k-NN
+ uses: actions/checkout@v1
+ # Setup git user so that patches for native libraries can be applied and committed
+ - name: Setup git user
+ run: |
+ su `id -un 1000` -c 'git config --global user.name "github-actions[bot]"'
+ su `id -un 1000` -c 'git config --global user.email "github-actions[bot]@users.noreply.github.com"'
+
+ - name: Setup Java ${{ matrix.java }}
+ uses: actions/setup-java@v1
+ with:
+ java-version: ${{ matrix.java }}
+
+ - name: Run build
+ # switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
+ run: |
+ chown -R 1000:1000 `pwd`
+ if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
+ then
+ echo "avx512 available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
+ elif lscpu | grep -i avx2
+ then
+ echo "avx2 available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false"
+ else
+ echo "avx512 and avx2 not available on system"
+ su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`"
+ fi
diff --git a/.gitignore b/.gitignore
index 160757c41..34f7ff154 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,9 +17,7 @@ oss/*
jni/CMakeCache.txt
jni/CMakeFiles
jni/Makefile
-jni/cmake*
-jni/CPack*
-jni/_CPack*
+jni/cmake_install.cmake
jni/release
jni/packages
jni/CTestTestfile.cmake
@@ -30,6 +28,12 @@ jni/bin/
jni/lib/
jni/jni_test*
jni/googletest*
+jni/cmake/*.cmake-e
+jni/.cmake
+jni/.idea
+jni/build.ninja
+jni/.ninja_deps
+jni/.ninja_log
benchmarks/perf-tool/okpt/output
benchmarks/perf-tool/okpt/dev
diff --git a/.idea/copyright/SPDX_ALv2.xml b/.idea/copyright/SPDX_ALv2.xml
index a2485beef..3c6f9c768 100644
--- a/.idea/copyright/SPDX_ALv2.xml
+++ b/.idea/copyright/SPDX_ALv2.xml
@@ -1,6 +1,6 @@
-
+
\ No newline at end of file
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 000000000..a39546843
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,29 @@
+
+# CHANGELOG
+All notable changes to this project are documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). See the [CONTRIBUTING guide](./CONTRIBUTING.md#Changelog) for instructions on how to add changelog entries.
+
+## [Unreleased 3.0](https://github.com/opensearch-project/k-NN/compare/2.x...HEAD)
+### Features
+### Enhancements
+### Bug Fixes
+### Infrastructure
+### Documentation
+### Maintenance
+### Refactoring
+
+## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.18...2.x)
+### Features
+### Enhancements
+- Introduced a writing layer in native engines where relies on the writing interface to process IO. (#2241)[https://github.com/opensearch-project/k-NN/pull/2241]
+### Bug Fixes
+* Fixing the bug when a segment has no vector field present for disk based vector search (#2282)[https://github.com/opensearch-project/k-NN/pull/2282]
+### Infrastructure
+* Updated C++ version in JNI from c++11 to c++17 [#2259](https://github.com/opensearch-project/k-NN/pull/2259)
+* Upgrade bytebuddy and objenesis version to match OpenSearch core and, update github ci runner for macos [#2279](https://github.com/opensearch-project/k-NN/pull/2279)
+### Documentation
+### Maintenance
+* Select index settings based on cluster version[2236](https://github.com/opensearch-project/k-NN/pull/2236)
+* Added null checks for fieldInfo in ExactSearcher to avoid NPE while running exact search for segments with no vector field (#2278)[https://github.com/opensearch-project/k-NN/pull/2278]
+### Refactoring
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 958c2986d..49c6b1d9b 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -8,6 +8,7 @@ OpenSearch k-NN is a community project that is built and maintained by people ju
- [Ways to Contribute](#ways-to-contribute)
- [Developer Certificate of Origin](#developer-certificate-of-origin)
- [License Headers](#license-headers)
+- [Changelog](#changelog)
- [Review Process](#review-process)
@@ -129,6 +130,42 @@ New files in your code contributions should contain the following license header
# SPDX-License-Identifier: Apache-2.0
```
+## Changelog
+
+OpenSearch maintains version specific changelog by enforcing a change to the ongoing [CHANGELOG](CHANGELOG.md) file adhering to the [Keep A Changelog](https://keepachangelog.com/en/1.0.0/) format. The purpose of the changelog is for the contributors and maintainers to incrementally build the release notes throughout the development process to avoid a painful and error-prone process of attempting to compile the release notes at release time. On each release the "unreleased" entries of the changelog are moved to the appropriate release notes document in the `./release-notes` folder. Also, incrementally building the changelog provides a concise, human-readable list of significant features that have been added to the unreleased version under development.
+
+### Which changes require a CHANGELOG entry?
+Changelogs are intended for operators/administrators, developers integrating with libraries and APIs, and end-users interacting with OpenSearch Dashboards and/or the REST API (collectively referred to as "user"). In short, any change that a user of OpenSearch might want to be aware of should be included in the changelog. The changelog is _not_ intended to replace the git commit log that developers of OpenSearch itself rely upon. The following are some examples of changes that should be in the changelog:
+
+- A newly added feature
+- A fix for a user-facing bug
+- Dependency updates
+- Fixes for security issues
+
+The following are some examples where a changelog entry is not necessary:
+
+- Adding, modifying, or fixing tests
+- An incremental PR for a larger feature (such features should include _one_ changelog entry for the feature)
+- Documentation changes or code refactoring
+- Build-related changes
+
+Any PR that does not include a changelog entry will result in a failure of the validation workflow in GitHub. If the contributor and maintainers agree that no changelog entry is required, then the `skip-changelog` label can be applied to the PR which will result in the workflow passing.
+
+### How to add my changes to [CHANGELOG](CHANGELOG.md)?
+
+Adding in the change is two step process:
+1. Add your changes to the corresponding section within the CHANGELOG file with dummy pull request information, publish the PR
+2. Update the entry for your change in [`CHANGELOG.md`](CHANGELOG.md) and make sure that you reference the pull request there.
+
+### Where should I put my CHANGELOG entry?
+Please review the [branching strategy](https://github.com/opensearch-project/.github/blob/main/RELEASING.md#opensearch-branching) document. The changelog on the `main` branch will contain sections for the _next major_ and _next minor_ releases. Your entry should go into the section it is intended to be released in. In practice, most changes to `main` will be backported to the next minor release so most entries will likely be in that section.
+
+The following examples assume the _next major_ release on main is 3.0, then _next minor_ release is 2.5, and the _current_ release is 2.4.
+
+- **Add a new feature to release in next minor:** Add a changelog entry to `[Unreleased 2.x]` on main, then backport to 2.x (including the changelog entry).
+- **Introduce a breaking API change to release in next major:** Add a changelog entry to `[Unreleased 3.0]` on main, do not backport.
+- **Upgrade a dependency to fix a CVE:** Add a changelog entry to `[Unreleased 2.x]` on main, then backport to 2.x (including the changelog entry), then backport to 2.4 and ensure the changelog entry is added to `[Unreleased 2.4.1]`.
+
## Review Process
We deeply appreciate everyone who takes the time to make a contribution. We will review all contributions as quickly as possible. As a reminder, [opening an issue](https://github.com/opensearch-project/k-NN/issues/new/choose) discussing your change before you make it is the best way to smooth the PR process. This will prevent a rejection because someone else is already working on the problem, or because the solution is incompatible with the architectural direction.
diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md
index 291b902b8..c39a6f1b6 100644
--- a/DEVELOPER_GUIDE.md
+++ b/DEVELOPER_GUIDE.md
@@ -3,11 +3,16 @@
- [Fork OpenSearch k-NN Repo](#fork-opensearch-k-nn-repo)
- [Install Prerequisites](#install-prerequisites)
- [JDK 11](#jdk-11)
+ - [CMake](#cmake)
+ - [Faiss Dependencies](#Faiss-Dependencies)
+ - [Environment](#Environment)
- [Use an Editor](#use-an-editor)
- [IntelliJ IDEA](#intellij-idea)
- [Build](#build)
- [JNI Library](#jni-library)
- [JNI Library Artifacts](#jni-library-artifacts)
+ - [Parallelize make](#parallelize-make)
+ - [Enable SIMD Optimization](#enable-simd-optimization)
- [Run OpenSearch k-NN](#run-opensearch-k-nn)
- [Run Single-node Cluster Locally](#run-single-node-cluster-locally)
- [Run Multi-node Cluster Locally](#run-multi-node-cluster-locally)
@@ -53,13 +58,90 @@ In addition to this, the plugin has been tested with JDK 17, and this JDK versio
#### CMake
-The plugin requires that cmake >= 3.17.2 is installed in order to build the JNI libraries.
+The plugin requires that cmake >= 3.24.0 is installed in order to build the JNI libraries.
+
+One easy way to install on mac or linux is to use pip:
+```bash
+pip install cmake==3.24.0
+```
+
+On Mac M series machines, install cmake using:
+```bash
+brew install cmake
+```
#### Faiss Dependencies
To build the *faiss* JNI library, you need to have openmp, lapack and blas installed. For more information on *faiss*
dependencies, please refer to [their documentation](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md).
+[Openblas](https://www.openblas.net/) can be used for both lapack and blas. To install on Mac, run:
+```bash
+brew install openblas
+```
+
+Additionally, the `gcc` toolchain needs to be installed on Mac. To install, run:
+```bash
+brew install gcc
+```
+
+#### Additional setup for Mac M series Machines
+
+The following commands enable running/building k-NN on M series machines:
+
+```bash
+// Go to k-NN folder
+cd k-NN
+
+// Build to generate the necessary files to be modified below (will fail)
+./gradlew build
+
+//Go to jni folder
+cd jni
+
+// File changes required
+sed -i -e 's/\/usr\/local\/opt\/libomp\//\/opt\/homebrew\/opt\/llvm\//g' cmake/init-faiss.cmake
+sed -i -e 's/__aarch64__/__undefine_aarch64__/g' external/faiss/faiss/utils/distances_simd.cpp
+sed -i -e 's/pragma message WARN/pragma message /g' external/nmslib/similarity_search/src/distcomp_scalar.cc
+// Change -mcpu value to use chip version according to your M series, for example, -mcpu=apple-m1
+sed -i -e 's/-march=native/-mcpu=apple-m1/g' external/nmslib/similarity_search/CMakeLists.txt
+sed -i -e 's/-mcpu=apple-a14/-mcpu=apple-m1/g' external/nmslib/python_bindings/setup.py
+
+// Install llvm
+brew install llvm
+echo 'export PATH="/opt/homebrew/opt/llvm/bin:$PATH"' >> ~/.zshrc
+source ~/.zshrc
+
+// Set compiler path for CMAKE
+export CC=/opt/homebrew/opt/llvm/bin/clang
+export CXX=/opt/homebrew/opt/llvm/bin/clang++
+
+// In case of linking issues with the external libraries and clang, you can try setting the CMAKE compiler to gcc/g++ instead through the following commands:
+export CC=gcc
+export CXX=g++
+sed -i '' '/set(CMAKE_CXX_STANDARD_REQUIRED True)/a\'$'\n''set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Xclang -fopenmp -L/opt/homebrew/opt/libomp/lib -I/opt/homebrew/opt/libomp/include -lomp -arch arm64 -fexceptions")'$'\n''' CMakeLists.txt
+
+// Build
+cmake . --fresh
+make
+```
+
+Next, obtain a minimum distribution tarball of the k-NN version you want to build:
+
+1. Fork the [OpenSearch Repo](https://github.com/opensearch-project/OpenSearch) into your github account.
+2. Clone the repository locally
+3. Run the following commands:
+```cd OpenSearch && ./gradlew -p distribution/archives/darwin-tar assemble```
+4. You should see a opensearch-min--SNAPSHOT-darwin-x64.tar.gz file present in distribution/archives/darwin-tar/build/distributions/
+5. Build k-NN by passing the OpenSearch distribution path in `./gradlew -PcustomDistributionUrl=""`
+
+If you want to start OpenSearch directly on Mac M series, make sure to use JDK for ARM. Otherwise, you will see the following error: `mach-o file, but is an incompatible architecture (have 'arm64', need 'x86_64')`. It is better to start OpenSearch by running `bash opensearch-tar-install.sh` instead of `./bin/opensearch`. To run `./bin/opensearch`, the environment variable `JAVA_LIBRARY_PATH` needs to be set correctly so that OpenSearch can find the JNI library:
+
+```
+export OPENSEARCH_HOME=the directory of opensearch...
+export JAVA_LIBRARY_PATH=$JAVA_LIBRARY_PATH:$OPENSEARCH_HOME/plugins/opensearch-knn/lib
+```
+
#### Environment
Currently, the plugin only supports Linux on x64 and arm platforms.
@@ -96,19 +178,32 @@ Please follow these formatting guidelines:
* Wildcard imports (`import foo.bar.baz.*`) are forbidden and will cause the build to fail.
* If *absolutely* necessary, you can disable formatting for regions of code with the `// tag::NAME` and `// end::NAME` directives, but note that these are intended for use in documentation, so please make it clear what you have done, and only do this where the benefit clearly outweighs the decrease in consistency.
* Note that JavaDoc and block comments i.e. `/* ... */` are not formatted, but line comments i.e `// ...` are.
-* There is an implicit rule that negative boolean expressions should use the form `foo == false` instead of `!foo` for better readability of the code. While this isn't strictly enforced, if might get called out in PR reviews as something to change.
+* There is an implicit rule that negative boolean expressions should use the form `foo == false` instead of `!foo` for better readability of the code. While this isn't strictly enforced, it might get called out in PR reviews as something to change.
## Build
OpenSearch k-NN uses a [Gradle](https://docs.gradle.org/6.6.1/userguide/userguide.html) wrapper for its build.
Run `gradlew` on Unix systems.
+Tests use `JAVA11_HOME` environment variable, make sure to add it in the export path else the tests might fail.
+e.g
+```
+echo "export JAVA11_HOME=" >> ~/.zshrc
+source ~/.zshrc
+```
+
Build OpenSearch k-NN using `gradlew build`
```
./gradlew build
```
+For Mac M series machines use
+```
+./gradlew build -PcustomDistributionUrl=""
+```
+
+
### JNI Library
The plugin relies on 2 JNI libraries to perform approximate k-NN search. `./gradlew build` will first build the
@@ -121,11 +216,11 @@ To build the JNI Library manually, follow these steps:
cd jni
cmake .
-# To build everything, including tests
+# To build everything, including tests. If your computer has multiple cores you can speed it up by building in parallel using make -j 2 (or a higher number for more parallelism)
make
# To just build the libraries
-make opensearchknn_nmslib opensearchknn_nmslib
+make opensearchknn_faiss opensearchknn_nmslib
```
The libraries will be placed in the `jni/release` directory.
@@ -138,10 +233,10 @@ run:
./bin/jni_test
# To run nmslib tests
-./bin/jni_test --gtest_filter=Nmslib*
+./bin/jni_test --gtest_filter='Nmslib*'
# To run faiss tests
-./bin/jni_test --gtest_filter=Faiss*
+./bin/jni_test --gtest_filter='Faiss*'
```
### JNI Library Artifacts
@@ -155,14 +250,67 @@ For users that want to get the most out of the libraries, they should follow [th
and build the libraries from source in their production environment, so that if their environment has optimized
instruction sets, they take advantage of them.
+### Custom patch on JNI Library
+If you want to make a custom patch on JNI library
+1. Make a change on top of current version of JNI library and push the commit locally.
+2. Create a patch file for the change using `git format-patch -o patches HEAD^`
+3. Place the patch file under `jni/patches`
+4. Make a change in `jni/CmakeLists.txt`, `.github/workflows/CI.yml` to apply the patch during build
+
+By default, in the cmake build system, these patches will be applied and committed to the native libraries. In order to
+successfully make the commits the `user.name` and `user.email` git configurations need to be setup. If you cannot set
+these in your environment, you can disable committing the changes to the library by passing gradle this flag:
+`build.lib.commit_patches=false`. For example, `gradlew build -Dbuild.lib.commit_patches=false`. If the patches are
+not committed, then the full library build process will run each time `cmake` is invoked. In a development environment,
+it is recommended to setup the user git configuration to avoid this cost.
+
+### Parallelize make
+When we are building the plugin for the first time, it takes some time to build the JNI libraries. We can parallelize make and speed up the build time by setting and passing
+this flag to gradle, `nproc.count` if your computer has more number of cores (greater than or equal to 2).
+```
+# While building OpenSearch k-NN
+./gradlew build -Dnproc.count=4
+
+# While running OpenSearch k-NN
+./gradlew run -Dnproc.count=4
+
+# When building the JNI library manually
+cd jni
+cmake .
+# Pass the processor count with make using `-j`
+make -j 4
+```
+
+### Enable SIMD Optimization
+SIMD(Single Instruction/Multiple Data) Optimization is enabled by default on Linux and Mac which boosts the performance
+by enabling `AVX2` and `AVX512` on `x86 architecture` and `NEON` on `ARM64 architecture` where applicable while building the Faiss library. But to enable SIMD,
+the underlying processor should support these capabilities (AVX512, AVX2 or NEON). It can be disabled by setting the parameter `avx2.enabled` to `false` and
+`avx512.enabled` to `false`. If your processor supports `AVX512` or `AVX2`, they can be set by enabling the setting . By default, these values are enabled on
+OpenSearch. Some exceptions: As of now, SIMD support is not supported on Windows OS, and AVX512 is not present on MAC systems due to hardware not supporting the
+feature.
+
+```
+# While building OpenSearch k-NN
+./gradlew build -Davx2.enabled=true -Davx512.enabled=true
+
+# While running OpenSearch k-NN
+./gradlew run -Davx2.enabled=true -Davx512.enabled=true
+
+# While building the JNI libraries
+cd jni
+cmake . -DAVX2_ENABLED=true -DAVX512_ENABLED=true
+```
+
## Run OpenSearch k-NN
### Run Single-node Cluster Locally
-Run OpenSearch k-NN using `gradlew run`.
+Run OpenSearch k-NN using `gradlew run`. For Mac M series add ```-PcustomDistributionUrl=``` argument.
```shell script
./gradlew run
```
+
+
That will build OpenSearch and start it, writing its log above Gradle's status message. We log a lot of stuff on startup, specifically these lines tell you that plugin is ready.
```
[2020-05-29T14:50:35,167][INFO ][o.e.h.AbstractHttpServerTransport] [runTask-0] publish_address {127.0.0.1:9200}, bound_addresses {[::1]:9200}, {127.0.0.1:9200}
@@ -190,6 +338,35 @@ curl localhost:9200
}
}
```
+
+Additionally, it is also possible to run a cluster with security enabled:
+```shell script
+./gradlew run -Dsecurity.enabled=true -Dhttps=true -Duser=admin -Dpassword=
+```
+
+Then, to access the cluster, we can run
+```bash
+curl https://localhost:9200 --insecure -u admin:
+
+{
+ "name" : "integTest-0",
+ "cluster_name" : "integTest",
+ "cluster_uuid" : "kLsNk4JDTMyp1yQRqog-3g",
+ "version" : {
+ "distribution" : "opensearch",
+ "number" : "3.0.0-SNAPSHOT",
+ "build_type" : "tar",
+ "build_hash" : "9d85e566894ef53e5f2093618b3d455e4d0a04ce",
+ "build_date" : "2023-10-30T18:34:06.996519Z",
+ "build_snapshot" : true,
+ "lucene_version" : "9.8.0",
+ "minimum_wire_compatibility_version" : "2.12.0",
+ "minimum_index_compatibility_version" : "2.0.0"
+ },
+ "tagline" : "The OpenSearch Project: https://opensearch.org/"
+}
+```
+
### Run Multi-node Cluster Locally
It can be useful to test and debug on a multi-node cluster. In order to launch a 3 node cluster with the KNN plugin installed, run the following command:
@@ -198,12 +375,17 @@ It can be useful to test and debug on a multi-node cluster. In order to launch a
./gradlew run -PnumNodes=3
```
-In order to run the integration tests with a 3 node cluster, run this command:
+In order to run the integration tests, run this command:
```
./gradlew :integTest -PnumNodes=3
```
+Additionally, to run integration tests with security enabled, run
+```
+./gradlew :integTest -Dsecurity.enabled=true -PnumNodes=3
+```
+
Integration tests can be run with remote cluster. For that run the following command and replace host/port/cluster name values with ones for the target cluster:
```
@@ -213,7 +395,7 @@ Integration tests can be run with remote cluster. For that run the following com
In case remote cluster is secured it's possible to pass username and password with the following command:
```
-./gradlew :integTestRemote -Dtests.rest.cluster=localhost:9200 -Dtests.cluster=localhost:9200 -Dtests.clustername="integTest-0" -Dhttps=true -Duser=admin -Dpassword=admin
+./gradlew :integTestRemote -Dtests.rest.cluster=localhost:9200 -Dtests.cluster=localhost:9200 -Dtests.clustername="integTest-0" -Dhttps=true -Duser=admin -Dpassword=
```
### Debugging
@@ -268,10 +450,12 @@ Before adding any new tests to Backward Compatibility Tests, we should be aware
Starting from 2.0 release the new versioning for codec has been introduced. Two positions will be used to define the version,
in format 'X.Y', where 'X' corresponds to underlying version of Lucene and 'Y' is the version of the format.
+Please note that Lucene version along with corresponding Lucene codec is part of the core OpenSearch. KNN codec should be in sync with Lucene codec version from core OpenSearch.
Codec version is used in following classes and methods:
- org.opensearch.knn.index.codec.KNNXYCodec.KNNXYCodec
-- org.opensearch.knn.index.codec.KNNFormatFactory.createKNNXYFormat
+- org.opensearch.knn.index.codec.KNNXYCodec.KNNXYPerFieldKnnVectorsFormat
+- org.opensearch.knn.index.codec.KNNCodecVersion
These classes and methods are tied directly to Lucene version represented by 'X' part.
Other classes use the delegate pattern so no direct tie to Lucene version are related to format and represented by 'Y'
diff --git a/LICENSE.txt b/LICENSE.txt
index d64569567..62364a94c 100644
--- a/LICENSE.txt
+++ b/LICENSE.txt
@@ -200,3 +200,205 @@
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+
+---------------
+OpenSearch k-NN plugin includes the following third-party software/licensing:
+
+** faiss; version 1.7.6 --
+MIT License
+
+Copyright (c) Facebook, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+** OpenBLAS; version 0.3.3, 0.3.21
+Copyright (c) 2011-2014, The OpenBLAS Project
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+ 1. Redistributions of source code must retain the above copyright
+ notice, this list of conditions and the following disclaimer.
+
+ 2. Redistributions in binary form must reproduce the above copyright
+ notice, this list of conditions and the following disclaimer in
+ the documentation and/or other materials provided with the
+ distribution.
+ 3. Neither the name of the OpenBLAS project nor the names of
+ its contributors may be used to endorse or promote products
+ derived from this software without specific prior written
+ permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
+USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+Cerberus
+
+ISC License
+
+Copyright (c) 2012-2016 Nicola Iarocci.
+
+Permission to use, copy, modify, and/or distribute this software for any
+purpose with or without fee is hereby granted, provided that the above
+copyright notice and this permission notice appear in all copies.
+
+THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
+FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+PERFORMANCE OF THIS SOFTWARE.
+
+------
+PyYAML
+
+
+Copyright (c) 2017-2021 Ingy döt Net
+Copyright (c) 2006-2016 Kirill Simonov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+------
+
+Copyright (c) 2005-2021, NumPy Developers.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+ * Redistributions of source code must retain the above copyright
+ notice, this list of conditions and the following disclaimer.
+
+ * Redistributions in binary form must reproduce the above
+ copyright notice, this list of conditions and the following
+ disclaimer in the documentation and/or other materials provided
+ with the distribution.
+
+ * Neither the name of the NumPy Developers nor the names of any
+ contributors may be used to endorse or promote products derived
+ from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+------
+
+** libwinpthread; version 12.2.0 -- https://www.mingw-w64.org/
+Copyright (c) 2011 mingw-w64 project
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.
+
+
+/*
+ * Parts of this library are derived by:
+ *
+ * Posix Threads library for Microsoft Windows
+ *
+ * Use at own risk, there is no implied warranty to this code.
+ * It uses undocumented features of Microsoft Windows that can change
+ * at any time in the future.
+ *
+ * (C) 2010 Lockless Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+modification,
+ * are permitted provided that the following conditions are met:
+ *
+ *
+ * * Redistributions of source code must retain the above copyright notice,
+ * this list of conditions and the following disclaimer.
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ * * Neither the name of Lockless Inc. nor the names of its contributors may be
+ * used to endorse or promote products derived from this software without
+ * specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AN
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED.
+ * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
+DIRECT,
+ * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+OF
+ * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+NEGLIGENCE
+ * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+ADVISED
+ * OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
diff --git a/MAINTAINERS.md b/MAINTAINERS.md
index 6c8ef5ec2..73fdc028a 100644
--- a/MAINTAINERS.md
+++ b/MAINTAINERS.md
@@ -1,71 +1,16 @@
-- [Overview](#overview)
-- [Current Maintainers](#current-maintainers)
-- [Maintainer Responsibilities](#maintainer-responsibilities)
- - [Uphold Code of Conduct](#uphold-code-of-conduct)
- - [Prioritize Security](#prioritize-security)
- - [Review Pull Requests](#review-pull-requests)
- - [Triage Open Issues](#triage-open-issues)
- - [Be Responsive](#be-responsive)
- - [Maintain Overall Health of the Repo](#maintain-overall-health-of-the-repo)
- - [Use Semver](#use-semver)
- - [Release Frequently](#release-frequently)
- - [Promote Other Maintainers](#promote-other-maintainers)
-
## Overview
-This document explains who the maintainers are (see below), what they do in this repo, and how they should be doing it. If you're interested in contributing, see [CONTRIBUTING](CONTRIBUTING.md).
-
+This document contains a list of maintainers in this repo. See [opensearch-project/.github/RESPONSIBILITIES.md](https://github.com/opensearch-project/.github/blob/main/RESPONSIBILITIES.md#maintainer-responsibilities) that explains what the role of maintainer means, what maintainers do in this and other repos, and how they should be doing it. If you're interested in contributing, and becoming a maintainer, see [CONTRIBUTING](CONTRIBUTING.md).
## Current Maintainers
-| Maintainer | GitHub ID | Affiliation |
-| ------------------------| --------------------------------------------| ----------|
-| Jack Mazanec | [jmazanec15](https://github.com/jmazanec15) | Amazon |
-| Vamshi Vijay Nakkirtha | [vamshin](https://github.com/vamshin) | Amazon |
-| Vijayan Balasubramanian | [VijayanB](https://github.com/VijayanB) | Amazon |
-
-## Maintainer Responsibilities
-
-Maintainers are active and visible members of the community, and have [maintain-level permissions on a repository](https://docs.github.com/en/organizations/managing-access-to-your-organizations-repositories/repository-permission-levels-for-an-organization). Use those privileges to serve the community and evolve code as follows.
-
-### Uphold Code of Conduct
-
-Model the behavior set forward by the [Code of Conduct](CODE_OF_CONDUCT.md) and raise any violations to other maintainers and admins.
-
-### Prioritize Security
-
-Security is your number one priority. Maintainer's Github keys must be password protected securely and any reported security vulnerabilities are addressed before features or bugs.
-
-Note that this repository is monitored and supported 24/7 by Amazon Security, see [Reporting a Vulnerability](SECURITY.md) for details.
-
-### Review Pull Requests
-
-Review pull requests regularly, comment, suggest, reject, merge and close. Accept only high quality pull-requests. Provide code reviews and guidance on incomming pull requests. Don't let PRs be stale and do your best to be helpful to contributors.
-
-### Triage Open Issues
-
-Manage labels, review issues regularly, and triage by labelling them.
-
-All repositories in this organization have a standard set of labels, including `bug`, `documentation`, `duplicate`, `enhancement`, `good first issue`, `help wanted`, `blocker`, `invalid`, `question`, `wontfix`, and `untriaged`, along with release labels, such as `v1.0.0`, `v1.1.0` and `v2.0.0`, and `backport`.
-
-Use labels to target an issue or a PR for a given release, add `help wanted` to good issues for new community members, and `blocker` for issues that scare you or need immediate attention. Request for more information from a submitter if an issue is not clear. Create new labels as needed by the project.
-
-### Be Responsive
-
-Respond to enhancement requests, and forum posts. Allocate time to reviewing and commenting on issues and conversations as they come in.
-
-### Maintain Overall Health of the Repo
-
-Keep the `main` branch at production quality at all times. Backport features as needed. Cut release branches and tags to enable future patches.
-
-### Use Semver
-
-Use and enforce [semantic versioning](https://semver.org/) and do not let breaking changes be made outside of major releases.
-
-### Release Frequently
-
-Make frequent project releases to the community.
-
-### Promote Other Maintainers
-
-Assist, add, and remove [MAINTAINERS](MAINTAINERS.md). Exercise good judgement, and propose high quality contributors to become co-maintainers.
-
+| Maintainer | GitHub ID | Affiliation |
+|-------------------------|-------------------------------------------------------|-------------|
+| Heemin Kim | [heemin32](https://github.com/heemin32) | Amazon |
+| Jack Mazanec | [jmazanec15](https://github.com/jmazanec15) | Amazon |
+| Junqiu Lei | [junqiu-lei](https://github.com/junqiu-lei) | Amazon |
+| Martin Gaievski | [martin-gaievski](https://github.com/martin-gaievski) | Amazon |
+| Naveen Tatikonda | [naveentatikonda](https://github.com/naveentatikonda) | Amazon |
+| Navneet Verma | [navneet1v](https://github.com/navneet1v) | Amazon |
+| Vamshi Vijay Nakkirtha | [vamshin](https://github.com/vamshin) | Amazon |
+| Vijayan Balasubramanian | [VijayanB](https://github.com/VijayanB) | Amazon |
+| Yuncheng Lu | [luyuncheng](https://github.com/luyuncheng) | Bytedance |
diff --git a/README.md b/README.md
index bb8ca928b..bd4baac00 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
[](https://github.com/opensearch-project/k-NN/actions/workflows/CI.yml)
[](https://codecov.io/gh/opensearch-project/k-NN)
[](https://opensearch.org/docs/search-plugins/knn/index/)
-[](https://discuss.opendistrocommunity.dev/c/k-NN/)
+[](https://forum.opensearch.org/c/plugins/k-nn/48)

# OpenSearch k-NN
@@ -21,7 +21,7 @@
* [Project Website](https://opensearch.org/)
* [Downloads](https://opensearch.org/downloads.html).
* [Documentation](https://opensearch.org/docs/search-plugins/knn/index/)
-* Need help? Try [Forums](https://discuss.opendistrocommunity.dev/c/k-nn/)
+* Need help? Try the [Forum](https://forum.opensearch.org/c/plugins/k-nn/48)
* [Project Principles](https://opensearch.org/#principles)
* [Contributing to OpenSearch k-NN](CONTRIBUTING.md)
* [Maintainer Responsibilities](MAINTAINERS.md)
diff --git a/SECURITY.md b/SECURITY.md
index 0b85ca04e..be4ac7463 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -1,3 +1,3 @@
## Reporting a Vulnerability
-If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/) or directly via email to aws-security@amazon.com. Please do **not** create a public GitHub issue.
\ No newline at end of file
+If you discover a potential security issue in this project we ask that you notify OpenSearch Security directly via email to security@opensearch.org. Please do **not** create a public GitHub issue.
diff --git a/THIRD-PARTY b/THIRD-PARTY
deleted file mode 100644
index 2add2c0c2..000000000
--- a/THIRD-PARTY
+++ /dev/null
@@ -1,129 +0,0 @@
-** faiss; version 1.7.6 --
-MIT License
-
-Copyright (c) Facebook, Inc. and its affiliates.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
-** OpenBLAS; version 0.3.3
-Copyright (c) 2011-2014, The OpenBLAS Project
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
- 1. Redistributions of source code must retain the above copyright
- notice, this list of conditions and the following disclaimer.
-
- 2. Redistributions in binary form must reproduce the above copyright
- notice, this list of conditions and the following disclaimer in
- the documentation and/or other materials provided with the
- distribution.
- 3. Neither the name of the OpenBLAS project nor the names of
- its contributors may be used to endorse or promote products
- derived from this software without specific prior written
- permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
-LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
-USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-Cerberus
-
-ISC License
-
-Copyright (c) 2012-2016 Nicola Iarocci.
-
-Permission to use, copy, modify, and/or distribute this software for any
-purpose with or without fee is hereby granted, provided that the above
-copyright notice and this permission notice appear in all copies.
-
-THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
-FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-PERFORMANCE OF THIS SOFTWARE.
-
-------
-PyYAML
-
-
-Copyright (c) 2017-2021 Ingy döt Net
-Copyright (c) 2006-2016 Kirill Simonov
-
-Permission is hereby granted, free of charge, to any person obtaining a copy of
-this software and associated documentation files (the "Software"), to deal in
-the Software without restriction, including without limitation the rights to
-use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
-of the Software, and to permit persons to whom the Software is furnished to do
-so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
-------
-
-Copyright (c) 2005-2021, NumPy Developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are
-met:
-
- * Redistributions of source code must retain the above copyright
- notice, this list of conditions and the following disclaimer.
-
- * Redistributions in binary form must reproduce the above
- copyright notice, this list of conditions and the following
- disclaimer in the documentation and/or other materials provided
- with the distribution.
-
- * Neither the name of the NumPy Developers nor the names of any
- contributors may be used to endorse or promote products derived
- from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/TRIAGING.md b/TRIAGING.md
new file mode 100644
index 000000000..9949b3390
--- /dev/null
+++ b/TRIAGING.md
@@ -0,0 +1,89 @@
+
+
+The maintainers of the k-NN/neural-search Repo's seek to promote an inclusive and engaged community of contributors. In
+order to facilitate this, bi-weekly triage meetings are open-to-all and attendance is encouraged for anyone who hopes to
+contribute, discuss an issue, or learn more about the project. To learn more about contributing to the
+k-NN/neural-search Repo visit the [Contributing](./CONTRIBUTING.md) documentation.
+
+### Do I need to attend for my issue to be addressed/triaged?
+
+Attendance is not required for your issue to be triaged or addressed. All new issues are triaged bi-weekly.
+
+### What happens if my issue does not get covered this time?
+
+Each meeting we seek to address all new issues. However, should we run out of time before your issue is discussed, you
+are always welcome to attend the next meeting or to follow up on the issue post itself.
+
+### How do I join the Backlog & Triage meeting?
+
+Meetings are hosted regularly at 5 PM Pacific Time on Tuesdays bi-weekly and can be joined via the links posted on the
+[OpenSearch Meetup Group](https://www.meetup.com/opensearch/events/) list of events. The event will be titled
+`Development Backlog & Triage Meeting - k-NN/neural-search`.
+
+After joining the Chime meeting, you can enable your video / voice to join the discussion. If you do not have a webcam
+or microphone available, you can still join in via the text chat.
+
+If you have an issue you'd like to bring forth please consider getting a link to the issue so it can be presented to
+everyone in the meeting.
+
+### Is there an agenda for each week?
+
+Meetings are 60 minutes and structured as follows:
+
+1. Initial Gathering: As we gather, feel free to turn on video and engage in informal and open-to-all conversation. After a bit a volunteer will share their screen and proceed with the agenda.
+2. Announcements: If there are any announcements to be made they will happen at the start of the meeting.
+3. Review of New Issues: The meetings always start with reviewing all untriaged/recent issues for the k-NN and neural-search repositories.
+4. Member Requests: Opportunity for any meeting member to ask for consideration of an issue or pull request.
+5. Pull Request Discussion: Then, we review the status of outstanding pull requests from the k-NN and neural-search repositories.
+6. Open Discussion: Allow for members of the meeting to surface any topics without issues filed or pull request created.
+
+
+There is no specific ordering within each category.
+
+If you have an issue you would like to discuss but do not have the ability to attend the entire meeting please attend when is best for you and signal that you have an issue to discuss when you arrive.
+
+### Do I need to have already contributed to the project to attend a triage meeting?
+
+No, all are welcome and encouraged to attend. Attending the Backlog & Triage meetings is a great way for a new contributor to learn about the project as well as explore different avenues of contribution.
+
+### What if I have an issue that is almost a duplicate, should I open a new one to be triaged?
+
+You can always open an issue including one that you think may be a duplicate. However, in cases where you believe there
+is an important distinction to be made between an existing issue and your newly created one, you are encouraged to
+attend the triaging meeting to explain.
+
+### What if I have follow-up questions on an issue?
+
+If you have an existing issue you would like to discuss, you can always comment on the issue itself. Alternatively, you
+are welcome to come to the triage meeting to discuss.
+
+### Is this meeting a good place to get help setting up k-NN/neural-search features on my OpenSearch instance?
+
+While we are always happy to help the community, the best resource for implementation questions is [the OpenSearch forum](https://forum.opensearch.org/c/plugins/k-nn/48).
+
+There you can find answers to many common questions as well as speak with implementation experts.
+
+### What are the issue labels associated with triaging?
+
+Yes, there are several labels that are used to identify the 'state' of issues filed in OpenSearch and the Security Plugin.
+
+| Label | When applied | Meaning |
+| ----- | ------------ | ------- |
+| Untriaged | When issues are created or re-opened. | Issues labeled as 'Untriaged' require the attention of the repository maintainers and may need to be prioritized for quicker resolution. It's crucial to keep the count of 'Untriaged' labels low to ensure all potential security issues are addressed in a timely manner. See [SECURITY.md](https://github.com/opensearch-project/security/blob/main/SECURITY.md) for more details on handling these issues. |
+| Triaged | During triage meetings. | Issues labeled as 'Triaged' have been reviewed and are deemed actionable. Opening a pull request for an issue with the 'Triaged' label has a higher likelihood of approval from the project maintainers, particularly in novel areas. |
+| Neither Label | During triage meetings. | This category is for issues that lack sufficient details to formulate a potential solution. Until more details are provided, it's difficult to ascertain if a proposed solution would be acceptable. When dealing with an 'Untriaged' issue that falls into this category, the triage team should provide further insights so the issue can be appropriately closed or labeled as 'Triaged'. Issues in this state are reviewed during every triage meeting. |
+| Help Wanted | Anytime. | Issues marked as 'Help Wanted' signal that they are actionable and not the current focus of the project maintainers. Community contributions are especially encouraged for these issues. |
+| Good First Issue | Anytime. | Issues labeled as 'Good First Issue' are small in scope and can be resolved with a single pull request. These are recommended starting points for newcomers looking to make their first contributions. |
+
+
+### What if my issue is critical to OpenSearch operations, do I have to wait for the bi-weekly meeting for it to be addressed?
+
+All new issues for the [k-NN](https://github.com/opensearch-project/k-NN/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged) repo and [neural-search](https://github.com/opensearch-project/neural-search/issues?q=is%3Aissue+is%3Aopen+-label%3Atriaged) repo are reviewed daily to check for critical issues which require immediate triaging. If an issue relates to a severe concern for OpenSearch operation, it will be triaged by a maintainer mid-week. You can still come to discuss an issue at the following meeting even if it has already been triaged during the week.
+
+### Is this where I should bring up potential security vulnerabilities?
+
+Due to the sensitive nature of security vulnerabilities, please report all potential vulnerabilities directly by following the steps outlined on the [SECURITY.md](https://github.com/opensearch-project/k-NN/blob/main/SECURITY.md) document.
+
+### Who should I contact if I have further questions?
+
+You can always file an issue for any question you have about the project.
diff --git a/benchmarks/README.md b/benchmarks/README.md
new file mode 100644
index 000000000..2e642d41b
--- /dev/null
+++ b/benchmarks/README.md
@@ -0,0 +1,4 @@
+## Benchmark Folder Tools Deprecated
+All benchmark workloads have been moved to [OpenSearch Benchmark Workloads](https://github.com/opensearch-project/opensearch-benchmark-workloads/tree/main/vectorsearch). Please use OSB tool to run the benchmarks.
+
+If you are still interested in using the old tool, the benchmarks are moved to the [branch](https://github.com/opensearch-project/k-NN/tree/old-benchmarks/benchmarks).
diff --git a/benchmarks/osb/README.md b/benchmarks/osb/README.md
deleted file mode 100644
index 92272e20b..000000000
--- a/benchmarks/osb/README.md
+++ /dev/null
@@ -1,468 +0,0 @@
-# OpenSearch Benchmarks for k-NN
-
-## Overview
-
-This directory contains code and configurations to run k-NN benchmarking
-workloads using OpenSearch Benchmarks.
-
-The [extensions](extensions) directory contains common code shared between
-procedures. The [procedures](procedures) directory contains the individual
-test procedures for this workload.
-
-## Getting Started
-
-### OpenSearch Benchmarks Background
-
-OpenSearch Benchmark is a framework for performance benchmarking an OpenSearch
-cluster. For more details, checkout their
-[repo](https://github.com/opensearch-project/opensearch-benchmark/).
-
-Before getting into the benchmarks, it is helpful to know a few terms:
-1. Workload - Top level description of a benchmark suite. A workload will have a `workload.json` file that defines different components of the tests
-2. Test Procedures - A workload can have a schedule of operations that run the test. However, a workload can also have several test procedures that define their own schedule of operations. This is helpful for sharing code between tests
-3. Operation - An action against the OpenSearch cluster
-4. Parameter source - Producers of parameters for OpenSearch operations
-5. Runners - Code that actually will execute the OpenSearch operations
-
-### Setup
-
-OpenSearch Benchmarks requires Python 3.8 or greater to be installed. One of
-the easier ways to do this is through Conda, a package and environment
-management system for Python.
-
-First, follow the
-[installation instructions](https://docs.conda.io/projects/conda/en/latest/user-guide/install/index.html)
-to install Conda on your system.
-
-Next, create a Python 3.8 environment:
-```
-conda create -n knn-osb python=3.8
-```
-
-After the environment is created, activate it:
-```
-source activate knn-osb
-```
-
-Lastly, clone the k-NN repo and install all required python packages:
-```
-git clone https://github.com/opensearch-project/k-NN.git
-cd k-NN/benchmarks/osb
-pip install -r requirements.txt
-```
-
-After all of this completes, you should be ready to run your first benchmark!
-
-### Running a benchmark
-
-Before running a benchmark, make sure you have the endpoint of your cluster and
- the machine you are running the benchmarks from can access it.
- Additionally, ensure that all data has been pulled to the client.
-
-Currently, we support 2 test procedures for the k-NN workload: train-test and
-no-train-test. The train test has steps to train a model included in the
-schedule, while no train does not. Both test procedures will index a data set
-of vectors into an OpenSearch index and then run a set of queries against them.
-
-Once you have decided which test procedure you want to use, open up
-[params/train-params.json](params/train-params.json) or
-[params/no-train-params.json](params/no-train-params.json) and
-fill out the parameters. Notice, at the bottom of `no-train-params.json` there
-are several parameters that relate to training. Ignore these. They need to be
-defined for the workload but not used.
-
-Once the parameters are set, set the URL and PORT of your cluster and run the
-command to run the test procedure.
-
-```
-export URL=
-export PORT=
-export PARAMS_FILE=
-export PROCEDURE={no-train-test | train-test}
-
-opensearch-benchmark execute_test \
- --target-hosts $URL:$PORT \
- --workload-path ./workload.json \
- --workload-params ${PARAMS_FILE} \
- --test-procedure=${PROCEDURE} \
- --pipeline benchmark-only
-```
-
-## Current Procedures
-
-### No Train Test
-
-The No Train Test procedure is used to test `knn_vector` indices that do not
-use an algorithm that requires training.
-
-#### Workflow
-
-1. Delete old resources in the cluster if they are present
-2. Create an OpenSearch index with `knn_vector` configured to use the HNSW algorithm
-3. Wait for cluster to be green
-4. Ingest data set into the cluster
-5. Refresh the index
-6. Run queries from data set against the cluster
-
-#### Parameters
-
-| Name | Description |
-|-----------------------------------------|----------------------------------------------------------------------------------|
-| target_index_name | Name of index to add vectors to |
-| target_field_name | Name of field to add vectors to |
-| target_index_body | Path to target index definition |
-| target_index_primary_shards | Target index primary shards |
-| target_index_replica_shards | Target index replica shards |
-| target_index_dimension | Dimension of target index |
-| target_index_space_type | Target index space type |
-| target_index_bulk_size | Target index bulk size |
-| target_index_bulk_index_data_set_format | Format of vector data set |
-| target_index_bulk_index_data_set_path | Path to vector data set |
-| target_index_bulk_index_clients | Clients to be used for bulk ingestion (must be divisor of data set size) |
-| hnsw_ef_search | HNSW ef search parameter |
-| hnsw_ef_construction | HNSW ef construction parameter |
-| hnsw_m | HNSW m parameter |
-| query_k | The number of neighbors to return for the search |
-| query_clients | Number of clients to use for running queries |
-| query_data_set_format | Format of vector data set for queries |
-| query_data_set_path | Path to vector data set for queries |
-
-#### Metrics
-
-The result metrics of this procedure will look like:
-```
-------------------------------------------------------
- _______ __ _____
- / ____(_)___ ____ _/ / / ___/_________ ________
- / /_ / / __ \/ __ `/ / \__ \/ ___/ __ \/ ___/ _ \
- / __/ / / / / / /_/ / / ___/ / /__/ /_/ / / / __/
-/_/ /_/_/ /_/\__,_/_/ /____/\___/\____/_/ \___/
-------------------------------------------------------
-
-| Metric | Task | Value | Unit |
-|---------------------------------------------------------------:|------------------------:|------------:|-------:|
-| Cumulative indexing time of primary shards | | 0.00173333 | min |
-| Min cumulative indexing time across primary shards | | 0 | min |
-| Median cumulative indexing time across primary shards | | 0 | min |
-| Max cumulative indexing time across primary shards | | 0.000616667 | min |
-| Cumulative indexing throttle time of primary shards | | 0 | min |
-| Min cumulative indexing throttle time across primary shards | | 0 | min |
-| Median cumulative indexing throttle time across primary shards | | 0 | min |
-| Max cumulative indexing throttle time across primary shards | | 0 | min |
-| Cumulative merge time of primary shards | | 0 | min |
-| Cumulative merge count of primary shards | | 0 | |
-| Min cumulative merge time across primary shards | | 0 | min |
-| Median cumulative merge time across primary shards | | 0 | min |
-| Max cumulative merge time across primary shards | | 0 | min |
-| Cumulative merge throttle time of primary shards | | 0 | min |
-| Min cumulative merge throttle time across primary shards | | 0 | min |
-| Median cumulative merge throttle time across primary shards | | 0 | min |
-| Max cumulative merge throttle time across primary shards | | 0 | min |
-| Cumulative refresh time of primary shards | | 0.00271667 | min |
-| Cumulative refresh count of primary shards | | 115 | |
-| Min cumulative refresh time across primary shards | | 0 | min |
-| Median cumulative refresh time across primary shards | | 0 | min |
-| Max cumulative refresh time across primary shards | | 0.00135 | min |
-| Cumulative flush time of primary shards | | 0 | min |
-| Cumulative flush count of primary shards | | 43 | |
-| Min cumulative flush time across primary shards | | 0 | min |
-| Median cumulative flush time across primary shards | | 0 | min |
-| Max cumulative flush time across primary shards | | 0 | min |
-| Total Young Gen GC time | | 0.849 | s |
-| Total Young Gen GC count | | 20 | |
-| Total Old Gen GC time | | 0 | s |
-| Total Old Gen GC count | | 0 | |
-| Store size | | 0.647921 | GB |
-| Translog size | | 0.00247511 | GB |
-| Heap used for segments | | 0.284451 | MB |
-| Heap used for doc values | | 0.0872688 | MB |
-| Heap used for terms | | 0.0714417 | MB |
-| Heap used for norms | | 6.10352e-05 | MB |
-| Heap used for points | | 0 | MB |
-| Heap used for stored fields | | 0.125679 | MB |
-| Segment count | | 257 | |
-| Min Throughput | custom-vector-bulk | 18018.5 | docs/s |
-| Mean Throughput | custom-vector-bulk | 18018.5 | docs/s |
-| Median Throughput | custom-vector-bulk | 18018.5 | docs/s |
-| Max Throughput | custom-vector-bulk | 18018.5 | docs/s |
-| 50th percentile latency | custom-vector-bulk | 98.5565 | ms |
-| 90th percentile latency | custom-vector-bulk | 100.033 | ms |
-| 100th percentile latency | custom-vector-bulk | 103.792 | ms |
-| 50th percentile service time | custom-vector-bulk | 98.5565 | ms |
-| 90th percentile service time | custom-vector-bulk | 100.033 | ms |
-| 100th percentile service time | custom-vector-bulk | 103.792 | ms |
-| error rate | custom-vector-bulk | 0 | % |
-| Min Throughput | refresh-target-index | 76.22 | ops/s |
-| Mean Throughput | refresh-target-index | 76.22 | ops/s |
-| Median Throughput | refresh-target-index | 76.22 | ops/s |
-| Max Throughput | refresh-target-index | 76.22 | ops/s |
-| 100th percentile latency | refresh-target-index | 12.7619 | ms |
-| 100th percentile service time | refresh-target-index | 12.7619 | ms |
-| error rate | refresh-target-index | 0 | % |
-| Min Throughput | knn-query-from-data-set | 1587.47 | ops/s |
-| Mean Throughput | knn-query-from-data-set | 1649.97 | ops/s |
-| Median Throughput | knn-query-from-data-set | 1661.79 | ops/s |
-| Max Throughput | knn-query-from-data-set | 1677.06 | ops/s |
-| 50th percentile latency | knn-query-from-data-set | 4.79125 | ms |
-| 90th percentile latency | knn-query-from-data-set | 5.38 | ms |
-| 99th percentile latency | knn-query-from-data-set | 46.8965 | ms |
-| 99.9th percentile latency | knn-query-from-data-set | 58.2049 | ms |
-| 99.99th percentile latency | knn-query-from-data-set | 59.6476 | ms |
-| 100th percentile latency | knn-query-from-data-set | 60.9245 | ms |
-| 50th percentile service time | knn-query-from-data-set | 4.79125 | ms |
-| 90th percentile service time | knn-query-from-data-set | 5.38 | ms |
-| 99th percentile service time | knn-query-from-data-set | 46.8965 | ms |
-| 99.9th percentile service time | knn-query-from-data-set | 58.2049 | ms |
-| 99.99th percentile service time | knn-query-from-data-set | 59.6476 | ms |
-| 100th percentile service time | knn-query-from-data-set | 60.9245 | ms |
-| error rate | knn-query-from-data-set | 0 | % |
-
-
---------------------------------
-[INFO] SUCCESS (took 46 seconds)
---------------------------------
-```
-
-### Train Test
-
-The Train Test procedure is used to test `knn_vector` indices that do use an
-algorithm that requires training.
-
-#### Workflow
-
-1. Delete old resources in the cluster if they are present
-2. Create an OpenSearch index with `knn_vector` configured to load with training data
-3. Wait for cluster to be green
-4. Ingest data set into the training index
-5. Refresh the index
-6. Train a model based on user provided input parameters
-7. Create an OpenSearch index with `knn_vector` configured to use the model
-8. Ingest vectors into the target index
-9. Refresh the target index
-10. Run queries from data set against the cluster
-
-#### Parameters
-
-| Name | Description |
-|-----------------------------------------|----------------------------------------------------------------------------------|
-| target_index_name | Name of index to add vectors to |
-| target_field_name | Name of field to add vectors to |
-| target_index_body | Path to target index definition |
-| target_index_primary_shards | Target index primary shards |
-| target_index_replica_shards | Target index replica shards |
-| target_index_dimension | Dimension of target index |
-| target_index_space_type | Target index space type |
-| target_index_bulk_size | Target index bulk size |
-| target_index_bulk_index_data_set_format | Format of vector data set for ingestion |
-| target_index_bulk_index_data_set_path | Path to vector data set for ingestion |
-| target_index_bulk_index_clients | Clients to be used for bulk ingestion (must be divisor of data set size) |
-| ivf_nlists | IVF nlist parameter |
-| ivf_nprobes | IVF nprobe parameter |
-| pq_code_size | PQ code_size parameter |
-| pq_m | PQ m parameter |
-| train_model_method | Method to be used for model (ivf or ivfpq) |
-| train_model_id | Model ID |
-| train_index_name | Name of index to put training data into |
-| train_field_name | Name of field to put training data into |
-| train_index_body | Path to train index definition |
-| train_search_size | Search size to use when pulling training data |
-| train_timeout | Timeout to wait for training to finish |
-| train_index_primary_shards | Train index primary shards |
-| train_index_replica_shards | Train index replica shards |
-| train_index_bulk_size | Train index bulk size |
-| train_index_data_set_format | Format of vector data set for training |
-| train_index_data_set_path | Path to vector data set for training |
-| train_index_num_vectors | Number of vectors to use from vector data set for training |
-| train_index_bulk_index_clients | Clients to be used for bulk ingestion (must be divisor of data set size) |
-| query_k | The number of neighbors to return for the search |
-| query_clients | Number of clients to use for running queries |
-| query_data_set_format | Format of vector data set for queries |
-| query_data_set_path | Path to vector data set for queries |
-
-#### Metrics
-
-The result metrics of this procedure will look like:
-```
-------------------------------------------------------
- _______ __ _____
- / ____(_)___ ____ _/ / / ___/_________ ________
- / /_ / / __ \/ __ `/ / \__ \/ ___/ __ \/ ___/ _ \ [63/1855]
- / __/ / / / / / /_/ / / ___/ / /__/ /_/ / / / __/
-/_/ /_/_/ /_/\__,_/_/ /____/\___/\____/_/ \___/
-------------------------------------------------------
-| Metric | Task | Value | Unit |
-|---------------------------------------------------------------:|------------------------:|------------:|-----------------:|
-| Cumulative indexing time of primary shards | | 2.92355 | min |
-| Min cumulative indexing time across primary shards | | 0 | min |
-| Median cumulative indexing time across primary shards | | 0.497817 | min |
-| Max cumulative indexing time across primary shards | | 1.37717 | min |
-| Cumulative indexing throttle time of primary shards | | 0 | min |
-| Min cumulative indexing throttle time across primary shards | | 0 | min |
-| Median cumulative indexing throttle time across primary shards | | 0 | min |
-| Max cumulative indexing throttle time across primary shards | | 0 | min |
-| Cumulative merge time of primary shards | | 1.34895 | min |
-| Cumulative merge count of primary shards | | 39 | |
-| Min cumulative merge time across primary shards | | 0 | min |
-| Median cumulative merge time across primary shards | | 0.292033 | min |
-| Max cumulative merge time across primary shards | | 0.6268 | min |
-| Cumulative merge throttle time of primary shards | | 0.62845 | min |
-| Min cumulative merge throttle time across primary shards | | 0 | min |
-| Median cumulative merge throttle time across primary shards | | 0.155617 | min |
-| Max cumulative merge throttle time across primary shards | | 0.290117 | min |
-| Cumulative refresh time of primary shards | | 0.369433 | min |
-| Cumulative refresh count of primary shards | | 96 | |
-| Min cumulative refresh time across primary shards | | 0 | min |
-| Median cumulative refresh time across primary shards | | 0.0903833 | min |
-| Max cumulative refresh time across primary shards | | 0.10365 | min |
-| Cumulative flush time of primary shards | | 0.0278667 | min |
-| Cumulative flush count of primary shards | | 2 | |
-| Min cumulative flush time across primary shards | | 0 | min |
-| Median cumulative flush time across primary shards | | 0 | min |
-| Max cumulative flush time across primary shards | | 0.0278667 | min |
-| Total Young Gen GC time | | 13.106 | s |
-| Total Young Gen GC count | | 263 | |
-| Total Old Gen GC time | | 0 | s |
-| Total Old Gen GC count | | 0 | |
-| Store size | | 2.60183 | GB |
-| Translog size | | 1.34787 | GB |
-| Heap used for segments | | 0.0646248 | MB |
-| Heap used for doc values | | 0.00899887 | MB |
-| Heap used for terms | | 0.0203552 | MB |
-| Heap used for norms | | 6.10352e-05 | MB |
-| Heap used for points | | 0 | MB |
-| Heap used for stored fields | | 0.0352097 | MB |
-| Segment count | | 71 | |
-| Min Throughput | delete-model | 10.55 | ops/s |
-| Mean Throughput | delete-model | 10.55 | ops/s |
-| Median Throughput | delete-model | 10.55 | ops/s |
-| Max Throughput | delete-model | 10.55 | ops/s |
-| 100th percentile latency | delete-model | 94.4726 | ms |
-| 100th percentile service time | delete-model | 94.4726 | ms |
-| error rate | delete-model | 0 | % |
-| Min Throughput | train-vector-bulk | 44763.1 | docs/s |
-| Mean Throughput | train-vector-bulk | 52022.4 | docs/s |
-| Median Throughput | train-vector-bulk | 52564.8 | docs/s |
-| Max Throughput | train-vector-bulk | 53833 | docs/s |
-| 50th percentile latency | train-vector-bulk | 22.3364 | ms |
-| 90th percentile latency | train-vector-bulk | 47.799 | ms |
-| 99th percentile latency | train-vector-bulk | 195.954 | ms |
-| 99.9th percentile latency | train-vector-bulk | 495.217 | ms |
-| 100th percentile latency | train-vector-bulk | 663.48 | ms |
-| 50th percentile service time | train-vector-bulk | 22.3364 | ms |
-| 90th percentile service time | train-vector-bulk | 47.799 | ms |
-| 99th percentile service time | train-vector-bulk | 195.954 | ms |
-| 99.9th percentile service time | train-vector-bulk | 495.217 | ms |
-| 100th percentile service time | train-vector-bulk | 663.48 | ms |
-| error rate | train-vector-bulk | 0 | % |
-| Min Throughput | refresh-train-index | 0.98 | ops/s |
-| Mean Throughput | refresh-train-index | 0.98 | ops/s |
-| Median Throughput | refresh-train-index | 0.98 | ops/s |
-| Max Throughput | refresh-train-index | 0.98 | ops/s |
-| 100th percentile latency | refresh-train-index | 1019.54 | ms |
-| 100th percentile service time | refresh-train-index | 1019.54 | ms |
-| error rate | refresh-train-index | 0 | % |
-| Min Throughput | ivfpq-train-model | 0.01 | models_trained/s |
-| Mean Throughput | ivfpq-train-model | 0.01 | models_trained/s |
-| Median Throughput | ivfpq-train-model | 0.01 | models_trained/s |
-| Max Throughput | ivfpq-train-model | 0.01 | models_trained/s |
-| 100th percentile latency | ivfpq-train-model | 150952 | ms |
-| 100th percentile service time | ivfpq-train-model | 150952 | ms |
-| error rate | ivfpq-train-model | 0 | % |
-| Min Throughput | custom-vector-bulk | 32367.4 | docs/s |
-| Mean Throughput | custom-vector-bulk | 36027.5 | docs/s |
-| Median Throughput | custom-vector-bulk | 35276.7 | docs/s |
-| Max Throughput | custom-vector-bulk | 41095 | docs/s |
-| 50th percentile latency | custom-vector-bulk | 22.2419 | ms |
-| 90th percentile latency | custom-vector-bulk | 70.163 | ms |
-| 99th percentile latency | custom-vector-bulk | 308.395 | ms |
-| 99.9th percentile latency | custom-vector-bulk | 548.558 | ms |
-| 100th percentile latency | custom-vector-bulk | 655.628 | ms |
-| 50th percentile service time | custom-vector-bulk | 22.2419 | ms |
-| 90th percentile service time | custom-vector-bulk | 70.163 | ms |
-| 99th percentile service time | custom-vector-bulk | 308.395 | ms |
-| 99.9th percentile service time | custom-vector-bulk | 548.558 | ms |
-| 100th percentile service time | custom-vector-bulk | 655.628 | ms |
-| error rate | custom-vector-bulk | 0 | % |
-| Min Throughput | refresh-target-index | 0.23 | ops/s |
-| Mean Throughput | refresh-target-index | 0.23 | ops/s |
-| Median Throughput | refresh-target-index | 0.23 | ops/s |
-| Max Throughput | refresh-target-index | 0.23 | ops/s |
-| 100th percentile latency | refresh-target-index | 4331.17 | ms |
-| 100th percentile service time | refresh-target-index | 4331.17 | ms |
-| error rate | refresh-target-index | 0 | % |
-| Min Throughput | knn-query-from-data-set | 455.19 | ops/s |
-| Mean Throughput | knn-query-from-data-set | 511.74 | ops/s |
-| Median Throughput | knn-query-from-data-set | 510.85 | ops/s |
-| Max Throughput | knn-query-from-data-set | 570.07 | ops/s |
-| 50th percentile latency | knn-query-from-data-set | 14.1626 | ms |
-| 90th percentile latency | knn-query-from-data-set | 30.2389 | ms |
-| 99th percentile latency | knn-query-from-data-set | 71.2793 | ms |
-| 99.9th percentile latency | knn-query-from-data-set | 104.733 | ms |
-| 99.99th percentile latency | knn-query-from-data-set | 127.298 | ms |
-| 100th percentile latency | knn-query-from-data-set | 145.229 | ms |
-| 50th percentile service time | knn-query-from-data-set | 14.1626 | ms |
-| 90th percentile service time | knn-query-from-data-set | 30.2389 | ms |
-| 99th percentile service time | knn-query-from-data-set | 71.2793 | ms |
-| 99.9th percentile service time | knn-query-from-data-set | 104.733 | ms |
-| 99.99th percentile service time | knn-query-from-data-set | 127.298 | ms |
-| 100th percentile service time | knn-query-from-data-set | 145.229 | ms |
-| error rate | knn-query-from-data-set | 0 | % |
-
-
----------------------------------
-[INFO] SUCCESS (took 295 seconds)
----------------------------------
-```
-
-## Adding a procedure
-
-Adding additional benchmarks is very simple. First, place any custom parameter
-sources or runners in the [extensions](extensions) directory so that other tests
-can use them and also update the [documentation](#custom-extensions)
-accordingly.
-
-Next, create a new test procedure file and add the operations you want your test
-to run. Lastly, be sure to update documentation.
-
-## Custom Extensions
-
-OpenSearch Benchmarks is very extendable. To fit the plugins needs, we add
-customer parameter sources and custom runners. Parameter sources allow users to
-supply custom parameters to an operation. Runners are what actually performs
-the operations against OpenSearch.
-
-### Custom Parameter Sources
-
-Custom parameter sources are defined in [extensions/param_sources.py](extensions/param_sources.py).
-
-| Name | Description | Parameters |
-|-------------------------|------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| bulk-from-data-set | Provides bulk payloads containing vectors from a data set for indexing | 1. data_set_format - (hdf5, bigann)
2. data_set_path - path to data set
3. index - name of index for bulk ingestion
4. field - field to place vector in
5. bulk_size - vectors per bulk request
6. num_vectors - number of vectors to use from the data set. Defaults to the whole data set. |
-| knn-query-from-data-set | Provides a query generated from a data set | 1. data_set_format - (hdf5, bigann)
2. data_set_path - path to data set
3. index - name of index to query against
4. field - field to to query against
5. k - number of results to return
6. dimension - size of vectors to produce
7. num_vectors - number of vectors to use from the data set. Defaults to the whole data set. |
-
-
-### Custom Runners
-
-Custom runners are defined in [extensions/runners.py](extensions/runners.py).
-
-| Syntax | Description | Parameters |
-|--------------------|-----------------------------------------------------|:-------------------------------------------------------------------------------------------------------------|
-| custom-vector-bulk | Bulk index a set of vectors in an OpenSearch index. | 1. bulk-from-data-set |
-| custom-refresh | Run refresh with retry capabilities. | 1. index - name of index to refresh
2. retries - number of times to retry the operation |
-| train-model | Trains a model. | 1. body - model definition
2. timeout - time to wait for model to finish
3. model_id - ID of model |
-| delete-model | Deletes a model if it exists. | 1. model_id - ID of model |
-
-### Testing
-
-We have a set of unit tests for our extensions in
-[tests](tests). To run all the tests, run the following
-command:
-
-```commandline
-python -m unittest discover ./tests
-```
-
-To run an individual test:
-```commandline
-python -m unittest tests.test_param_sources.VectorsFromDataSetParamSourceTestCase.test_partition_hdf5
-```
diff --git a/benchmarks/osb/__init__.py b/benchmarks/osb/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/benchmarks/osb/extensions/__init__.py b/benchmarks/osb/extensions/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/benchmarks/osb/extensions/data_set.py b/benchmarks/osb/extensions/data_set.py
deleted file mode 100644
index 7e8058844..000000000
--- a/benchmarks/osb/extensions/data_set.py
+++ /dev/null
@@ -1,202 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-import os
-import numpy as np
-from abc import ABC, ABCMeta, abstractmethod
-from enum import Enum
-from typing import cast
-import h5py
-import struct
-
-
-class Context(Enum):
- """DataSet context enum. Can be used to add additional context for how a
- data-set should be interpreted.
- """
- INDEX = 1
- QUERY = 2
- NEIGHBORS = 3
-
-
-class DataSet(ABC):
- """DataSet interface. Used for reading data-sets from files.
-
- Methods:
- read: Read a chunk of data from the data-set
- seek: Get to position in the data-set
- size: Gets the number of items in the data-set
- reset: Resets internal state of data-set to beginning
- """
- __metaclass__ = ABCMeta
-
- BEGINNING = 0
-
- @abstractmethod
- def read(self, chunk_size: int):
- pass
-
- @abstractmethod
- def seek(self, offset: int):
- pass
-
- @abstractmethod
- def size(self):
- pass
-
- @abstractmethod
- def reset(self):
- pass
-
-
-class HDF5DataSet(DataSet):
- """ Data-set format corresponding to `ANN Benchmarks
- `_
- """
-
- FORMAT_NAME = "hdf5"
-
- def __init__(self, dataset_path: str, context: Context):
- file = h5py.File(dataset_path)
- self.data = cast(h5py.Dataset, file[self.parse_context(context)])
- self.current = self.BEGINNING
-
- def read(self, chunk_size: int):
- if self.current >= self.size():
- return None
-
- end_offset = self.current + chunk_size
- if end_offset > self.size():
- end_offset = self.size()
-
- v = cast(np.ndarray, self.data[self.current:end_offset])
- self.current = end_offset
- return v
-
- def seek(self, offset: int):
-
- if offset < self.BEGINNING:
- raise Exception("Offset must be greater than or equal to 0")
-
- if offset >= self.size():
- raise Exception("Offset must be less than the data set size")
-
- self.current = offset
-
- def size(self):
- return self.data.len()
-
- def reset(self):
- self.current = self.BEGINNING
-
- @staticmethod
- def parse_context(context: Context) -> str:
- if context == Context.NEIGHBORS:
- return "neighbors"
-
- if context == Context.INDEX:
- return "train"
-
- if context == Context.QUERY:
- return "test"
-
- raise Exception("Unsupported context")
-
-
-class BigANNVectorDataSet(DataSet):
- """ Data-set format for vector data-sets for `Big ANN Benchmarks
- `_
- """
-
- DATA_SET_HEADER_LENGTH = 8
- U8BIN_EXTENSION = "u8bin"
- FBIN_EXTENSION = "fbin"
- FORMAT_NAME = "bigann"
-
- BYTES_PER_U8INT = 1
- BYTES_PER_FLOAT = 4
-
- def __init__(self, dataset_path: str):
- self.file = open(dataset_path, 'rb')
- self.file.seek(BigANNVectorDataSet.BEGINNING, os.SEEK_END)
- num_bytes = self.file.tell()
- self.file.seek(BigANNVectorDataSet.BEGINNING)
-
- if num_bytes < BigANNVectorDataSet.DATA_SET_HEADER_LENGTH:
- raise Exception("File is invalid")
-
- self.num_points = int.from_bytes(self.file.read(4), "little")
- self.dimension = int.from_bytes(self.file.read(4), "little")
- self.bytes_per_num = self._get_data_size(dataset_path)
-
- if (num_bytes - BigANNVectorDataSet.DATA_SET_HEADER_LENGTH) != self.num_points * \
- self.dimension * self.bytes_per_num:
- raise Exception("File is invalid")
-
- self.reader = self._value_reader(dataset_path)
- self.current = BigANNVectorDataSet.BEGINNING
-
- def read(self, chunk_size: int):
- if self.current >= self.size():
- return None
-
- end_offset = self.current + chunk_size
- if end_offset > self.size():
- end_offset = self.size()
-
- v = np.asarray([self._read_vector() for _ in
- range(end_offset - self.current)])
- self.current = end_offset
- return v
-
- def seek(self, offset: int):
-
- if offset < self.BEGINNING:
- raise Exception("Offset must be greater than or equal to 0")
-
- if offset >= self.size():
- raise Exception("Offset must be less than the data set size")
-
- bytes_offset = BigANNVectorDataSet.DATA_SET_HEADER_LENGTH + \
- self.dimension * self.bytes_per_num * offset
- self.file.seek(bytes_offset)
- self.current = offset
-
- def _read_vector(self):
- return np.asarray([self.reader(self.file) for _ in
- range(self.dimension)])
-
- def size(self):
- return self.num_points
-
- def reset(self):
- self.file.seek(BigANNVectorDataSet.DATA_SET_HEADER_LENGTH)
- self.current = BigANNVectorDataSet.BEGINNING
-
- def __del__(self):
- self.file.close()
-
- @staticmethod
- def _get_data_size(file_name):
- ext = file_name.split('.')[-1]
- if ext == BigANNVectorDataSet.U8BIN_EXTENSION:
- return BigANNVectorDataSet.BYTES_PER_U8INT
-
- if ext == BigANNVectorDataSet.FBIN_EXTENSION:
- return BigANNVectorDataSet.BYTES_PER_FLOAT
-
- raise Exception("Unknown extension")
-
- @staticmethod
- def _value_reader(file_name):
- ext = file_name.split('.')[-1]
- if ext == BigANNVectorDataSet.U8BIN_EXTENSION:
- return lambda file: float(int.from_bytes(file.read(BigANNVectorDataSet.BYTES_PER_U8INT), "little"))
-
- if ext == BigANNVectorDataSet.FBIN_EXTENSION:
- return lambda file: struct.unpack('= self.num_vectors + self.offset:
- raise StopIteration
-
- if self.vector_batch is None or len(self.vector_batch) == 0:
- self.vector_batch = self._batch_read(self.data_set)
- if self.vector_batch is None:
- raise StopIteration
- vector = self.vector_batch.pop(0)
- self.current += 1
- self.percent_completed = self.current / self.total
-
- return self._build_query_body(self.index_name, self.field_name, self.k,
- vector)
-
- def _batch_read(self, data_set: DataSet):
- return list(data_set.read(self.VECTOR_READ_BATCH_SIZE))
-
- def _build_query_body(self, index_name: str, field_name: str, k: int,
- vector) -> dict:
- """Builds a k-NN query that can be used to execute an approximate nearest
- neighbor search against a k-NN plugin index
- Args:
- index_name: name of index to search
- field_name: name of field to search
- k: number of results to return
- vector: vector used for query
- Returns:
- A dictionary containing the body used for search, a set of request
- parameters to attach to the search and the name of the index.
- """
- return {
- "index": index_name,
- "request-params": {
- "_source": {
- "exclude": [field_name]
- }
- },
- "body": {
- "size": k,
- "query": {
- "knn": {
- field_name: {
- "vector": vector,
- "k": k
- }
- }
- }
- }
- }
-
-
-class BulkVectorsFromDataSetParamSource(VectorsFromDataSetParamSource):
- """ Create bulk index requests from a data set of vectors.
-
- Attributes:
- bulk_size: number of vectors per request
- retries: number of times to retry the request when it fails
- """
-
- DEFAULT_RETRIES = 10
-
- def __init__(self, workload, params, **kwargs):
- super().__init__(params, Context.INDEX)
- self.bulk_size: int = parse_int_parameter("bulk_size", params)
- self.retries: int = parse_int_parameter("retries", params,
- self.DEFAULT_RETRIES)
-
- def params(self):
- """
- Returns: A bulk index parameter with vectors from a data set.
- """
- if self.current >= self.num_vectors + self.offset:
- raise StopIteration
-
- def action(doc_id):
- return {'index': {'_index': self.index_name, '_id': doc_id}}
-
- partition = self.data_set.read(self.bulk_size)
- body = bulk_transform(partition, self.field_name, action, self.current)
- size = len(body) // 2
- self.current += size
- self.percent_completed = self.current / self.total
-
- return {
- "body": body,
- "retries": self.retries,
- "size": size
- }
diff --git a/benchmarks/osb/extensions/registry.py b/benchmarks/osb/extensions/registry.py
deleted file mode 100644
index 5ce17ab6f..000000000
--- a/benchmarks/osb/extensions/registry.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-from .param_sources import register as param_sources_register
-from .runners import register as runners_register
-
-
-def register(registry):
- param_sources_register(registry)
- runners_register(registry)
diff --git a/benchmarks/osb/extensions/runners.py b/benchmarks/osb/extensions/runners.py
deleted file mode 100644
index d048f80b0..000000000
--- a/benchmarks/osb/extensions/runners.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-from opensearchpy.exceptions import ConnectionTimeout
-from .util import parse_int_parameter, parse_string_parameter
-import logging
-import time
-
-
-def register(registry):
- registry.register_runner(
- "custom-vector-bulk", BulkVectorsFromDataSetRunner(), async_runner=True
- )
- registry.register_runner(
- "custom-refresh", CustomRefreshRunner(), async_runner=True
- )
- registry.register_runner(
- "train-model", TrainModelRunner(), async_runner=True
- )
- registry.register_runner(
- "delete-model", DeleteModelRunner(), async_runner=True
- )
-
-
-class BulkVectorsFromDataSetRunner:
-
- async def __call__(self, opensearch, params):
- size = parse_int_parameter("size", params)
- retries = parse_int_parameter("retries", params, 0) + 1
-
- for _ in range(retries):
- try:
- await opensearch.bulk(
- body=params["body"],
- timeout='5m'
- )
-
- return size, "docs"
- except ConnectionTimeout:
- logging.getLogger(__name__)\
- .warning("Bulk vector ingestion timed out. Retrying")
-
- raise TimeoutError("Failed to submit bulk request in specified number "
- "of retries: {}".format(retries))
-
- def __repr__(self, *args, **kwargs):
- return "custom-vector-bulk"
-
-
-class CustomRefreshRunner:
-
- async def __call__(self, opensearch, params):
- retries = parse_int_parameter("retries", params, 0) + 1
-
- for _ in range(retries):
- try:
- await opensearch.indices.refresh(
- index=parse_string_parameter("index", params)
- )
-
- return
- except ConnectionTimeout:
- logging.getLogger(__name__)\
- .warning("Custom refresh timed out. Retrying")
-
- raise TimeoutError("Failed to refresh the index in specified number "
- "of retries: {}".format(retries))
-
- def __repr__(self, *args, **kwargs):
- return "custom-refresh"
-
-
-class TrainModelRunner:
-
- async def __call__(self, opensearch, params):
- # Train a model and wait for it training to complete
- body = params["body"]
- timeout = parse_int_parameter("timeout", params)
- model_id = parse_string_parameter("model_id", params)
-
- method = "POST"
- model_uri = "/_plugins/_knn/models/{}".format(model_id)
- await opensearch.transport.perform_request(method, "{}/_train".format(model_uri), body=body)
-
- start_time = time.time()
- while time.time() < start_time + timeout:
- time.sleep(1)
- model_response = await opensearch.transport.perform_request("GET", model_uri)
-
- if 'state' not in model_response.keys():
- continue
-
- if model_response['state'] == 'created':
- #TODO: Return model size as well
- return 1, "models_trained"
-
- if model_response['state'] == 'failed':
- raise Exception("Failed to create model: {}".format(model_response))
-
- raise Exception('Failed to create model: {} within timeout {} seconds'
- .format(model_id, timeout))
-
- def __repr__(self, *args, **kwargs):
- return "train-model"
-
-
-class DeleteModelRunner:
-
- async def __call__(self, opensearch, params):
- # Delete model provided by model id
- method = "DELETE"
- model_id = parse_string_parameter("model_id", params)
- uri = "/_plugins/_knn/models/{}".format(model_id)
-
- # Ignore if model doesnt exist
- await opensearch.transport.perform_request(method, uri, params={"ignore": [400, 404]})
-
- def __repr__(self, *args, **kwargs):
- return "delete-model"
diff --git a/benchmarks/osb/extensions/util.py b/benchmarks/osb/extensions/util.py
deleted file mode 100644
index f7f6aab62..000000000
--- a/benchmarks/osb/extensions/util.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-import numpy as np
-from typing import List
-from typing import Dict
-from typing import Any
-
-
-def bulk_transform(partition: np.ndarray, field_name: str, action,
- offset: int) -> List[Dict[str, Any]]:
- """Partitions and transforms a list of vectors into OpenSearch's bulk
- injection format.
- Args:
- offset: to start counting from
- partition: An array of vectors to transform.
- field_name: field name for action
- action: Bulk API action.
- Returns:
- An array of transformed vectors in bulk format.
- """
- actions = []
- _ = [
- actions.extend([action(i + offset), None])
- for i in range(len(partition))
- ]
- actions[1::2] = [{field_name: vec} for vec in partition.tolist()]
- return actions
-
-
-def parse_string_parameter(key: str, params: dict, default: str = None) -> str:
- if key not in params:
- if default is not None:
- return default
- raise ConfigurationError(
- "Value cannot be None for param {}".format(key)
- )
-
- if type(params[key]) is str:
- return params[key]
-
- raise ConfigurationError("Value must be a string for param {}".format(key))
-
-
-def parse_int_parameter(key: str, params: dict, default: int = None) -> int:
- if key not in params:
- if default:
- return default
- raise ConfigurationError(
- "Value cannot be None for param {}".format(key)
- )
-
- if type(params[key]) is int:
- return params[key]
-
- raise ConfigurationError("Value must be a int for param {}".format(key))
-
-
-class ConfigurationError(Exception):
- """Exception raised for errors configuration.
-
- Attributes:
- message -- explanation of the error
- """
-
- def __init__(self, message: str):
- self.message = f'{message}'
- super().__init__(self.message)
diff --git a/benchmarks/osb/indices/faiss-index.json b/benchmarks/osb/indices/faiss-index.json
deleted file mode 100644
index 2db4d34d4..000000000
--- a/benchmarks/osb/indices/faiss-index.json
+++ /dev/null
@@ -1,27 +0,0 @@
-{
- "settings": {
- "index": {
- "knn": true,
- "number_of_shards": {{ target_index_primary_shards }},
- "number_of_replicas": {{ target_index_replica_shards }}
- }
- },
- "mappings": {
- "properties": {
- "target_field": {
- "type": "knn_vector",
- "dimension": {{ target_index_dimension }},
- "method": {
- "name": "hnsw",
- "space_type": "{{ target_index_space_type }}",
- "engine": "faiss",
- "parameters": {
- "ef_search": {{ hnsw_ef_search }},
- "ef_construction": {{ hnsw_ef_construction }},
- "m": {{ hnsw_m }}
- }
- }
- }
- }
- }
-}
diff --git a/benchmarks/osb/indices/model-index.json b/benchmarks/osb/indices/model-index.json
deleted file mode 100644
index 0e92c8903..000000000
--- a/benchmarks/osb/indices/model-index.json
+++ /dev/null
@@ -1,17 +0,0 @@
-{
- "settings": {
- "index": {
- "knn": true,
- "number_of_shards": {{ target_index_primary_shards | default(1) }},
- "number_of_replicas": {{ target_index_replica_shards | default(0) }}
- }
- },
- "mappings": {
- "properties": {
- "{{ target_field_name }}": {
- "type": "knn_vector",
- "model_id": "{{ train_model_id }}"
- }
- }
- }
-}
diff --git a/benchmarks/osb/indices/nmslib-index.json b/benchmarks/osb/indices/nmslib-index.json
deleted file mode 100644
index 4ceb57977..000000000
--- a/benchmarks/osb/indices/nmslib-index.json
+++ /dev/null
@@ -1,27 +0,0 @@
-{
- "settings": {
- "index": {
- "knn": true,
- "knn.algo_param.ef_search": {{ hnsw_ef_search }},
- "number_of_shards": {{ target_index_primary_shards }},
- "number_of_replicas": {{ target_index_replica_shards }}
- }
- },
- "mappings": {
- "properties": {
- "target_field": {
- "type": "knn_vector",
- "dimension": {{ target_index_dimension }},
- "method": {
- "name": "hnsw",
- "space_type": "{{ target_index_space_type }}",
- "engine": "nmslib",
- "parameters": {
- "ef_construction": {{ hnsw_ef_construction }},
- "m": {{ hnsw_m }}
- }
- }
- }
- }
- }
-}
diff --git a/benchmarks/osb/indices/train-index.json b/benchmarks/osb/indices/train-index.json
deleted file mode 100644
index 82af8215e..000000000
--- a/benchmarks/osb/indices/train-index.json
+++ /dev/null
@@ -1,16 +0,0 @@
-{
- "settings": {
- "index": {
- "number_of_shards": {{ train_index_primary_shards }},
- "number_of_replicas": {{ train_index_replica_shards }}
- }
- },
- "mappings": {
- "properties": {
- "{{ train_field_name }}": {
- "type": "knn_vector",
- "dimension": {{ target_index_dimension }}
- }
- }
- }
-}
diff --git a/benchmarks/osb/operations/default.json b/benchmarks/osb/operations/default.json
deleted file mode 100644
index ee33166f0..000000000
--- a/benchmarks/osb/operations/default.json
+++ /dev/null
@@ -1,53 +0,0 @@
-[
- {
- "name": "ivfpq-train-model",
- "operation-type": "train-model",
- "model_id": "{{ train_model_id }}",
- "timeout": {{ train_timeout }},
- "body": {
- "training_index": "{{ train_index_name }}",
- "training_field": "{{ train_field_name }}",
- "dimension": {{ target_index_dimension }},
- "search_size": {{ train_search_size }},
- "max_training_vector_count": {{ train_index_num_vectors }},
- "method": {
- "name":"ivf",
- "engine":"faiss",
- "space_type": "{{ target_index_space_type }}",
- "parameters":{
- "nlist": {{ ivf_nlists }},
- "nprobes": {{ ivf_nprobes }},
- "encoder":{
- "name":"pq",
- "parameters":{
- "code_size": {{ pq_code_size }},
- "m": {{ pq_m }}
- }
- }
- }
- }
- }
- },
- {
- "name": "ivf-train-model",
- "operation-type": "train-model",
- "model_id": "{{ train_model_id }}",
- "timeout": {{ train_timeout | default(1000) }},
- "body": {
- "training_index": "{{ train_index_name }}",
- "training_field": "{{ train_field_name }}",
- "search_size": {{ train_search_size }},
- "dimension": {{ target_index_dimension }},
- "max_training_vector_count": {{ train_index_num_vectors }},
- "method": {
- "name":"ivf",
- "engine":"faiss",
- "space_type": "{{ target_index_space_type }}",
- "parameters":{
- "nlist": {{ ivf_nlists }},
- "nprobes": {{ ivf_nprobes }}
- }
- }
- }
- }
-]
diff --git a/benchmarks/osb/params/no-train-params.json b/benchmarks/osb/params/no-train-params.json
deleted file mode 100644
index 64fe3c296..000000000
--- a/benchmarks/osb/params/no-train-params.json
+++ /dev/null
@@ -1,38 +0,0 @@
-{
- "target_index_name": "target_index",
- "target_field_name": "target_field",
- "target_index_body": "indices/nmslib-index.json",
- "target_index_primary_shards": 3,
- "target_index_replica_shards": 1,
- "target_index_dimension": 128,
- "target_index_space_type": "l2",
- "target_index_bulk_size": 200,
- "target_index_bulk_index_data_set_format": "hdf5",
- "target_index_bulk_index_data_set_path": "",
- "target_index_bulk_index_clients": 10,
- "hnsw_ef_search": 512,
- "hnsw_ef_construction": 512,
- "hnsw_m": 16,
-
- "query_k": 10,
- "query_clients": 10,
- "query_data_set_format": "hdf5",
- "query_data_set_path": "",
-
- "ivf_nlists": 1,
- "ivf_nprobes": 1,
- "pq_code_size": 1,
- "pq_m": 1,
- "train_model_method": "",
- "train_model_id": "",
- "train_index_name": "",
- "train_field_name": "",
- "train_index_body": "",
- "train_search_size": 1,
- "train_timeout": 1,
- "train_index_bulk_size": 1,
- "train_index_data_set_format": "",
- "train_index_data_set_path": "",
- "train_index_num_vectors": 1,
- "train_index_bulk_index_clients": 1
-}
diff --git a/benchmarks/osb/params/train-params.json b/benchmarks/osb/params/train-params.json
deleted file mode 100644
index 4c598d25b..000000000
--- a/benchmarks/osb/params/train-params.json
+++ /dev/null
@@ -1,36 +0,0 @@
-{
- "target_index_name": "target_index",
- "target_field_name": "target_field",
- "target_index_body": "indices/model-index.json",
- "target_index_primary_shards": 3,
- "target_index_replica_shards": 1,
- "target_index_dimension": 128,
- "target_index_space_type": "l2",
- "target_index_bulk_size": 200,
- "target_index_bulk_index_data_set_format": "hdf5",
- "target_index_bulk_index_data_set_path": "",
- "target_index_bulk_index_clients": 10,
- "ivf_nlists": 10,
- "ivf_nprobes": 1,
- "pq_code_size": 8,
- "pq_m": 8,
- "train_model_method": "ivfpq",
- "train_model_id": "test-model",
- "train_index_name": "train_index",
- "train_field_name": "train_field",
- "train_index_body": "indices/train-index.json",
- "train_search_size": 500,
- "train_timeout": 5000,
- "train_index_primary_shards": 1,
- "train_index_replica_shards": 0,
- "train_index_bulk_size": 200,
- "train_index_data_set_format": "hdf5",
- "train_index_data_set_path": "",
- "train_index_num_vectors": 1000000,
- "train_index_bulk_index_clients": 10,
-
- "query_k": 10,
- "query_clients": 10,
- "query_data_set_format": "hdf5",
- "query_data_set_path": ""
-}
diff --git a/benchmarks/osb/procedures/no-train-test.json b/benchmarks/osb/procedures/no-train-test.json
deleted file mode 100644
index 03d72d6bd..000000000
--- a/benchmarks/osb/procedures/no-train-test.json
+++ /dev/null
@@ -1,63 +0,0 @@
-{% import "benchmark.helpers" as benchmark with context %}
-{
- "name": "no-train-test",
- "default": true,
- "schedule": [
- {
- "operation": {
- "name": "delete-target-index",
- "operation-type": "delete-index",
- "only-if-exists": true,
- "index": "{{ target_index_name }}"
- }
- },
- {
- "operation": {
- "name": "create-target-index",
- "operation-type": "create-index",
- "index": "{{ target_index_name }}"
- }
- },
- {
- "name": "wait-for-cluster-to-be-green",
- "operation": "cluster-health",
- "request-params": {
- "wait_for_status": "green"
- }
- },
- {
- "operation": {
- "name": "custom-vector-bulk",
- "operation-type": "custom-vector-bulk",
- "param-source": "bulk-from-data-set",
- "index": "{{ target_index_name }}",
- "field": "{{ target_field_name }}",
- "bulk_size": {{ target_index_bulk_size }},
- "data_set_format": "{{ target_index_bulk_index_data_set_format }}",
- "data_set_path": "{{ target_index_bulk_index_data_set_path }}"
- },
- "clients": {{ target_index_bulk_index_clients }}
- },
- {
- "operation": {
- "name": "refresh-target-index",
- "operation-type": "custom-refresh",
- "index": "{{ target_index_name }}",
- "retries": 100
- }
- },
- {
- "operation": {
- "name": "knn-query-from-data-set",
- "operation-type": "search",
- "index": "{{ target_index_name }}",
- "param-source": "knn-query-from-data-set",
- "k": {{ query_k }},
- "field": "{{ target_field_name }}",
- "data_set_format": "{{ query_data_set_format }}",
- "data_set_path": "{{ query_data_set_path }}"
- },
- "clients": {{ query_clients }}
- }
- ]
-}
diff --git a/benchmarks/osb/procedures/train-test.json b/benchmarks/osb/procedures/train-test.json
deleted file mode 100644
index 49930044a..000000000
--- a/benchmarks/osb/procedures/train-test.json
+++ /dev/null
@@ -1,117 +0,0 @@
-{% import "benchmark.helpers" as benchmark with context %}
-{
- "name": "train-test",
- "default": false,
- "schedule": [
- {
- "operation": {
- "name": "delete-target-index",
- "operation-type": "delete-index",
- "only-if-exists": true,
- "index": "{{ target_index_name }}"
- }
- },
- {
- "operation": {
- "name": "delete-train-index",
- "operation-type": "delete-index",
- "only-if-exists": true,
- "index": "{{ train_index_name }}"
- }
- },
- {
- "operation": {
- "operation-type": "delete-model",
- "name": "delete-model",
- "model_id": "{{ train_model_id }}"
- }
- },
- {
- "operation": {
- "name": "create-train-index",
- "operation-type": "create-index",
- "index": "{{ train_index_name }}"
- }
- },
- {
- "name": "wait-for-train-index-to-be-green",
- "operation": "cluster-health",
- "request-params": {
- "wait_for_status": "green"
- }
- },
- {
- "operation": {
- "name": "train-vector-bulk",
- "operation-type": "custom-vector-bulk",
- "param-source": "bulk-from-data-set",
- "index": "{{ train_index_name }}",
- "field": "{{ train_field_name }}",
- "bulk_size": {{ train_index_bulk_size }},
- "data_set_format": "{{ train_index_data_set_format }}",
- "data_set_path": "{{ train_index_data_set_path }}",
- "num_vectors": {{ train_index_num_vectors }}
- },
- "clients": {{ train_index_bulk_index_clients }}
- },
- {
- "operation": {
- "name": "refresh-train-index",
- "operation-type": "custom-refresh",
- "index": "{{ train_index_name }}",
- "retries": 100
- }
- },
- {
- "operation": "{{ train_model_method }}-train-model"
- },
- {
- "operation": {
- "name": "create-target-index",
- "operation-type": "create-index",
- "index": "{{ target_index_name }}"
- }
- },
- {
- "name": "wait-for-target-index-to-be-green",
- "operation": "cluster-health",
- "request-params": {
- "wait_for_status": "green"
- }
- },
- {
- "operation": {
- "name": "custom-vector-bulk",
- "operation-type": "custom-vector-bulk",
- "param-source": "bulk-from-data-set",
- "index": "{{ target_index_name }}",
- "field": "{{ target_field_name }}",
- "bulk_size": {{ target_index_bulk_size }},
- "data_set_format": "{{ target_index_bulk_index_data_set_format }}",
- "data_set_path": "{{ target_index_bulk_index_data_set_path }}"
- },
- "clients": {{ target_index_bulk_index_clients }}
- },
- {
- "operation": {
- "name": "refresh-target-index",
- "operation-type": "custom-refresh",
- "index": "{{ target_index_name }}",
- "retries": 100
- }
- },
- {
- "operation": {
- "name": "knn-query-from-data-set",
- "operation-type": "search",
- "index": "{{ target_index_name }}",
- "param-source": "knn-query-from-data-set",
- "k": {{ query_k }},
- "field": "{{ target_field_name }}",
- "data_set_format": "{{ query_data_set_format }}",
- "data_set_path": "{{ query_data_set_path }}"
- },
- "clients": {{ query_clients }}
- }
- ]
-}
diff --git a/benchmarks/osb/requirements.in b/benchmarks/osb/requirements.in
deleted file mode 100644
index a9e12b5d3..000000000
--- a/benchmarks/osb/requirements.in
+++ /dev/null
@@ -1,4 +0,0 @@
-opensearch-py
-numpy
-h5py
-opensearch-benchmark
diff --git a/benchmarks/osb/requirements.txt b/benchmarks/osb/requirements.txt
deleted file mode 100644
index 271e8ab07..000000000
--- a/benchmarks/osb/requirements.txt
+++ /dev/null
@@ -1,98 +0,0 @@
-#
-# This file is autogenerated by pip-compile with python 3.8
-# To update, run:
-#
-# pip-compile
-#
-aiohttp==3.8.1
- # via opensearch-py
-aiosignal==1.2.0
- # via aiohttp
-async-timeout==4.0.2
- # via aiohttp
-attrs==21.4.0
- # via
- # aiohttp
- # jsonschema
-cachetools==4.2.4
- # via google-auth
-certifi==2021.10.8
- # via
- # opensearch-benchmark
- # opensearch-py
-charset-normalizer==2.0.12
- # via aiohttp
-frozenlist==1.3.0
- # via
- # aiohttp
- # aiosignal
-google-auth==1.22.1
- # via opensearch-benchmark
-google-crc32c==1.3.0
- # via google-resumable-media
-google-resumable-media==1.1.0
- # via opensearch-benchmark
-h5py==3.6.0
- # via -r requirements.in
-idna==3.3
- # via yarl
-ijson==2.6.1
- # via opensearch-benchmark
-importlib-metadata==4.11.3
- # via jsonschema
-jinja2==2.11.3
- # via opensearch-benchmark
-jsonschema==3.1.1
- # via opensearch-benchmark
-markupsafe==2.0.1
- # via
- # jinja2
- # opensearch-benchmark
-multidict==6.0.2
- # via
- # aiohttp
- # yarl
-numpy==1.22.3
- # via
- # -r requirements.in
- # h5py
-opensearch-benchmark==0.0.2
- # via -r requirements.in
-opensearch-py[async]==1.0.0
- # via
- # -r requirements.in
- # opensearch-benchmark
-psutil==5.8.0
- # via opensearch-benchmark
-py-cpuinfo==7.0.0
- # via opensearch-benchmark
-pyasn1==0.4.8
- # via
- # pyasn1-modules
- # rsa
-pyasn1-modules==0.2.8
- # via google-auth
-pyrsistent==0.18.1
- # via jsonschema
-rsa==4.8
- # via google-auth
-six==1.16.0
- # via
- # google-auth
- # google-resumable-media
- # jsonschema
-tabulate==0.8.7
- # via opensearch-benchmark
-thespian==3.10.1
- # via opensearch-benchmark
-urllib3==1.26.9
- # via opensearch-py
-yappi==1.2.3
- # via opensearch-benchmark
-yarl==1.7.2
- # via aiohttp
-zipp==3.7.0
- # via importlib-metadata
-
-# The following packages are considered to be unsafe in a requirements file:
-# setuptools
\ No newline at end of file
diff --git a/benchmarks/osb/tests/__init__.py b/benchmarks/osb/tests/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/benchmarks/osb/tests/data_set_helper.py b/benchmarks/osb/tests/data_set_helper.py
deleted file mode 100644
index 2b144da49..000000000
--- a/benchmarks/osb/tests/data_set_helper.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-from abc import ABC, abstractmethod
-
-import h5py
-import numpy as np
-
-from osb.extensions.data_set import Context, HDF5DataSet, BigANNVectorDataSet
-
-""" Module containing utility classes and functions for working with data sets.
-
-Included are utilities that can be used to build data sets and write them to
-paths.
-"""
-
-
-class DataSetBuildContext:
- """ Data class capturing information needed to build a particular data set
-
- Attributes:
- data_set_context: Indicator of what the data set is used for,
- vectors: A 2D array containing vectors that are used to build data set.
- path: string representing path where data set should be serialized to.
- """
- def __init__(self, data_set_context: Context, vectors: np.ndarray, path: str):
- self.data_set_context: Context = data_set_context
- self.vectors: np.ndarray = vectors #TODO: Validate shape
- self.path: str = path
-
- def get_num_vectors(self) -> int:
- return self.vectors.shape[0]
-
- def get_dimension(self) -> int:
- return self.vectors.shape[1]
-
- def get_type(self) -> np.dtype:
- return self.vectors.dtype
-
-
-class DataSetBuilder(ABC):
- """ Abstract builder used to create a build a collection of data sets
-
- Attributes:
- data_set_build_contexts: list of data set build contexts that builder
- will build.
- """
- def __init__(self):
- self.data_set_build_contexts = list()
-
- def add_data_set_build_context(self, data_set_build_context: DataSetBuildContext):
- """ Adds a data set build context to list of contexts to be built.
-
- Args:
- data_set_build_context: DataSetBuildContext to be added to list
-
- Returns: Updated DataSetBuilder
-
- """
- self._validate_data_set_context(data_set_build_context)
- self.data_set_build_contexts.append(data_set_build_context)
- return self
-
- def build(self):
- """ Builds and serializes all data sets build contexts
-
- Returns:
-
- """
- [self._build_data_set(data_set_build_context) for data_set_build_context
- in self.data_set_build_contexts]
-
- @abstractmethod
- def _build_data_set(self, context: DataSetBuildContext):
- """ Builds an individual data set
-
- Args:
- context: DataSetBuildContext of data set to be built
-
- Returns:
-
- """
- pass
-
- @abstractmethod
- def _validate_data_set_context(self, context: DataSetBuildContext):
- """ Validates that data set context can be added to this builder
-
- Args:
- context: DataSetBuildContext to be validated
-
- Returns:
-
- """
- pass
-
-
-class HDF5Builder(DataSetBuilder):
-
- def __init__(self):
- super(HDF5Builder, self).__init__()
- self.data_set_meta_data = dict()
-
- def _validate_data_set_context(self, context: DataSetBuildContext):
- if context.path not in self.data_set_meta_data.keys():
- self.data_set_meta_data[context.path] = {
- context.data_set_context: context
- }
- return
-
- if context.data_set_context in \
- self.data_set_meta_data[context.path].keys():
- raise IllegalDataSetBuildContext("Path and context for data set "
- "are already present in builder.")
-
- self.data_set_meta_data[context.path][context.data_set_context] = \
- context
-
- @staticmethod
- def _validate_extension(context: DataSetBuildContext):
- ext = context.path.split('.')[-1]
-
- if ext != HDF5DataSet.FORMAT_NAME:
- raise IllegalDataSetBuildContext("Invalid file extension")
-
- def _build_data_set(self, context: DataSetBuildContext):
- # For HDF5, because multiple data sets can be grouped in the same file,
- # we will build data sets in memory and not write to disk until
- # _flush_data_sets_to_disk is called
- with h5py.File(context.path, 'a') as hf:
- hf.create_dataset(
- HDF5DataSet.parse_context(context.data_set_context),
- data=context.vectors
- )
-
-
-class BigANNBuilder(DataSetBuilder):
-
- def _validate_data_set_context(self, context: DataSetBuildContext):
- self._validate_extension(context)
-
- # prevent the duplication of paths for data sets
- data_set_paths = [c.path for c in self.data_set_build_contexts]
- if any(data_set_paths.count(x) > 1 for x in data_set_paths):
- raise IllegalDataSetBuildContext("Build context paths have to be "
- "unique.")
-
- @staticmethod
- def _validate_extension(context: DataSetBuildContext):
- ext = context.path.split('.')[-1]
-
- if ext != BigANNVectorDataSet.U8BIN_EXTENSION and ext != \
- BigANNVectorDataSet.FBIN_EXTENSION:
- raise IllegalDataSetBuildContext("Invalid file extension")
-
- if ext == BigANNVectorDataSet.U8BIN_EXTENSION and context.get_type() != \
- np.u8int:
- raise IllegalDataSetBuildContext("Invalid data type for {} ext."
- .format(BigANNVectorDataSet
- .U8BIN_EXTENSION))
-
- if ext == BigANNVectorDataSet.FBIN_EXTENSION and context.get_type() != \
- np.float32:
- print(context.get_type())
- raise IllegalDataSetBuildContext("Invalid data type for {} ext."
- .format(BigANNVectorDataSet
- .FBIN_EXTENSION))
-
- def _build_data_set(self, context: DataSetBuildContext):
- num_vectors = context.get_num_vectors()
- dimension = context.get_dimension()
-
- with open(context.path, 'wb') as f:
- f.write(int.to_bytes(num_vectors, 4, "little"))
- f.write(int.to_bytes(dimension, 4, "little"))
- context.vectors.tofile(f)
-
-
-def create_random_2d_array(num_vectors: int, dimension: int) -> np.ndarray:
- rng = np.random.default_rng()
- return rng.random(size=(num_vectors, dimension), dtype=np.float32)
-
-
-class IllegalDataSetBuildContext(Exception):
- """Exception raised when passed in DataSetBuildContext is illegal
-
- Attributes:
- message -- explanation of the error
- """
-
- def __init__(self, message: str):
- self.message = f'{message}'
- super().__init__(self.message)
-
diff --git a/benchmarks/osb/tests/test_param_sources.py b/benchmarks/osb/tests/test_param_sources.py
deleted file mode 100644
index cda730cee..000000000
--- a/benchmarks/osb/tests/test_param_sources.py
+++ /dev/null
@@ -1,353 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-import os
-import random
-import shutil
-import string
-import sys
-import tempfile
-import unittest
-
-# Add parent directory to path
-import numpy as np
-
-sys.path.append(os.path.abspath(os.path.join(os.getcwd(), os.pardir)))
-
-from osb.tests.data_set_helper import HDF5Builder, create_random_2d_array, \
- DataSetBuildContext, BigANNBuilder
-from osb.extensions.data_set import Context, HDF5DataSet
-from osb.extensions.param_sources import VectorsFromDataSetParamSource, \
- QueryVectorsFromDataSetParamSource, BulkVectorsFromDataSetParamSource
-from osb.extensions.util import ConfigurationError
-
-DEFAULT_INDEX_NAME = "test-index"
-DEFAULT_FIELD_NAME = "test-field"
-DEFAULT_CONTEXT = Context.INDEX
-DEFAULT_TYPE = HDF5DataSet.FORMAT_NAME
-DEFAULT_NUM_VECTORS = 10
-DEFAULT_DIMENSION = 10
-DEFAULT_RANDOM_STRING_LENGTH = 8
-
-
-class VectorsFromDataSetParamSourceTestCase(unittest.TestCase):
-
- def setUp(self) -> None:
- self.data_set_dir = tempfile.mkdtemp()
-
- # Create a data set we know to be valid for convenience
- self.valid_data_set_path = _create_data_set(
- DEFAULT_NUM_VECTORS,
- DEFAULT_DIMENSION,
- DEFAULT_TYPE,
- DEFAULT_CONTEXT,
- self.data_set_dir
- )
-
- def tearDown(self):
- shutil.rmtree(self.data_set_dir)
-
- def test_missing_params(self):
- empty_params = dict()
- self.assertRaises(
- ConfigurationError,
- lambda: VectorsFromDataSetParamSourceTestCase.
- TestVectorsFromDataSetParamSource(empty_params, DEFAULT_CONTEXT)
- )
-
- def test_invalid_data_set_format(self):
- invalid_data_set_format = "invalid-data-set-format"
-
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": invalid_data_set_format,
- "data_set_path": self.valid_data_set_path,
- }
- self.assertRaises(
- ConfigurationError,
- lambda: self.TestVectorsFromDataSetParamSource(
- test_param_source_params,
- DEFAULT_CONTEXT
- )
- )
-
- def test_invalid_data_set_path(self):
- invalid_data_set_path = "invalid-data-set-path"
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": HDF5DataSet.FORMAT_NAME,
- "data_set_path": invalid_data_set_path,
- }
- self.assertRaises(
- FileNotFoundError,
- lambda: self.TestVectorsFromDataSetParamSource(
- test_param_source_params,
- DEFAULT_CONTEXT
- )
- )
-
- def test_partition_hdf5(self):
- num_vectors = 100
-
- hdf5_data_set_path = _create_data_set(
- num_vectors,
- DEFAULT_DIMENSION,
- HDF5DataSet.FORMAT_NAME,
- DEFAULT_CONTEXT,
- self.data_set_dir
- )
-
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": HDF5DataSet.FORMAT_NAME,
- "data_set_path": hdf5_data_set_path,
- }
- test_param_source = self.TestVectorsFromDataSetParamSource(
- test_param_source_params,
- DEFAULT_CONTEXT
- )
-
- num_partitions = 10
- vecs_per_partition = test_param_source.num_vectors // num_partitions
-
- self._test_partition(
- test_param_source,
- num_partitions,
- vecs_per_partition
- )
-
- def test_partition_bigann(self):
- num_vectors = 100
- float_extension = "fbin"
-
- bigann_data_set_path = _create_data_set(
- num_vectors,
- DEFAULT_DIMENSION,
- float_extension,
- DEFAULT_CONTEXT,
- self.data_set_dir
- )
-
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": "bigann",
- "data_set_path": bigann_data_set_path,
- }
- test_param_source = self.TestVectorsFromDataSetParamSource(
- test_param_source_params,
- DEFAULT_CONTEXT
- )
-
- num_partitions = 10
- vecs_per_partition = test_param_source.num_vectors // num_partitions
-
- self._test_partition(
- test_param_source,
- num_partitions,
- vecs_per_partition
- )
-
- def _test_partition(
- self,
- test_param_source: VectorsFromDataSetParamSource,
- num_partitions: int,
- vec_per_partition: int
- ):
- for i in range(num_partitions):
- test_param_source_i = test_param_source.partition(i, num_partitions)
- self.assertEqual(test_param_source_i.num_vectors, vec_per_partition)
- self.assertEqual(test_param_source_i.offset, i * vec_per_partition)
-
- class TestVectorsFromDataSetParamSource(VectorsFromDataSetParamSource):
- """
- Empty implementation of ABC VectorsFromDataSetParamSource so that we can
- test the concrete methods.
- """
-
- def params(self):
- pass
-
-
-class QueryVectorsFromDataSetParamSourceTestCase(unittest.TestCase):
-
- def setUp(self) -> None:
- self.data_set_dir = tempfile.mkdtemp()
-
- def tearDown(self):
- shutil.rmtree(self.data_set_dir)
-
- def test_params(self):
- # Create a data set
- k = 12
- data_set_path = _create_data_set(
- DEFAULT_NUM_VECTORS,
- DEFAULT_DIMENSION,
- DEFAULT_TYPE,
- Context.QUERY,
- self.data_set_dir
- )
-
- # Create a QueryVectorsFromDataSetParamSource with relevant params
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": DEFAULT_TYPE,
- "data_set_path": data_set_path,
- "k": k,
- }
- query_param_source = QueryVectorsFromDataSetParamSource(
- None, test_param_source_params
- )
-
- # Check each
- for i in range(DEFAULT_NUM_VECTORS):
- self._check_params(
- query_param_source.params(),
- DEFAULT_INDEX_NAME,
- DEFAULT_FIELD_NAME,
- DEFAULT_DIMENSION,
- k
- )
-
- # Assert last call creates stop iteration
- self.assertRaises(
- StopIteration,
- lambda: query_param_source.params()
- )
-
- def _check_params(
- self,
- params: dict,
- expected_index: str,
- expected_field: str,
- expected_dimension: int,
- expected_k: int
- ):
- index_name = params.get("index")
- self.assertEqual(expected_index, index_name)
- body = params.get("body")
- self.assertIsInstance(body, dict)
- query = body.get("query")
- self.assertIsInstance(query, dict)
- query_knn = query.get("knn")
- self.assertIsInstance(query_knn, dict)
- field = query_knn.get(expected_field)
- self.assertIsInstance(field, dict)
- vector = field.get("vector")
- self.assertIsInstance(vector, np.ndarray)
- self.assertEqual(len(list(vector)), expected_dimension)
- k = field.get("k")
- self.assertEqual(k, expected_k)
-
-
-class BulkVectorsFromDataSetParamSourceTestCase(unittest.TestCase):
-
- def setUp(self) -> None:
- self.data_set_dir = tempfile.mkdtemp()
-
- def tearDown(self):
- shutil.rmtree(self.data_set_dir)
-
- def test_params(self):
- num_vectors = 49
- bulk_size = 10
- data_set_path = _create_data_set(
- num_vectors,
- DEFAULT_DIMENSION,
- DEFAULT_TYPE,
- Context.INDEX,
- self.data_set_dir
- )
-
- test_param_source_params = {
- "index": DEFAULT_INDEX_NAME,
- "field": DEFAULT_FIELD_NAME,
- "data_set_format": DEFAULT_TYPE,
- "data_set_path": data_set_path,
- "bulk_size": bulk_size
- }
- bulk_param_source = BulkVectorsFromDataSetParamSource(
- None, test_param_source_params
- )
-
- # Check each payload returned
- vectors_consumed = 0
- while vectors_consumed < num_vectors:
- expected_num_vectors = min(num_vectors - vectors_consumed, bulk_size)
- self._check_params(
- bulk_param_source.params(),
- DEFAULT_INDEX_NAME,
- DEFAULT_FIELD_NAME,
- DEFAULT_DIMENSION,
- expected_num_vectors
- )
- vectors_consumed += expected_num_vectors
-
- # Assert last call creates stop iteration
- self.assertRaises(
- StopIteration,
- lambda: bulk_param_source.params()
- )
-
- def _check_params(
- self,
- params: dict,
- expected_index: str,
- expected_field: str,
- expected_dimension: int,
- expected_num_vectors_in_payload: int
- ):
- size = params.get("size")
- self.assertEqual(size, expected_num_vectors_in_payload)
- body = params.get("body")
- self.assertIsInstance(body, list)
- self.assertEqual(len(body) // 2, expected_num_vectors_in_payload)
-
- # Bulk payload has 2 parts: first one is the header and the second one
- # is the body. The header will have the index name and the body will
- # have the vector
- for header, req_body in zip(*[iter(body)] * 2):
- index = header.get("index")
- self.assertIsInstance(index, dict)
- index_name = index.get("_index")
- self.assertEqual(index_name, expected_index)
-
- vector = req_body.get(expected_field)
- self.assertIsInstance(vector, list)
- self.assertEqual(len(vector), expected_dimension)
-
-
-def _create_data_set(
- num_vectors: int,
- dimension: int,
- extension: str,
- data_set_context: Context,
- data_set_dir
-) -> str:
-
- file_name_base = ''.join(random.choice(string.ascii_letters) for _ in
- range(DEFAULT_RANDOM_STRING_LENGTH))
- data_set_file_name = "{}.{}".format(file_name_base, extension)
- data_set_path = os.path.join(data_set_dir, data_set_file_name)
- context = DataSetBuildContext(
- data_set_context,
- create_random_2d_array(num_vectors, dimension),
- data_set_path)
-
- if extension == HDF5DataSet.FORMAT_NAME:
- HDF5Builder().add_data_set_build_context(context).build()
- else:
- BigANNBuilder().add_data_set_build_context(context).build()
-
- return data_set_path
-
-
-if __name__ == '__main__':
- unittest.main()
diff --git a/benchmarks/osb/workload.json b/benchmarks/osb/workload.json
deleted file mode 100644
index bd0d84195..000000000
--- a/benchmarks/osb/workload.json
+++ /dev/null
@@ -1,17 +0,0 @@
-{% import "benchmark.helpers" as benchmark with context %}
-{
- "version": 2,
- "description": "k-NN Plugin train workload",
- "indices": [
- {
- "name": "{{ target_index_name }}",
- "body": "{{ target_index_body }}"
- },
- {
- "name": "{{ train_index_name }}",
- "body": "{{ train_index_body }}"
- }
- ],
- "operations": {{ benchmark.collect(parts="operations/*.json") }},
- "test_procedures": [{{ benchmark.collect(parts="procedures/*.json") }}]
-}
diff --git a/benchmarks/osb/workload.py b/benchmarks/osb/workload.py
deleted file mode 100644
index 32e6ad02c..000000000
--- a/benchmarks/osb/workload.py
+++ /dev/null
@@ -1,18 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-# This code needs to be included at the top of every workload.py file.
-# OpenSearch Benchmarks is not able to find other helper files unless the path
-# is updated.
-import os
-import sys
-sys.path.append(os.path.abspath(os.getcwd()))
-
-from extensions.registry import register as custom_register
-
-
-def register(registry):
- custom_register(registry)
diff --git a/benchmarks/perf-tool/.pylintrc b/benchmarks/perf-tool/.pylintrc
deleted file mode 100644
index 15bf4ccc3..000000000
--- a/benchmarks/perf-tool/.pylintrc
+++ /dev/null
@@ -1,443 +0,0 @@
-# This Pylint rcfile contains a best-effort configuration to uphold the
-# best-practices and style described in the Google Python style guide:
-# https://google.github.io/styleguide/pyguide.html
-#
-# Its canonical open-source location is:
-# https://google.github.io/styleguide/pylintrc
-
-[MASTER]
-
-fail-under=9.0
-
-# Files or directories to be skipped. They should be base names, not paths.
-ignore=third_party
-
-# Files or directories matching the regex patterns are skipped. The regex
-# matches against base names, not paths.
-ignore-patterns=
-
-# Pickle collected data for later comparisons.
-persistent=no
-
-# List of plugins (as comma separated values of python modules names) to load,
-# usually to register additional checkers.
-load-plugins=
-
-# Use multiple processes to speed up Pylint.
-jobs=4
-
-# Allow loading of arbitrary C extensions. Extensions are imported into the
-# active Python interpreter and may run arbitrary code.
-unsafe-load-any-extension=no
-
-
-[MESSAGES CONTROL]
-
-# Only show warnings with the listed confidence levels. Leave empty to show
-# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED
-confidence=
-
-# Enable the message, report, category or checker with the given id(s). You can
-# either give multiple identifier separated by comma (,) or put this option
-# multiple time (only on the command line, not in the configuration file where
-# it should appear only once). See also the "--disable" option for examples.
-#enable=
-
-# Disable the message, report, category or checker with the given id(s). You
-# can either give multiple identifiers separated by comma (,) or put this
-# option multiple times (only on the command line, not in the configuration
-# file where it should appear only once).You can also use "--disable=all" to
-# disable everything first and then reenable specific checks. For example, if
-# you want to run only the similarities checker, you can use "--disable=all
-# --enable=similarities". If you want to run only the classes checker, but have
-# no Warning level messages displayed, use"--disable=all --enable=classes
-# --disable=W"
-disable=abstract-method,
- apply-builtin,
- arguments-differ,
- attribute-defined-outside-init,
- backtick,
- bad-option-value,
- basestring-builtin,
- buffer-builtin,
- c-extension-no-member,
- consider-using-enumerate,
- cmp-builtin,
- cmp-method,
- coerce-builtin,
- coerce-method,
- delslice-method,
- div-method,
- duplicate-code,
- eq-without-hash,
- execfile-builtin,
- file-builtin,
- filter-builtin-not-iterating,
- fixme,
- getslice-method,
- global-statement,
- hex-method,
- idiv-method,
- implicit-str-concat-in-sequence,
- import-error,
- import-self,
- import-star-module-level,
- inconsistent-return-statements,
- input-builtin,
- intern-builtin,
- invalid-str-codec,
- locally-disabled,
- long-builtin,
- long-suffix,
- map-builtin-not-iterating,
- misplaced-comparison-constant,
- missing-function-docstring,
- metaclass-assignment,
- next-method-called,
- next-method-defined,
- no-absolute-import,
- no-else-break,
- no-else-continue,
- no-else-raise,
- no-else-return,
- no-init, # added
- no-member,
- no-name-in-module,
- no-self-use,
- nonzero-method,
- oct-method,
- old-division,
- old-ne-operator,
- old-octal-literal,
- old-raise-syntax,
- parameter-unpacking,
- print-statement,
- raising-string,
- range-builtin-not-iterating,
- raw_input-builtin,
- rdiv-method,
- reduce-builtin,
- relative-import,
- reload-builtin,
- round-builtin,
- setslice-method,
- signature-differs,
- standarderror-builtin,
- suppressed-message,
- sys-max-int,
- too-few-public-methods,
- too-many-ancestors,
- too-many-arguments,
- too-many-boolean-expressions,
- too-many-branches,
- too-many-instance-attributes,
- too-many-locals,
- too-many-nested-blocks,
- too-many-public-methods,
- too-many-return-statements,
- too-many-statements,
- trailing-newlines,
- unichr-builtin,
- unicode-builtin,
- unnecessary-pass,
- unpacking-in-except,
- useless-else-on-loop,
- useless-object-inheritance,
- useless-suppression,
- using-cmp-argument,
- wrong-import-order,
- xrange-builtin,
- zip-builtin-not-iterating,
-
-
-[REPORTS]
-
-# Set the output format. Available formats are text, parseable, colorized, msvs
-# (visual studio) and html. You can also give a reporter class, eg
-# mypackage.mymodule.MyReporterClass.
-output-format=text
-
-# Put messages in a separate file for each module / package specified on the
-# command line instead of printing them on stdout. Reports (if any) will be
-# written in a file name "pylint_global.[txt|html]". This option is deprecated
-# and it will be removed in Pylint 2.0.
-files-output=no
-
-# Tells whether to display a full report or only the messages
-reports=no
-
-# Python expression which should return a note less than 10 (10 is the highest
-# note). You have access to the variables errors warning, statement which
-# respectively contain the number of errors / warnings messages and the total
-# number of statements analyzed. This is used by the global evaluation report
-# (RP0004).
-evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
-
-# Template used to display messages. This is a python new-style format string
-# used to format the message information. See doc for all details
-#msg-template=
-
-
-[BASIC]
-
-# Good variable names which should always be accepted, separated by a comma
-good-names=main,_
-
-# Bad variable names which should always be refused, separated by a comma
-bad-names=
-
-# Colon-delimited sets of names that determine each other's naming style when
-# the name regexes allow several styles.
-name-group=
-
-# Include a hint for the correct naming format with invalid-name
-include-naming-hint=no
-
-# List of decorators that produce properties, such as abc.abstractproperty. Add
-# to this list to register other decorators that produce valid properties.
-property-classes=abc.abstractproperty,cached_property.cached_property,cached_property.threaded_cached_property,cached_property.cached_property_with_ttl,cached_property.threaded_cached_property_with_ttl
-
-# Regular expression matching correct function names
-function-rgx=^(?:(?PsetUp|tearDown|setUpModule|tearDownModule)|(?P_?[A-Z][a-zA-Z0-9]*)|(?P_?[a-z][a-z0-9_]*))$
-
-# Regular expression matching correct variable names
-variable-rgx=^[a-z][a-z0-9_]*$
-
-# Regular expression matching correct constant names
-const-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
-
-# Regular expression matching correct attribute names
-attr-rgx=^_{0,2}[a-z][a-z0-9_]*$
-
-# Regular expression matching correct argument names
-argument-rgx=^[a-z][a-z0-9_]*$
-
-# Regular expression matching correct class attribute names
-class-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
-
-# Regular expression matching correct inline iteration names
-inlinevar-rgx=^[a-z][a-z0-9_]*$
-
-# Regular expression matching correct class names
-class-rgx=^_?[A-Z][a-zA-Z0-9]*$
-
-# Regular expression matching correct module names
-module-rgx=^(_?[a-z][a-z0-9_]*|__init__)$
-
-# Regular expression matching correct method names
-method-rgx=(?x)^(?:(?P_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase|tearDownTestCase|setupSelf|tearDownClass|setUpClass|(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next)|(?P_{0,2}[A-Z][a-zA-Z0-9_]*)|(?P_{0,2}[a-z][a-z0-9_]*))$
-
-# Regular expression which should only match function or class names that do
-# not require a docstring.
-no-docstring-rgx=(__.*__|main|test.*|.*test|.*Test)$
-
-# Minimum line length for functions/classes that require docstrings, shorter
-# ones are exempt.
-docstring-min-length=10
-
-
-[TYPECHECK]
-
-# List of decorators that produce context managers, such as
-# contextlib.contextmanager. Add to this list to register other decorators that
-# produce valid context managers.
-contextmanager-decorators=contextlib.contextmanager,contextlib2.contextmanager
-
-# Tells whether missing members accessed in mixin class should be ignored. A
-# mixin class is detected if its name ends with "mixin" (case insensitive).
-ignore-mixin-members=yes
-
-# List of module names for which member attributes should not be checked
-# (useful for modules/projects where namespaces are manipulated during runtime
-# and thus existing member attributes cannot be deduced by static analysis. It
-# supports qualified module names, as well as Unix pattern matching.
-ignored-modules=
-
-# List of class names for which member attributes should not be checked (useful
-# for classes with dynamically set attributes). This supports the use of
-# qualified names.
-ignored-classes=optparse.Values,thread._local,_thread._local
-
-# List of members which are set dynamically and missed by pylint inference
-# system, and so shouldn't trigger E1101 when accessed. Python regular
-# expressions are accepted.
-generated-members=
-
-
-[FORMAT]
-
-# Maximum number of characters on a single line.
-max-line-length=80
-
-# TODO(https://github.com/PyCQA/pylint/issues/3352): Direct pylint to exempt
-# lines made too long by directives to pytype.
-
-# Regexp for a line that is allowed to be longer than the limit.
-ignore-long-lines=(?x)(
- ^\s*(\#\ )??$|
- ^\s*(from\s+\S+\s+)?import\s+.+$)
-
-# Allow the body of an if to be on the same line as the test if there is no
-# else.
-single-line-if-stmt=yes
-
-# List of optional constructs for which whitespace checking is disabled. `dict-
-# separator` is used to allow tabulation in dicts, etc.: {1 : 1,\n222: 2}.
-# `trailing-comma` allows a space between comma and closing bracket: (a, ).
-# `empty-line` allows space-only lines.
-no-space-check=
-
-# Maximum number of lines in a module
-max-module-lines=99999
-
-# String used as indentation unit. The internal Google style guide mandates 2
-# spaces. Google's externaly-published style guide says 4, consistent with
-# PEP 8. Here, we use 2 spaces, for conformity with many open-sourced Google
-# projects (like TensorFlow).
-indent-string=' '
-
-# Number of spaces of indent required inside a hanging or continued line.
-indent-after-paren=4
-
-# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
-expected-line-ending-format=
-
-
-[MISCELLANEOUS]
-
-# List of note tags to take in consideration, separated by a comma.
-notes=TODO
-
-
-[STRING]
-
-# This flag controls whether inconsistent-quotes generates a warning when the
-# character used as a quote delimiter is used inconsistently within a module.
-check-quote-consistency=yes
-
-
-[VARIABLES]
-
-# Tells whether we should check for unused import in __init__ files.
-init-import=no
-
-# A regular expression matching the name of dummy variables (i.e. expectedly
-# not used).
-dummy-variables-rgx=^\*{0,2}(_$|unused_|dummy_)
-
-# List of additional names supposed to be defined in builtins. Remember that
-# you should avoid to define new builtins when possible.
-additional-builtins=
-
-# List of strings which can identify a callback function by name. A callback
-# name must start or end with one of those strings.
-callbacks=cb_,_cb
-
-# List of qualified module names which can have objects that can redefine
-# builtins.
-redefining-builtins-modules=six,six.moves,past.builtins,future.builtins,functools
-
-
-[LOGGING]
-
-# Logging modules to check that the string format arguments are in logging
-# function parameter format
-logging-modules=logging,absl.logging,tensorflow.io.logging
-
-
-[SIMILARITIES]
-
-# Minimum lines number of a similarity.
-min-similarity-lines=4
-
-# Ignore comments when computing similarities.
-ignore-comments=yes
-
-# Ignore docstrings when computing similarities.
-ignore-docstrings=yes
-
-# Ignore imports when computing similarities.
-ignore-imports=no
-
-
-[SPELLING]
-
-# Spelling dictionary name. Available dictionaries: none. To make it working
-# install python-enchant package.
-spelling-dict=
-
-# List of comma separated words that should not be checked.
-spelling-ignore-words=
-
-# A path to a file that contains private dictionary; one word per line.
-spelling-private-dict-file=
-
-# Tells whether to store unknown words to indicated private dictionary in
-# --spelling-private-dict-file option instead of raising a message.
-spelling-store-unknown-words=no
-
-
-[IMPORTS]
-
-# Deprecated modules which should not be used, separated by a comma
-deprecated-modules=regsub,
- TERMIOS,
- Bastion,
- rexec,
- sets
-
-# Create a graph of every (i.e. internal and external) dependencies in the
-# given file (report RP0402 must not be disabled)
-import-graph=
-
-# Create a graph of external dependencies in the given file (report RP0402 must
-# not be disabled)
-ext-import-graph=
-
-# Create a graph of internal dependencies in the given file (report RP0402 must
-# not be disabled)
-int-import-graph=
-
-# Force import order to recognize a module as part of the standard
-# compatibility libraries.
-known-standard-library=
-
-# Force import order to recognize a module as part of a third party library.
-known-third-party=enchant, absl
-
-# Analyse import fallback blocks. This can be used to support both Python 2 and
-# 3 compatible code, which means that the block might have code that exists
-# only in one or another interpreter, leading to false positives when analysed.
-analyse-fallback-blocks=no
-
-
-[CLASSES]
-
-# List of method names used to declare (i.e. assign) instance attributes.
-defining-attr-methods=__init__,
- __new__,
- setUp
-
-# List of member names, which should be excluded from the protected access
-# warning.
-exclude-protected=_asdict,
- _fields,
- _replace,
- _source,
- _make
-
-# List of valid names for the first argument in a class method.
-valid-classmethod-first-arg=cls,
- class_
-
-# List of valid names for the first argument in a metaclass class method.
-valid-metaclass-classmethod-first-arg=mcs
-
-
-[EXCEPTIONS]
-
-# Exceptions that will emit a warning when being caught. Defaults to
-# "Exception"
-overgeneral-exceptions=StandardError,
- Exception,
- BaseException
diff --git a/benchmarks/perf-tool/.style.yapf b/benchmarks/perf-tool/.style.yapf
deleted file mode 100644
index 39b663a7a..000000000
--- a/benchmarks/perf-tool/.style.yapf
+++ /dev/null
@@ -1,10 +0,0 @@
-[style]
-COLUMN_LIMIT: 80
-DEDENT_CLOSING_BRACKETS: True
-INDENT_DICTIONARY_VALUE: True
-SPLIT_ALL_COMMA_SEPARATED_VALUES: True
-SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED: True
-SPLIT_BEFORE_CLOSING_BRACKET: True
-SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN: True
-SPLIT_BEFORE_FIRST_ARGUMENT: True
-SPLIT_BEFORE_NAMED_ASSIGNS: True
diff --git a/benchmarks/perf-tool/README.md b/benchmarks/perf-tool/README.md
deleted file mode 100644
index eb4ac0dc1..000000000
--- a/benchmarks/perf-tool/README.md
+++ /dev/null
@@ -1,279 +0,0 @@
-# OpenSearch k-NN Benchmarking
-- [Welcome!](#welcome)
-- [Install Prerequisites](#install-prerequisites)
-- [Usage](#usage)
-- [Contributing](#contributing)
-
-## Welcome!
-
-This directory contains the code related to benchmarking the k-NN plugin.
-Benchmarks can be run against any OpenSearch cluster with the k-NN plugin
-installed. Benchmarks are highly configurable using the test configuration
-file.
-
-## Install Prerequisites
-
-### Python
-
-Python 3.7 or above is required.
-
-### Pip
-
-Use pip to install the necessary requirements:
-
-```
-pip install -r requirements.txt
-```
-
-## Usage
-
-### Quick Start
-
-In order to run a benchmark, you must first create a test configuration yml
-file. Checkout [this example](https://github.com/opensearch-project/k-NN/blob/main/benchmarks/perf-tool/sample-configs) file
-for benchmarking *faiss*'s IVF method. This file contains the definition for
-the benchmark that you want to run. At the top are
-[test parameters](#test-parameters). These define high level settings of the
-test, such as the endpoint of the OpenSearch cluster.
-
-Next, you define the actions that the test will perform. These actions are
-referred to as steps. First, you can define "setup" steps. These are steps that
-are run once at the beginning of the execution to configure the cluster how you
-want it. These steps do not contribute to the final metrics.
-
-After that, you define the "steps". These are the steps that the test will be
-collecting metrics on. Each step emits certain metrics. These are run
-multiple times, depending on the test parameter "num_runs". At the end of the
-execution of all of the runs, the metrics from each run are collected and
-averaged.
-
-Lastly, you define the "cleanup" steps. The "cleanup" steps are executed after
-each test run. For instance, if you are measuring index performance, you may
-want to delete the index after each run.
-
-To run the test, execute the following command:
-```
-python knn-perf-tool.py [--log LOGLEVEL] test config-path.yml output.json
-
---log log level of tool, options are: info, debug, warning, error, critical
-```
-
-The output will be a json document containing the results.
-
-Additionally, you can get the difference between two test runs using the diff
-command:
-```
-python knn-perf-tool.py [--log LOGLEVEL] diff result1.json result2.json
-
---log log level of tool, options are: info, debug, warning, error, critical
-```
-
-The output will be the delta between the two metrics.
-
-### Test Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| endpoint | Endpoint OpenSearch cluster is running on | localhost |
-| test_name | Name of test | No default |
-| test_id | String ID of test | No default |
-| num_runs | Number of runs to execute steps | 1 |
-| show_runs | Whether to output each run in addition to the total summary | false |
-| setup | List of steps to run once before metric collection starts | [] |
-| steps | List of steps that make up one test run. Metrics will be collected on these steps. | No default |
-| cleanup | List of steps to run after each test run | [] |
-
-### Steps
-
-Included are the list of steps that are currently supported. Each step contains
-a set of parameters that are passed in the test configuration file and a set
-of metrics that the test produces.
-
-#### create_index
-
-Creates an OpenSearch index.
-
-##### Parameters
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| index_name | Name of index to create | No default |
-| index_spec | Path to index specification | No default |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end. | ms |
-
-#### disable_refresh
-
-Disables refresh for all indices in the cluster.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end. | ms |
-
-#### refresh_index
-
-Refreshes an OpenSearch index.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| index_name | Name of index to refresh | No default |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end. | ms |
-| store_kb | Size of index after refresh completes | KB |
-
-#### force_merge
-
-Force merges an index to a specified number of segments.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| index_name | Name of index to force merge | No default |
-| max_num_segments | Number of segments to force merge to | No default |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end. | ms |
-
-#### train_model
-
-Trains a model.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| model_id | Model id to set | Test |
-| train_index | Index to pull training data from | No default |
-| train_field | Field to pull training data from | No default |
-| dimension | Dimension of model | No default |
-| description | Description of model | No default |
-| max_training_vector_count | Number of training vectors to used | No default |
-| method_spec | Path to method specification | No default |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end | ms |
-
-#### delete_model
-
-Deletes a model from the cluster.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| model_id | Model id to delete | Test |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end | ms |
-
-#### delete_index
-
-Deletes an index from the cluster.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| index_name | Name of index to delete | No default |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Time to execute step end to end | ms |
-
-#### ingest
-
-Ingests a dataset of vectors into the cluster.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| index_name | Name of index to ingest into | No default |
-| field_name | Name of field to ingest into | No default |
-| bulk_size | Documents per bulk request | 300 |
-| dataset_format | Format the data-set is in. Currently hdf5 and bigann is supported. The hdf5 file must be organized in the same way that the ann-benchmarks organizes theirs. | 'hdf5' |
-| dataset_path | Path to data-set | No default |
-| doc_count | Number of documents to create from data-set | Size of the data-set |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Total time to ingest the dataset into the index.| ms |
-
-#### query
-
-Runs a set of queries against an index.
-
-##### Parameters
-
-| Parameter Name | Description | Default |
-| ----------- | ----------- | ----------- |
-| k | Number of neighbors to return on search | 100 |
-| r | r value in Recall@R | 1 |
-| index_name | Name of index to search | No default |
-| field_name | Name field to search | No default |
-| calculate_recall | Whether to calculate recall values | False |
-| dataset_format | Format the dataset is in. Currently hdf5 and bigann is supported. The hdf5 file must be organized in the same way that the ann-benchmarks organizes theirs. | 'hdf5' |
-| dataset_path | Path to dataset | No default |
-| neighbors_format | Format the neighbors dataset is in. Currently hdf5 and bigann is supported. The hdf5 file must be organized in the same way that the ann-benchmarks organizes theirs. | 'hdf5' |
-| neighbors_path | Path to neighbors dataset | No default |
-| query_count | Number of queries to create from data-set | Size of the data-set |
-
-##### Metrics
-
-| Metric Name | Description | Unit |
-| ----------- | ----------- | ----------- |
-| took | Took times returned per query aggregated as total, p50, p90 and p99 (when applicable) | ms |
-| memory_kb | Native memory k-NN is using at the end of the query workload | KB |
-| recall@R | ratio of top R results from the ground truth neighbors that are in the K results returned by the plugin | float 0.0-1.0 |
-| recall@K | ratio of results returned that were ground truth nearest neighbors | float 0.0-1.0 |
-
-## Contributing
-
-### Linting
-
-Use pylint to lint the code:
-```
-pylint knn-perf-tool.py okpt/**/*.py okpt/**/**/*.py
-```
-
-### Formatting
-
-We use yapf and the google style to format our code. After installing yapf, you can format your code by running:
-
-```
-yapf --style google knn-perf-tool.py okpt/**/*.py okpt/**/**/*.py
-```
-
-### Updating requirements
-
-Add new requirements to "requirements.in" and run `pip-compile`
diff --git a/benchmarks/perf-tool/dataset/data.hdf5 b/benchmarks/perf-tool/dataset/data.hdf5
deleted file mode 100644
index c9268606d..000000000
Binary files a/benchmarks/perf-tool/dataset/data.hdf5 and /dev/null differ
diff --git a/benchmarks/perf-tool/knn-perf-tool.py b/benchmarks/perf-tool/knn-perf-tool.py
deleted file mode 100644
index 48eedc427..000000000
--- a/benchmarks/perf-tool/knn-perf-tool.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-"""Script for user to run the testing tool."""
-
-import okpt.main
-
-okpt.main.main()
diff --git a/benchmarks/perf-tool/okpt/diff/diff.py b/benchmarks/perf-tool/okpt/diff/diff.py
deleted file mode 100644
index 23f424ab9..000000000
--- a/benchmarks/perf-tool/okpt/diff/diff.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Provides the Diff class."""
-
-from enum import Enum
-from typing import Any, Dict, Tuple
-
-
-class InvalidTestResultsError(Exception):
- """Exception raised when the test results are invalid.
-
- The results can be invalid if they have different fields, non-numeric
- values, or if they don't follow the standard result format.
- """
- def __init__(self, msg: str):
- self.message = msg
- super().__init__(self.message)
-
-
-def _is_numeric(a) -> bool:
- return isinstance(a, (int, float))
-
-
-class TestResultFields(str, Enum):
- METADATA = 'metadata'
- RESULTS = 'results'
- TEST_PARAMETERS = 'test_parameters'
-
-
-class TestResultNames(str, Enum):
- BASE = 'base_result'
- CHANGED = 'changed_result'
-
-
-class Diff:
- """Diff class for validating and diffing two test result files.
-
- Methods:
- diff: Returns the diff between two test results. (changed - base)
- """
- def __init__(
- self,
- base_result: Dict[str,
- Any],
- changed_result: Dict[str,
- Any],
- metadata: bool
- ):
- """Initializes test results and validate them."""
- self.base_result = base_result
- self.changed_result = changed_result
- self.metadata = metadata
-
- # make sure results have proper test result fields
- is_valid, key, result = self._validate_keys()
- if not is_valid:
- raise InvalidTestResultsError(
- f'{result} has a missing or invalid key `{key}`.'
- )
-
- self.base_results = self.base_result[TestResultFields.RESULTS]
- self.changed_results = self.changed_result[TestResultFields.RESULTS]
-
- # make sure results have the same fields
- is_valid, key, result = self._validate_structure()
- if not is_valid:
- raise InvalidTestResultsError(
- f'key `{key}` is not present in {result}.'
- )
-
- # make sure results have numeric values
- is_valid, key, result = self._validate_types()
- if not is_valid:
- raise InvalidTestResultsError(
- f'key `{key}` in {result} points to a non-numeric value.'
- )
-
- def _validate_keys(self) -> Tuple[bool, str, str]:
- """Ensure both test results have `metadata` and `results` keys."""
- check_keydict = lambda key, res: key in res and isinstance(
- res[key], dict)
-
- # check if results have a `metadata` field and if `metadata` is a dict
- if self.metadata:
- if not check_keydict(TestResultFields.METADATA, self.base_result):
- return (False, TestResultFields.METADATA, TestResultNames.BASE)
- if not check_keydict(TestResultFields.METADATA,
- self.changed_result):
- return (
- False,
- TestResultFields.METADATA,
- TestResultNames.CHANGED
- )
- # check if results have a `results` field and `results` is a dict
- if not check_keydict(TestResultFields.RESULTS, self.base_result):
- return (False, TestResultFields.RESULTS, TestResultNames.BASE)
- if not check_keydict(TestResultFields.RESULTS, self.changed_result):
- return (False, TestResultFields.RESULTS, TestResultNames.CHANGED)
- return (True, '', '')
-
- def _validate_structure(self) -> Tuple[bool, str, str]:
- """Ensure both test results have the same keys."""
- for k in self.base_results:
- if not k in self.changed_results:
- return (False, k, TestResultNames.CHANGED)
- for k in self.changed_results:
- if not k in self.base_results:
- return (False, k, TestResultNames.BASE)
- return (True, '', '')
-
- def _validate_types(self) -> Tuple[bool, str, str]:
- """Ensure both test results have numeric values."""
- for k, v in self.base_results.items():
- if not _is_numeric(v):
- return (False, k, TestResultNames.BASE)
- for k, v in self.changed_results.items():
- if not _is_numeric(v):
- return (False, k, TestResultNames.BASE)
- return (True, '', '')
-
- def diff(self) -> Dict[str, Any]:
- """Return the diff between the two test results. (changed - base)"""
- results_diff = {
- key: self.changed_results[key] - self.base_results[key]
- for key in self.base_results
- }
-
- # add metadata if specified
- if self.metadata:
- return {
- f'{TestResultNames.BASE}_{TestResultFields.METADATA}':
- self.base_result[TestResultFields.METADATA],
- f'{TestResultNames.CHANGED}_{TestResultFields.METADATA}':
- self.changed_result[TestResultFields.METADATA],
- 'diff':
- results_diff
- }
- return results_diff
diff --git a/benchmarks/perf-tool/okpt/io/args.py b/benchmarks/perf-tool/okpt/io/args.py
deleted file mode 100644
index f8c5d8809..000000000
--- a/benchmarks/perf-tool/okpt/io/args.py
+++ /dev/null
@@ -1,178 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Parses and defines command line arguments for the program.
-
-Defines the subcommands `test` and `diff` and the corresponding
-files that are required by each command.
-
-Functions:
- define_args(): Define the command line arguments.
- get_args(): Returns a dictionary of the command line args.
-"""
-
-import argparse
-import sys
-from dataclasses import dataclass
-from io import TextIOWrapper
-from typing import Union
-
-_read_type = argparse.FileType('r')
-_write_type = argparse.FileType('w')
-
-
-def _add_config(parser, name, **kwargs):
- """"Add configuration file path argument."""
- opts = {
- 'type': _read_type,
- 'help': 'Path of configuration file.',
- 'metavar': 'config_path',
- **kwargs,
- }
- parser.add_argument(name, **opts)
-
-
-def _add_result(parser, name, **kwargs):
- """"Add results files paths argument."""
- opts = {
- 'type': _read_type,
- 'help': 'Path of one result file.',
- 'metavar': 'result_path',
- **kwargs,
- }
- parser.add_argument(name, **opts)
-
-
-def _add_results(parser, name, **kwargs):
- """"Add results files paths argument."""
- opts = {
- 'nargs': '+',
- 'type': _read_type,
- 'help': 'Paths of result files.',
- 'metavar': 'result_paths',
- **kwargs,
- }
- parser.add_argument(name, **opts)
-
-
-def _add_output(parser, name, **kwargs):
- """"Add output file path argument."""
- opts = {
- 'type': _write_type,
- 'help': 'Path of output file.',
- 'metavar': 'output_path',
- **kwargs,
- }
- parser.add_argument(name, **opts)
-
-
-def _add_metadata(parser, name, **kwargs):
- opts = {
- 'action': 'store_true',
- **kwargs,
- }
- parser.add_argument(name, **opts)
-
-
-def _add_test_cmd(subparsers):
- test_parser = subparsers.add_parser('test')
- _add_config(test_parser, 'config')
- _add_output(test_parser, 'output')
-
-
-def _add_diff_cmd(subparsers):
- diff_parser = subparsers.add_parser('diff')
- _add_metadata(diff_parser, '--metadata')
- _add_result(
- diff_parser,
- 'base_result',
- help='Base test result.',
- metavar='base_result'
- )
- _add_result(
- diff_parser,
- 'changed_result',
- help='Changed test result.',
- metavar='changed_result'
- )
- _add_output(diff_parser, '--output', default=sys.stdout)
-
-
-@dataclass
-class TestArgs:
- log: str
- command: str
- config: TextIOWrapper
- output: TextIOWrapper
-
-
-@dataclass
-class DiffArgs:
- log: str
- command: str
- metadata: bool
- base_result: TextIOWrapper
- changed_result: TextIOWrapper
- output: TextIOWrapper
-
-
-def get_args() -> Union[TestArgs, DiffArgs]:
- """Define, parse and return command line args.
-
- Returns:
- A dict containing the command line args.
- """
- parser = argparse.ArgumentParser(
- description=
- 'Run performance tests against the OpenSearch plugin and various ANN '
- 'libaries.'
- )
-
- def define_args():
- """Define tool commands."""
-
- # add log level arg
- parser.add_argument(
- '--log',
- default='info',
- type=str,
- choices=['debug',
- 'info',
- 'warning',
- 'error',
- 'critical'],
- help='Log level of the tool.'
- )
-
- subparsers = parser.add_subparsers(
- title='commands',
- dest='command',
- help='sub-command help'
- )
- subparsers.required = True
-
- # add subcommands
- _add_test_cmd(subparsers)
- _add_diff_cmd(subparsers)
-
- define_args()
- args = parser.parse_args()
- if args.command == 'test':
- return TestArgs(
- log=args.log,
- command=args.command,
- config=args.config,
- output=args.output
- )
- else:
- return DiffArgs(
- log=args.log,
- command=args.command,
- metadata=args.metadata,
- base_result=args.base_result,
- changed_result=args.changed_result,
- output=args.output
- )
diff --git a/benchmarks/perf-tool/okpt/io/config/parsers/base.py b/benchmarks/perf-tool/okpt/io/config/parsers/base.py
deleted file mode 100644
index 795aab1b2..000000000
--- a/benchmarks/perf-tool/okpt/io/config/parsers/base.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Base Parser class.
-
-Classes:
- BaseParser: Base class for config parsers.
-
-Exceptions:
- ConfigurationError: An error in the configuration syntax.
-"""
-
-import os
-from io import TextIOWrapper
-
-import cerberus
-
-from okpt.io.utils import reader
-
-
-class ConfigurationError(Exception):
- """Exception raised for errors in the tool configuration.
-
- Attributes:
- message -- explanation of the error
- """
-
- def __init__(self, message: str):
- self.message = f'{message}'
- super().__init__(self.message)
-
-
-def _get_validator_from_schema_name(schema_name: str):
- """Get the corresponding Cerberus validator from a schema name."""
- curr_file_dir = os.path.dirname(os.path.abspath(__file__))
- schemas_dir = os.path.join(os.path.dirname(curr_file_dir), 'schemas')
- schema_file_path = os.path.join(schemas_dir, f'{schema_name}.yml')
- schema_obj = reader.parse_yaml_from_path(schema_file_path)
- return cerberus.Validator(schema_obj)
-
-
-class BaseParser:
- """Base class for config parsers.
-
- Attributes:
- validator: Cerberus validator for a particular schema
- errors: Cerberus validation errors (if any are found during validation)
-
- Methods:
- parse: Parse config.
- """
-
- def __init__(self, schema_name: str):
- self.validator = _get_validator_from_schema_name(schema_name)
- self.errors = ''
-
- def parse(self, file_obj: TextIOWrapper):
- """Convert file object to dict, while validating against config schema."""
- config_obj = reader.parse_yaml(file_obj)
- is_config_valid = self.validator.validate(config_obj)
- if not is_config_valid:
- raise ConfigurationError(self.validator.errors)
-
- return self.validator.document
diff --git a/benchmarks/perf-tool/okpt/io/config/parsers/test.py b/benchmarks/perf-tool/okpt/io/config/parsers/test.py
deleted file mode 100644
index 34b1752c7..000000000
--- a/benchmarks/perf-tool/okpt/io/config/parsers/test.py
+++ /dev/null
@@ -1,74 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Provides ToolParser.
-
-Classes:
- ToolParser: Tool config parser.
-"""
-from dataclasses import dataclass
-from io import TextIOWrapper
-from typing import List
-
-from okpt.io.config.parsers import base
-from okpt.test.steps.base import Step, StepConfig
-from okpt.test.steps.factory import create_step
-
-
-@dataclass
-class TestConfig:
- test_name: str
- test_id: str
- endpoint: str
- num_runs: int
- show_runs: bool
- setup: List[Step]
- steps: List[Step]
- cleanup: List[Step]
-
-
-class TestParser(base.BaseParser):
- """Parser for Test config.
-
- Methods:
- parse: Parse and validate the Test config.
- """
-
- def __init__(self):
- super().__init__('test')
-
- def parse(self, file_obj: TextIOWrapper) -> TestConfig:
- """See base class."""
- config_obj = super().parse(file_obj)
-
- implicit_step_config = dict()
- if 'endpoint' in config_obj:
- implicit_step_config['endpoint'] = config_obj['endpoint']
-
- # Each step should have its own parse - take the config object and check if its valid
- setup = []
- if 'setup' in config_obj:
- setup = [create_step(StepConfig(step["name"], step, implicit_step_config)) for step in config_obj['setup']]
-
- steps = [create_step(StepConfig(step["name"], step, implicit_step_config)) for step in config_obj['steps']]
-
- cleanup = []
- if 'cleanup' in config_obj:
- cleanup = [create_step(StepConfig(step["name"], step, implicit_step_config)) for step
- in config_obj['cleanup']]
-
- test_config = TestConfig(
- endpoint=config_obj['endpoint'],
- test_name=config_obj['test_name'],
- test_id=config_obj['test_id'],
- num_runs=config_obj['num_runs'],
- show_runs=config_obj['show_runs'],
- setup=setup,
- steps=steps,
- cleanup=cleanup
- )
-
- return test_config
diff --git a/benchmarks/perf-tool/okpt/io/config/parsers/util.py b/benchmarks/perf-tool/okpt/io/config/parsers/util.py
deleted file mode 100644
index cecb9f2d0..000000000
--- a/benchmarks/perf-tool/okpt/io/config/parsers/util.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Utility functions for parsing"""
-
-
-from okpt.io.config.parsers.base import ConfigurationError
-from okpt.io.dataset import HDF5DataSet, BigANNNeighborDataSet, \
- BigANNVectorDataSet, DataSet, Context
-
-
-def parse_dataset(dataset_format: str, dataset_path: str,
- context: Context) -> DataSet:
- if dataset_format == 'hdf5':
- return HDF5DataSet(dataset_path, context)
-
- if dataset_format == 'bigann' and context == Context.NEIGHBORS:
- return BigANNNeighborDataSet(dataset_path)
-
- if dataset_format == 'bigann':
- return BigANNVectorDataSet(dataset_path)
-
- raise Exception("Unsupported data-set format")
-
-
-def parse_string_param(key: str, first_map, second_map, default) -> str:
- value = first_map.get(key)
- if value is not None:
- if type(value) is str:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- value = second_map.get(key)
- if value is not None:
- if type(value) is str:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- if default is None:
- raise ConfigurationError("{} must be set".format(key))
- return default
-
-
-def parse_int_param(key: str, first_map, second_map, default) -> int:
- value = first_map.get(key)
- if value is not None:
- if type(value) is int:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- value = second_map.get(key)
- if value is not None:
- if type(value) is int:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- if default is None:
- raise ConfigurationError("{} must be set".format(key))
- return default
-
-
-def parse_bool_param(key: str, first_map, second_map, default) -> bool:
- value = first_map.get(key)
- if value is not None:
- if type(value) is bool:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- value = second_map.get(key)
- if value is not None:
- if type(value) is bool:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- if default is None:
- raise ConfigurationError("{} must be set".format(key))
- return default
-
-
-def parse_dict_param(key: str, first_map, second_map, default) -> dict:
- value = first_map.get(key)
- if value is not None:
- if type(value) is dict:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- value = second_map.get(key)
- if value is not None:
- if type(value) is dict:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- if default is None:
- raise ConfigurationError("{} must be set".format(key))
- return default
-
-
-def parse_list_param(key: str, first_map, second_map, default) -> list:
- value = first_map.get(key)
- if value is not None:
- if type(value) is list:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- value = second_map.get(key)
- if value is not None:
- if type(value) is list:
- return value
- raise ConfigurationError("Invalid type for {}".format(key))
-
- if default is None:
- raise ConfigurationError("{} must be set".format(key))
- return default
diff --git a/benchmarks/perf-tool/okpt/io/config/schemas/test.yml b/benchmarks/perf-tool/okpt/io/config/schemas/test.yml
deleted file mode 100644
index 1939a8a31..000000000
--- a/benchmarks/perf-tool/okpt/io/config/schemas/test.yml
+++ /dev/null
@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-# defined using the cerberus validation API
-# https://docs.python-cerberus.org/en/stable/index.html
-endpoint:
- type: string
- default: "localhost"
-test_name:
- type: string
-test_id:
- type: string
-num_runs:
- type: integer
- default: 1
- min: 1
- max: 10000
-show_runs:
- type: boolean
- default: false
-setup:
- type: list
-steps:
- type: list
-cleanup:
- type: list
diff --git a/benchmarks/perf-tool/okpt/io/dataset.py b/benchmarks/perf-tool/okpt/io/dataset.py
deleted file mode 100644
index 4f8bc22a2..000000000
--- a/benchmarks/perf-tool/okpt/io/dataset.py
+++ /dev/null
@@ -1,218 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Defines DataSet interface and implements particular formats
-
-A DataSet is the basic functionality that it can be read in chunks, or
-read completely and reset to the start.
-
-Currently, we support HDF5 formats from ann-benchmarks and big-ann-benchmarks
-datasets.
-
-Classes:
- HDF5DataSet: Format used in ann-benchmarks
- BigANNNeighborDataSet: Neighbor format for big-ann-benchmarks
- BigANNVectorDataSet: Vector format for big-ann-benchmarks
-"""
-import os
-from abc import ABC, ABCMeta, abstractmethod
-from enum import Enum
-from typing import cast
-import h5py
-import numpy as np
-
-import struct
-
-
-class Context(Enum):
- """DataSet context enum. Can be used to add additional context for how a
- data-set should be interpreted.
- """
- INDEX = 1
- QUERY = 2
- NEIGHBORS = 3
-
-
-class DataSet(ABC):
- """DataSet interface. Used for reading data-sets from files.
-
- Methods:
- read: Read a chunk of data from the data-set
- size: Gets the number of items in the data-set
- reset: Resets internal state of data-set to beginning
- """
- __metaclass__ = ABCMeta
-
- @abstractmethod
- def read(self, chunk_size: int):
- pass
-
- @abstractmethod
- def size(self):
- pass
-
- @abstractmethod
- def reset(self):
- pass
-
-
-class HDF5DataSet(DataSet):
- """ Data-set format corresponding to `ANN Benchmarks
- `_
- """
-
- def __init__(self, dataset_path: str, context: Context):
- file = h5py.File(dataset_path)
- self.data = cast(h5py.Dataset, file[self._parse_context(context)])
- self.current = 0
-
- def read(self, chunk_size: int):
- if self.current >= self.size():
- return None
-
- end_i = self.current + chunk_size
- if end_i > self.size():
- end_i = self.size()
-
- v = cast(np.ndarray, self.data[self.current:end_i])
- self.current = end_i
- return v
-
- def size(self):
- return self.data.len()
-
- def reset(self):
- self.current = 0
-
- @staticmethod
- def _parse_context(context: Context) -> str:
- if context == Context.NEIGHBORS:
- return "neighbors"
-
- if context == Context.INDEX:
- return "train"
-
- if context == Context.QUERY:
- return "test"
-
- raise Exception("Unsupported context")
-
-
-class BigANNNeighborDataSet(DataSet):
- """ Data-set format for neighbor data-sets for `Big ANN Benchmarks
- `_"""
-
- def __init__(self, dataset_path: str):
- self.file = open(dataset_path, 'rb')
- self.file.seek(0, os.SEEK_END)
- num_bytes = self.file.tell()
- self.file.seek(0)
-
- if num_bytes < 8:
- raise Exception("File is invalid")
-
- self.num_queries = int.from_bytes(self.file.read(4), "little")
- self.k = int.from_bytes(self.file.read(4), "little")
-
- # According to the website, the number of bytes that will follow will
- # be: num_queries X K x sizeof(uint32_t) bytes + num_queries X K x
- # sizeof(float)
- if (num_bytes - 8) != 2 * (self.num_queries * self.k * 4):
- raise Exception("File is invalid")
-
- self.current = 0
-
- def read(self, chunk_size: int):
- if self.current >= self.size():
- return None
-
- end_i = self.current + chunk_size
- if end_i > self.size():
- end_i = self.size()
-
- v = [[int.from_bytes(self.file.read(4), "little") for _ in
- range(self.k)] for _ in range(end_i - self.current)]
-
- self.current = end_i
- return v
-
- def size(self):
- return self.num_queries
-
- def reset(self):
- self.file.seek(8)
- self.current = 0
-
-
-class BigANNVectorDataSet(DataSet):
- """ Data-set format for vector data-sets for `Big ANN Benchmarks
- `_
- """
-
- def __init__(self, dataset_path: str):
- self.file = open(dataset_path, 'rb')
- self.file.seek(0, os.SEEK_END)
- num_bytes = self.file.tell()
- self.file.seek(0)
-
- if num_bytes < 8:
- raise Exception("File is invalid")
-
- self.num_points = int.from_bytes(self.file.read(4), "little")
- self.dimension = int.from_bytes(self.file.read(4), "little")
- bytes_per_num = self._get_data_size(dataset_path)
-
- if (num_bytes - 8) != self.num_points * self.dimension * bytes_per_num:
- raise Exception("File is invalid")
-
- self.reader = self._value_reader(dataset_path)
- self.current = 0
-
- def read(self, chunk_size: int):
- if self.current >= self.size():
- return None
-
- end_i = self.current + chunk_size
- if end_i > self.size():
- end_i = self.size()
-
- v = np.asarray([self._read_vector() for _ in
- range(end_i - self.current)])
- self.current = end_i
- return v
-
- def _read_vector(self):
- return np.asarray([self.reader(self.file) for _ in
- range(self.dimension)])
-
- def size(self):
- return self.num_points
-
- def reset(self):
- self.file.seek(8) # Seek to 8 bytes to skip re-reading metadata
- self.current = 0
-
- @staticmethod
- def _get_data_size(file_name):
- ext = file_name.split('.')[-1]
- if ext == "u8bin":
- return 1
-
- if ext == "fbin":
- return 4
-
- raise Exception("Unknown extension")
-
- @staticmethod
- def _value_reader(file_name):
- ext = file_name.split('.')[-1]
- if ext == "u8bin":
- return lambda file: float(int.from_bytes(file.read(1), "little"))
-
- if ext == "fbin":
- return lambda file: struct.unpack(' TextIOWrapper:
- """Given a file path, get a readable file object.
-
- Args:
- file path
-
- Returns:
- Writeable file object
- """
- return open(path, 'r', encoding='UTF-8')
-
-
-def parse_yaml(file: TextIOWrapper) -> Dict[str, Any]:
- """Parses YAML file from file object.
-
- Args:
- file: file object to parse
-
- Returns:
- A dict representing the YAML file.
- """
- return yaml.load(file, Loader=yaml.SafeLoader)
-
-
-def parse_yaml_from_path(path: str) -> Dict[str, Any]:
- """Parses YAML file from file path.
-
- Args:
- path: file path to parse
-
- Returns:
- A dict representing the YAML file.
- """
- file = reader.get_file_obj(path)
- return parse_yaml(file)
-
-
-def parse_json(file: TextIOWrapper) -> Dict[str, Any]:
- """Parses JSON file from file object.
-
- Args:
- file: file object to parse
-
- Returns:
- A dict representing the JSON file.
- """
- return json.load(file)
-
-
-def parse_json_from_path(path: str) -> Dict[str, Any]:
- """Parses JSON file from file path.
-
- Args:
- path: file path to parse
-
- Returns:
- A dict representing the JSON file.
- """
- file = reader.get_file_obj(path)
- return json.load(file)
diff --git a/benchmarks/perf-tool/okpt/io/utils/writer.py b/benchmarks/perf-tool/okpt/io/utils/writer.py
deleted file mode 100644
index 1f14bfd94..000000000
--- a/benchmarks/perf-tool/okpt/io/utils/writer.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-"""Provides functions for writing to file.
-
-Functions:
- get_file_obj(): Get a writeable file object.
- write_json(): Writes a python dictionary to a JSON file
-"""
-
-import json
-from io import TextIOWrapper
-from typing import Any, Dict, TextIO, Union
-
-
-def get_file_obj(path: str) -> TextIOWrapper:
- """Get a writeable file object from a file path.
-
- Args:
- file path
-
- Returns:
- Writeable file object
- """
- return open(path, 'w', encoding='UTF-8')
-
-
-def write_json(data: Dict[str, Any],
- file: Union[TextIOWrapper, TextIO],
- pretty=False):
- """Writes a dictionary to a JSON file.
-
- Args:
- data: A dict to write to JSON.
- file: Path of output file.
- """
- indent = 2 if pretty else 0
- json.dump(data, file, indent=indent)
diff --git a/benchmarks/perf-tool/okpt/main.py b/benchmarks/perf-tool/okpt/main.py
deleted file mode 100644
index 3e6e022d4..000000000
--- a/benchmarks/perf-tool/okpt/main.py
+++ /dev/null
@@ -1,55 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-""" Runner script that serves as the main controller of the testing tool."""
-
-import logging
-import sys
-from typing import cast
-
-from okpt.diff import diff
-from okpt.io import args
-from okpt.io.config.parsers import test
-from okpt.io.utils import reader, writer
-from okpt.test import runner
-
-
-def main():
- """Main function of entry module."""
- cli_args = args.get_args()
- output = cli_args.output
- if cli_args.log:
- log_level = getattr(logging, cli_args.log.upper())
- logging.basicConfig(level=log_level)
-
- if cli_args.command == 'test':
- cli_args = cast(args.TestArgs, cli_args)
-
- # parse config
- parser = test.TestParser()
- test_config = parser.parse(cli_args.config)
- logging.info('Configs are valid.')
-
- # run tests
- test_runner = runner.TestRunner(test_config=test_config)
- test_result = test_runner.execute()
-
- # write test results
- logging.debug(
- f'Test Result:\n {writer.write_json(test_result, sys.stdout, pretty=True)}'
- )
- writer.write_json(test_result, output, pretty=True)
- elif cli_args.command == 'diff':
- cli_args = cast(args.DiffArgs, cli_args)
-
- # parse test results
- base_result = reader.parse_json(cli_args.base_result)
- changed_result = reader.parse_json(cli_args.changed_result)
-
- # get diff
- diff_result = diff.Diff(base_result, changed_result,
- cli_args.metadata).diff()
- writer.write_json(data=diff_result, file=output, pretty=True)
diff --git a/benchmarks/perf-tool/okpt/test/__init__.py b/benchmarks/perf-tool/okpt/test/__init__.py
deleted file mode 100644
index ff4fd04d1..000000000
--- a/benchmarks/perf-tool/okpt/test/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
diff --git a/benchmarks/perf-tool/okpt/test/profile.py b/benchmarks/perf-tool/okpt/test/profile.py
deleted file mode 100644
index d96860f9a..000000000
--- a/benchmarks/perf-tool/okpt/test/profile.py
+++ /dev/null
@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Provides decorators to profile functions.
-
-The decorators work by adding a `measureable` (time, memory, etc) field to a
-dictionary returned by the wrapped function. So the wrapped functions must
-return a dictionary in order to be profiled.
-"""
-import functools
-import time
-from typing import Callable
-
-
-class TimerStoppedWithoutStartingError(Exception):
- """Error raised when Timer is stopped without having been started."""
-
- def __init__(self):
- super().__init__()
- self.message = 'Timer must call start() before calling end().'
-
-
-class _Timer():
- """Timer class for timing.
-
- Methods:
- start: Starts the timer.
- end: Stops the timer and returns the time elapsed since start.
-
- Raises:
- TimerStoppedWithoutStartingError: Timer must start before ending.
- """
-
- def __init__(self):
- self.start_time = None
-
- def start(self):
- """Starts the timer."""
- self.start_time = time.perf_counter()
-
- def end(self) -> float:
- """Stops the timer.
-
- Returns:
- The time elapsed in milliseconds.
- """
- # ensure timer has started before ending
- if self.start_time is None:
- raise TimerStoppedWithoutStartingError()
-
- elapsed = (time.perf_counter() - self.start_time) * 1000
- self.start_time = None
- return elapsed
-
-
-def took(f: Callable):
- """Profiles a functions execution time.
-
- Args:
- f: Function to profile.
-
- Returns:
- A function that wraps the passed in function and adds a time took field
- to the return value.
- """
-
- @functools.wraps(f)
- def wrapper(*args, **kwargs):
- """Wrapper function."""
- timer = _Timer()
- timer.start()
- result = f(*args, **kwargs)
- time_took = timer.end()
-
- # if result already has a `took` field, don't modify the result
- if isinstance(result, dict) and 'took' in result:
- return result
- # `result` may not be a dictionary, so it may not be unpackable
- elif isinstance(result, dict):
- return {**result, 'took': time_took}
- return {'took': time_took}
-
- return wrapper
diff --git a/benchmarks/perf-tool/okpt/test/runner.py b/benchmarks/perf-tool/okpt/test/runner.py
deleted file mode 100644
index 150154691..000000000
--- a/benchmarks/perf-tool/okpt/test/runner.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Provides a test runner class."""
-import logging
-import platform
-import sys
-from datetime import datetime
-from typing import Any, Dict, List
-
-import psutil
-
-from okpt.io.config.parsers import test
-from okpt.test.test import Test, get_avg
-
-
-def _aggregate_runs(runs: List[Dict[str, Any]]):
- """Aggregates and averages a list of test results.
-
- Args:
- results: A list of test results.
- num_runs: Number of times the tests were ran.
-
- Returns:
- A dictionary containing the averages of the test results.
- """
- aggregate: Dict[str, Any] = {}
- for run in runs:
- for key, value in run.items():
- if key in aggregate:
- aggregate[key].append(value)
- else:
- aggregate[key] = [value]
-
- aggregate = {key: get_avg(value) for key, value in aggregate.items()}
- return aggregate
-
-
-class TestRunner:
- """Test runner class for running tests and aggregating the results.
-
- Methods:
- execute: Run the tests and aggregate the results.
- """
-
- def __init__(self, test_config: test.TestConfig):
- """"Initializes test state."""
- self.test_config = test_config
- self.test = Test(test_config)
-
- def _get_metadata(self):
- """"Retrieves the test metadata."""
- svmem = psutil.virtual_memory()
- return {
- 'test_name':
- self.test_config.test_name,
- 'test_id':
- self.test_config.test_id,
- 'date':
- datetime.now().strftime('%m/%d/%Y %H:%M:%S'),
- 'python_version':
- sys.version,
- 'os_version':
- platform.platform(),
- 'processor':
- platform.processor() + ', ' +
- str(psutil.cpu_count(logical=True)) + ' cores',
- 'memory':
- str(svmem.used) + ' (used) / ' + str(svmem.available) +
- ' (available) / ' + str(svmem.total) + ' (total)',
- }
-
- def execute(self) -> Dict[str, Any]:
- """Runs the tests and aggregates the results.
-
- Returns:
- A dictionary containing the aggregate of test results.
- """
- logging.info('Setting up tests.')
- self.test.setup()
- logging.info('Beginning to run tests.')
- runs = []
- for i in range(self.test_config.num_runs):
- logging.info(
- f'Running test {i + 1} of {self.test_config.num_runs}'
- )
- runs.append(self.test.execute())
-
- logging.info('Finished running tests.')
- aggregate = _aggregate_runs(runs)
-
- # add metadata to test results
- test_result = {
- 'metadata':
- self._get_metadata(),
- 'results':
- aggregate
- }
-
- # include info about all test runs if specified in config
- if self.test_config.show_runs:
- test_result['runs'] = runs
-
- return test_result
diff --git a/benchmarks/perf-tool/okpt/test/steps/base.py b/benchmarks/perf-tool/okpt/test/steps/base.py
deleted file mode 100644
index 829980421..000000000
--- a/benchmarks/perf-tool/okpt/test/steps/base.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-"""Provides base Step interface."""
-
-from dataclasses import dataclass
-from typing import Any, Dict, List
-
-from okpt.test import profile
-
-
-@dataclass
-class StepConfig:
- step_name: str
- config: Dict[str, object]
- implicit_config: Dict[str, object]
-
-
-class Step:
- """Test step interface.
-
- Attributes:
- label: Name of the step.
-
- Methods:
- execute: Run the step and return a step response with the label and
- corresponding measures.
- """
-
- label = 'base_step'
-
- def __init__(self, step_config: StepConfig):
- self.step_config = step_config
-
- def _action(self):
- """Step logic/behavior to be executed and profiled."""
- pass
-
- def _get_measures(self) -> List[str]:
- """Gets the measures for a particular test"""
- pass
-
- def execute(self) -> List[Dict[str, Any]]:
- """Execute step logic while profiling various measures.
-
- Returns:
- Dict containing step label and various step measures.
- """
- action = self._action
-
- # profile the action with measure decorators - add if necessary
- action = getattr(profile, 'took')(action)
-
- result = action()
- if isinstance(result, dict):
- return [{'label': self.label, **result}]
-
- raise ValueError('Invalid return by a step')
diff --git a/benchmarks/perf-tool/okpt/test/steps/factory.py b/benchmarks/perf-tool/okpt/test/steps/factory.py
deleted file mode 100644
index 2b7bcc68d..000000000
--- a/benchmarks/perf-tool/okpt/test/steps/factory.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-"""Factory for creating steps."""
-
-from okpt.io.config.parsers.base import ConfigurationError
-from okpt.test.steps.base import Step, StepConfig
-
-from okpt.test.steps.steps import CreateIndexStep, DisableRefreshStep, RefreshIndexStep, DeleteIndexStep, \
- TrainModelStep, DeleteModelStep, ForceMergeStep, ClearCacheStep, IngestStep, QueryStep
-
-
-def create_step(step_config: StepConfig) -> Step:
- if step_config.step_name == CreateIndexStep.label:
- return CreateIndexStep(step_config)
- elif step_config.step_name == DisableRefreshStep.label:
- return DisableRefreshStep(step_config)
- elif step_config.step_name == RefreshIndexStep.label:
- return RefreshIndexStep(step_config)
- elif step_config.step_name == TrainModelStep.label:
- return TrainModelStep(step_config)
- elif step_config.step_name == DeleteModelStep.label:
- return DeleteModelStep(step_config)
- elif step_config.step_name == DeleteIndexStep.label:
- return DeleteIndexStep(step_config)
- elif step_config.step_name == IngestStep.label:
- return IngestStep(step_config)
- elif step_config.step_name == QueryStep.label:
- return QueryStep(step_config)
- elif step_config.step_name == ForceMergeStep.label:
- return ForceMergeStep(step_config)
- elif step_config.step_name == ClearCacheStep.label:
- return ClearCacheStep(step_config)
-
- raise ConfigurationError(f'Invalid step {step_config.step_name}')
diff --git a/benchmarks/perf-tool/okpt/test/steps/steps.py b/benchmarks/perf-tool/okpt/test/steps/steps.py
deleted file mode 100644
index b61781a6e..000000000
--- a/benchmarks/perf-tool/okpt/test/steps/steps.py
+++ /dev/null
@@ -1,579 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-"""Provides steps for OpenSearch tests.
-
-Some of the OpenSearch operations return a `took` field in the response body,
-so the profiling decorators aren't needed for some functions.
-"""
-import json
-from typing import Any, Dict, List
-
-import numpy as np
-import requests
-import time
-
-from opensearchpy import OpenSearch, RequestsHttpConnection
-
-from okpt.io.config.parsers.base import ConfigurationError
-from okpt.io.config.parsers.util import parse_string_param, parse_int_param, parse_dataset, parse_bool_param
-from okpt.io.dataset import Context
-from okpt.io.utils.reader import parse_json_from_path
-from okpt.test.steps import base
-from okpt.test.steps.base import StepConfig
-
-
-class OpenSearchStep(base.Step):
- """See base class."""
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.endpoint = parse_string_param('endpoint', step_config.config,
- step_config.implicit_config,
- 'localhost')
- default_port = 9200 if self.endpoint == 'localhost' else 80
- self.port = parse_int_param('port', step_config.config,
- step_config.implicit_config, default_port)
- self.opensearch = get_opensearch_client(str(self.endpoint),
- int(self.port))
-
-
-class CreateIndexStep(OpenSearchStep):
- """See base class."""
-
- label = 'create_index'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
- index_spec = parse_string_param('index_spec', step_config.config, {},
- None)
- self.body = parse_json_from_path(index_spec)
- if self.body is None:
- raise ConfigurationError('Index body must be passed in')
-
- def _action(self):
- """Creates an OpenSearch index, applying the index settings/mappings.
-
- Returns:
- An OpenSearch index creation response body.
- """
- self.opensearch.indices.create(index=self.index_name, body=self.body)
- return {}
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class DisableRefreshStep(OpenSearchStep):
- """See base class."""
-
- label = 'disable_refresh'
-
- def _action(self):
- """Disables the refresh interval for an OpenSearch index.
-
- Returns:
- An OpenSearch index settings update response body.
- """
- self.opensearch.indices.put_settings(
- body={'index': {
- 'refresh_interval': -1
- }})
-
- return {}
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class RefreshIndexStep(OpenSearchStep):
- """See base class."""
-
- label = 'refresh_index'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
-
- def _action(self):
- while True:
- try:
- self.opensearch.indices.refresh(index=self.index_name)
- return {'store_kb': get_index_size_in_kb(self.opensearch,
- self.index_name)}
- except:
- pass
-
- def _get_measures(self) -> List[str]:
- return ['took', 'store_kb']
-
-
-class ForceMergeStep(OpenSearchStep):
- """See base class."""
-
- label = 'force_merge'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
- self.max_num_segments = parse_int_param('max_num_segments',
- step_config.config, {}, None)
-
- def _action(self):
- while True:
- try:
- self.opensearch.indices.forcemerge(
- index=self.index_name,
- max_num_segments=self.max_num_segments)
- return {}
- except:
- pass
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-class ClearCacheStep(OpenSearchStep):
- """See base class."""
-
- label = 'clear_cache'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
-
- def _action(self):
- while True:
- try:
- self.opensearch.indices.clear_cache(
- index=self.index_name)
- return {}
- except:
- pass
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class TrainModelStep(OpenSearchStep):
- """See base class."""
-
- label = 'train_model'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
-
- self.model_id = parse_string_param('model_id', step_config.config, {},
- 'Test')
- self.train_index_name = parse_string_param('train_index',
- step_config.config, {}, None)
- self.train_index_field = parse_string_param('train_field',
- step_config.config, {},
- None)
- self.dimension = parse_int_param('dimension', step_config.config, {},
- None)
- self.description = parse_string_param('description', step_config.config,
- {}, 'Default')
- self.max_training_vector_count = parse_int_param(
- 'max_training_vector_count', step_config.config, {}, 10000000000000)
-
- method_spec = parse_string_param('method_spec', step_config.config, {},
- None)
- self.method = parse_json_from_path(method_spec)
- if self.method is None:
- raise ConfigurationError('method must be passed in')
-
- def _action(self):
- """Train a model for an index.
-
- Returns:
- The trained model
- """
-
- # Build body
- body = {
- 'training_index': self.train_index_name,
- 'training_field': self.train_index_field,
- 'description': self.description,
- 'dimension': self.dimension,
- 'method': self.method,
- 'max_training_vector_count': self.max_training_vector_count
- }
-
- # So, we trained the model. Now we need to wait until we have to wait
- # until the model is created. Poll every
- # 1/10 second
- requests.post('http://' + self.endpoint + ':' + str(self.port) +
- '/_plugins/_knn/models/' + str(self.model_id) + '/_train',
- json.dumps(body),
- headers={'content-type': 'application/json'})
-
- sleep_time = 0.1
- timeout = 100000
- i = 0
- while i < timeout:
- time.sleep(sleep_time)
- model_response = get_model(self.endpoint, self.port, self.model_id)
- if 'state' in model_response.keys() and model_response['state'] == \
- 'created':
- return {}
- i += 1
-
- raise TimeoutError('Failed to create model')
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class DeleteModelStep(OpenSearchStep):
- """See base class."""
-
- label = 'delete_model'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
-
- self.model_id = parse_string_param('model_id', step_config.config, {},
- 'Test')
-
- def _action(self):
- """Train a model for an index.
-
- Returns:
- The trained model
- """
- delete_model(self.endpoint, self.port, self.model_id)
- return {}
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class DeleteIndexStep(OpenSearchStep):
- """See base class."""
-
- label = 'delete_index'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
-
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
-
- def _action(self):
- """Delete the index
-
- Returns:
- An empty dict
- """
- delete_index(self.opensearch, self.index_name)
- return {}
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class IngestStep(OpenSearchStep):
- """See base class."""
-
- label = 'ingest'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
- self.field_name = parse_string_param('field_name', step_config.config,
- {}, None)
- self.bulk_size = parse_int_param('bulk_size', step_config.config, {},
- 300)
- self.implicit_config = step_config.implicit_config
- dataset_format = parse_string_param('dataset_format',
- step_config.config, {}, 'hdf5')
- dataset_path = parse_string_param('dataset_path', step_config.config,
- {}, None)
- self.dataset = parse_dataset(dataset_format, dataset_path,
- Context.INDEX)
-
- input_doc_count = parse_int_param('doc_count', step_config.config, {},
- self.dataset.size())
- self.doc_count = min(input_doc_count, self.dataset.size())
-
- def _action(self):
-
- def action(doc_id):
- return {'index': {'_index': self.index_name, '_id': doc_id}}
-
- # Maintain minimal state outside of this loop. For large data sets, too
- # much state may cause out of memory failure
- for i in range(0, self.doc_count, self.bulk_size):
- partition = self.dataset.read(self.bulk_size)
- if partition is None:
- break
- body = bulk_transform(partition, self.field_name, action, i)
- bulk_index(self.opensearch, self.index_name, body)
-
- self.dataset.reset()
-
- return {}
-
- def _get_measures(self) -> List[str]:
- return ['took']
-
-
-class QueryStep(OpenSearchStep):
- """See base class."""
-
- label = 'query'
-
- def __init__(self, step_config: StepConfig):
- super().__init__(step_config)
- self.k = parse_int_param('k', step_config.config, {}, 100)
- self.r = parse_int_param('r', step_config.config, {}, 1)
- self.index_name = parse_string_param('index_name', step_config.config,
- {}, None)
- self.field_name = parse_string_param('field_name', step_config.config,
- {}, None)
- self.calculate_recall = parse_bool_param('calculate_recall',
- step_config.config, {}, False)
- dataset_format = parse_string_param('dataset_format',
- step_config.config, {}, 'hdf5')
- dataset_path = parse_string_param('dataset_path',
- step_config.config, {}, None)
- self.dataset = parse_dataset(dataset_format, dataset_path,
- Context.QUERY)
-
- input_query_count = parse_int_param('query_count',
- step_config.config, {},
- self.dataset.size())
- self.query_count = min(input_query_count, self.dataset.size())
-
- neighbors_format = parse_string_param('neighbors_format',
- step_config.config, {}, 'hdf5')
- neighbors_path = parse_string_param('neighbors_path',
- step_config.config, {}, None)
- self.neighbors = parse_dataset(neighbors_format, neighbors_path,
- Context.NEIGHBORS)
- self.implicit_config = step_config.implicit_config
-
- def _action(self):
-
- def get_body(vec):
- return {
- 'size': self.k,
- 'query': {
- 'knn': {
- self.field_name: {
- 'vector': vec,
- 'k': self.k
- }
- }
- }
- }
-
- results = {}
- query_responses = []
- for _ in range(self.query_count):
- query = self.dataset.read(1)
- if query is None:
- break
- query_responses.append(
- query_index(self.opensearch, self.index_name,
- get_body(query[0]), [self.field_name]))
-
- results['took'] = [
- float(query_response['took']) for query_response in query_responses
- ]
- results['memory_kb'] = get_cache_size_in_kb(self.endpoint, 80)
-
- if self.calculate_recall:
- ids = [[int(hit['_id'])
- for hit in query_response['hits']['hits']]
- for query_response in query_responses]
- results['recall@K'] = recall_at_r(ids, self.neighbors,
- self.k, self.k, self.query_count)
- self.neighbors.reset()
- results[f'recall@{str(self.r)}'] = recall_at_r(
- ids, self.neighbors, self.r, self.k, self.query_count)
- self.neighbors.reset()
-
- self.dataset.reset()
-
- return results
-
- def _get_measures(self) -> List[str]:
- measures = ['took', 'memory_kb']
-
- if self.calculate_recall:
- measures.extend(['recall@K', f'recall@{str(self.r)}'])
-
- return measures
-
-
-# Helper functions - (AKA not steps)
-def bulk_transform(partition: np.ndarray, field_name: str, action,
- offset: int) -> List[Dict[str, Any]]:
- """Partitions and transforms a list of vectors into OpenSearch's bulk
- injection format.
- Args:
- offset: to start counting from
- partition: An array of vectors to transform.
- field_name: field name for action
- action: Bulk API action.
- Returns:
- An array of transformed vectors in bulk format.
- """
- actions = []
- _ = [
- actions.extend([action(i + offset), None])
- for i in range(len(partition))
- ]
- actions[1::2] = [{field_name: vec} for vec in partition.tolist()]
- return actions
-
-
-def delete_index(opensearch: OpenSearch, index_name: str):
- """Deletes an OpenSearch index.
-
- Args:
- opensearch: An OpenSearch client.
- index_name: Name of the OpenSearch index to be deleted.
- """
- opensearch.indices.delete(index=index_name, ignore=[400, 404])
-
-
-def get_model(endpoint, port, model_id):
- """
- Retrieve a model from an OpenSearch cluster
- Args:
- endpoint: Endpoint OpenSearch is running on
- port: Port OpenSearch is running on
- model_id: ID of model to be deleted
- Returns:
- Get model response
- """
- response = requests.get('http://' + endpoint + ':' + str(port) +
- '/_plugins/_knn/models/' + model_id,
- headers={'content-type': 'application/json'})
- return response.json()
-
-
-def delete_model(endpoint, port, model_id):
- """
- Deletes a model from OpenSearch cluster
- Args:
- endpoint: Endpoint OpenSearch is running on
- port: Port OpenSearch is running on
- model_id: ID of model to be deleted
- Returns:
- Deleted model response
- """
- response = requests.delete('http://' + endpoint + ':' + str(port) +
- '/_plugins/_knn/models/' + model_id,
- headers={'content-type': 'application/json'})
- return response.json()
-
-
-def get_opensearch_client(endpoint: str, port: int):
- """
- Get an opensearch client from an endpoint and port
- Args:
- endpoint: Endpoint OpenSearch is running on
- port: Port OpenSearch is running on
- Returns:
- OpenSearch client
-
- """
- # TODO: fix for security in the future
- return OpenSearch(
- hosts=[{
- 'host': endpoint,
- 'port': port
- }],
- use_ssl=False,
- verify_certs=False,
- connection_class=RequestsHttpConnection,
- timeout=60,
- )
-
-
-def recall_at_r(results, neighbor_dataset, r, k, query_count):
- """
- Calculates the recall@R for a set of queries against a ground truth nearest
- neighbor set
- Args:
- results: 2D list containing ids of results returned by OpenSearch.
- results[i][j] i refers to query, j refers to
- result in the query
- neighbor_dataset: 2D dataset containing ids of the true nearest
- neighbors for a set of queries
- r: number of top results to check if they are in the ground truth k-NN
- set.
- k: k value for the query
- query_count: number of queries
- Returns:
- Recall at R
- """
- correct = 0.0
- for query in range(query_count):
- true_neighbors = neighbor_dataset.read(1)
- if true_neighbors is None:
- break
- true_neighbors_set = set(true_neighbors[0][:k])
- for j in range(r):
- if results[query][j] in true_neighbors_set:
- correct += 1.0
-
- return correct / (r * query_count)
-
-
-def get_index_size_in_kb(opensearch, index_name):
- """
- Gets the size of an index in kilobytes
- Args:
- opensearch: opensearch client
- index_name: name of index to look up
- Returns:
- size of index in kilobytes
- """
- return int(
- opensearch.indices.stats(index_name, metric='store')['indices']
- [index_name]['total']['store']['size_in_bytes']) / 1024
-
-
-def get_cache_size_in_kb(endpoint, port):
- """
- Gets the size of the k-NN cache in kilobytes
- Args:
- endpoint: endpoint of OpenSearch cluster
- port: port of endpoint OpenSearch is running on
- Returns:
- size of cache in kilobytes
- """
- response = requests.get('http://' + endpoint + ':' + str(port) +
- '/_plugins/_knn/stats',
- headers={'content-type': 'application/json'})
- stats = response.json()
-
- keys = stats['nodes'].keys()
-
- total_used = 0
- for key in keys:
- total_used += int(stats['nodes'][key]['graph_memory_usage'])
- return total_used
-
-
-def query_index(opensearch: OpenSearch, index_name: str, body: dict,
- excluded_fields: list):
- return opensearch.search(index=index_name,
- body=body,
- _source_excludes=excluded_fields)
-
-
-def bulk_index(opensearch: OpenSearch, index_name: str, body: List):
- return opensearch.bulk(index=index_name, body=body, timeout='5m')
diff --git a/benchmarks/perf-tool/okpt/test/test.py b/benchmarks/perf-tool/okpt/test/test.py
deleted file mode 100644
index dbd65d053..000000000
--- a/benchmarks/perf-tool/okpt/test/test.py
+++ /dev/null
@@ -1,180 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-
-"""Provides a base Test class."""
-from math import floor
-from typing import Any, Dict, List
-
-from okpt.io.config.parsers.test import TestConfig
-from okpt.test.steps.base import Step
-
-
-def get_avg(values: List[Any]):
- """Get average value of a list.
-
- Args:
- values: A list of values.
-
- Returns:
- The average value in the list.
- """
- valid_total = len(values)
- running_sum = 0.0
-
- for value in values:
- if value == -1:
- valid_total -= 1
- continue
- running_sum += value
-
- if valid_total == 0:
- return -1
- return running_sum / valid_total
-
-
-def _pxx(values: List[Any], p: float):
- """Calculates the pXX statistics for a given list.
-
- Args:
- values: List of values.
- p: Percentile (between 0 and 1).
-
- Returns:
- The corresponding pXX metric.
- """
- lowest_percentile = 1 / len(values)
- highest_percentile = (len(values) - 1) / len(values)
-
- # return -1 if p is out of range or if the list doesn't have enough elements
- # to support the specified percentile
- if p < 0 or p > 1:
- return -1.0
- elif p < lowest_percentile or p > highest_percentile:
- return -1.0
- else:
- return float(values[floor(len(values) * p)])
-
-
-def _aggregate_steps(step_results: List[Dict[str, Any]],
- measure_labels=None):
- """Aggregates the steps for a given Test.
-
- The aggregation process extracts the measures from each step and calculates
- the total time spent performing each step measure, including the
- percentile metrics, if possible.
-
- The aggregation process also extracts the test measures by simply summing
- up the respective step measures.
-
- A step measure is formatted as `{step_name}_{measure_name}`, for example,
- {bulk_index}_{took} or {query_index}_{memory}. The braces are not included
- in the actual key string.
-
- Percentile/Total step measures are give as
- `{step_name}_{measure_name}_{percentile|total}`.
-
- Test measures are just step measure sums so they just given as
- `test_{measure_name}`.
-
- Args:
- steps: List of test steps to be aggregated.
- measures: List of step metrics to account for.
-
- Returns:
- A complete test result.
- """
- if measure_labels is None:
- measure_labels = ['took']
- test_measures = {
- f'test_{measure_label}': 0
- for measure_label in measure_labels
- }
- step_measures: Dict[str, Any] = {}
-
- # iterate over all test steps
- for step in step_results:
- step_label = step['label']
-
- step_measure_labels = list(step.keys())
- step_measure_labels.remove('label')
-
- # iterate over all measures in each test step
- for measure_label in step_measure_labels:
-
- step_measure = step[measure_label]
- step_measure_label = f'{step_label}_{measure_label}'
-
- # Add cumulative test measures from steps to test measures
- if measure_label in measure_labels:
- test_measures[f'test_{measure_label}'] += sum(step_measure) if \
- isinstance(step_measure, list) else step_measure
-
- if step_measure_label in step_measures:
- _ = step_measures[step_measure_label].extend(step_measure) \
- if isinstance(step_measure, list) else \
- step_measures[step_measure_label].append(step_measure)
- else:
- step_measures[step_measure_label] = step_measure if \
- isinstance(step_measure, list) else [step_measure]
-
- aggregate = {**test_measures}
- # calculate the totals and percentile statistics for each step measure
- # where relevant
- for step_measure_label, step_measure in step_measures.items():
- step_measure.sort()
-
- aggregate[step_measure_label + '_total'] = float(sum(step_measure))
-
- p50 = _pxx(step_measure, 0.50)
- if p50 != -1:
- aggregate[step_measure_label + '_p50'] = p50
- p90 = _pxx(step_measure, 0.90)
- if p90 != -1:
- aggregate[step_measure_label + '_p90'] = p90
- p99 = _pxx(step_measure, 0.99)
- if p99 != -1:
- aggregate[step_measure_label + '_p99'] = p99
-
- return aggregate
-
-
-class Test:
- """A base Test class, representing a collection of steps to profiled and
- aggregated.
-
- Methods:
- setup: Performs test setup. Usually for steps not intended to be
- profiled.
- run_steps: Runs the test steps, aggregating the results into the
- `step_results` instance field.
- cleanup: Perform test cleanup. Useful for clearing the state of a
- persistent process like OpenSearch. Cleanup steps are executed after
- each run.
- execute: Runs steps, cleans up, and aggregates the test result.
- """
- def __init__(self, test_config: TestConfig):
- """Initializes the test state.
- """
- self.test_config = test_config
- self.setup_steps: List[Step] = test_config.setup
- self.test_steps: List[Step] = test_config.steps
- self.cleanup_steps: List[Step] = test_config.cleanup
-
- def setup(self):
- _ = [step.execute() for step in self.setup_steps]
-
- def _run_steps(self):
- step_results = []
- _ = [step_results.extend(step.execute()) for step in self.test_steps]
- return step_results
-
- def _cleanup(self):
- _ = [step.execute() for step in self.cleanup_steps]
-
- def execute(self):
- results = self._run_steps()
- self._cleanup()
- return _aggregate_steps(results)
diff --git a/benchmarks/perf-tool/requirements.in b/benchmarks/perf-tool/requirements.in
deleted file mode 100644
index fd3555aab..000000000
--- a/benchmarks/perf-tool/requirements.in
+++ /dev/null
@@ -1,7 +0,0 @@
-Cerberus
-opensearch-py
-PyYAML
-numpy
-h5py
-requests
-psutil
diff --git a/benchmarks/perf-tool/requirements.txt b/benchmarks/perf-tool/requirements.txt
deleted file mode 100644
index 2886795dc..000000000
--- a/benchmarks/perf-tool/requirements.txt
+++ /dev/null
@@ -1,39 +0,0 @@
-#
-# This file is autogenerated by pip-compile with python 3.9
-# To update, run:
-#
-# pip-compile
-#
-cached-property==1.5.2
- # via h5py
-cerberus==1.3.4
- # via -r requirements.in
-certifi==2021.5.30
- # via
- # opensearch-py
- # requests
-charset-normalizer==2.0.4
- # via requests
-h5py==3.3.0
- # via -r requirements.in
-idna==3.2
- # via requests
-numpy==1.22.1
- # via
- # -r requirements.in
- # h5py
-opensearch-py==1.0.0
- # via -r requirements.in
-psutil==5.8.0
- # via -r requirements.in
-pyyaml==5.4.1
- # via -r requirements.in
-requests==2.26.0
- # via -r requirements.in
-urllib3==1.26.6
- # via
- # opensearch-py
- # requests
-
-# The following packages are considered to be unsafe in a requirements file:
-# setuptools
diff --git a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/index-spec.json b/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/index-spec.json
deleted file mode 100644
index 5542ef387..000000000
--- a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/index-spec.json
+++ /dev/null
@@ -1,17 +0,0 @@
-{
- "settings": {
- "index": {
- "knn": true,
- "number_of_shards": 3,
- "number_of_replicas": 0
- }
- },
- "mappings": {
- "properties": {
- "target_field": {
- "type": "knn_vector",
- "model_id": "test-model"
- }
- }
- }
-}
diff --git a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/method-spec.json b/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/method-spec.json
deleted file mode 100644
index 1aa7f809f..000000000
--- a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/method-spec.json
+++ /dev/null
@@ -1,8 +0,0 @@
-{
- "name":"ivf",
- "engine":"faiss",
- "parameters":{
- "nlist":16,
- "nprobes": 4
- }
-}
diff --git a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/test.yml b/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/test.yml
deleted file mode 100644
index c8fb42ec4..000000000
--- a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/test.yml
+++ /dev/null
@@ -1,60 +0,0 @@
-endpoint: localhost
-test_name: faiss_sift_ivf
-test_id: "Test workflow for faiss ivf"
-num_runs: 3
-show_runs: true
-setup:
- - name: delete_model
- model_id: test-model
- - name: delete_index
- index_name: target_index
- - name: delete_index
- index_name: train_index
- - name: create_index
- index_name: train_index
- index_spec: sample-configs/faiss-sift-ivf/train-index-spec.json
- - name: ingest
- index_name: train_index
- field_name: train_field
- bulk_size: 500
- dataset_format: hdf5
- dataset_path: ../dataset/sift-128-euclidean.hdf5
- - name: refresh_index
- index_name: train_index
-steps:
- - name: train_model
- model_id: test-model
- train_index: train_index
- train_field: train_field
- dimension: 128
- method_spec: sample-configs/faiss-sift-ivf/method-spec.json
- max_training_vector_count: 1000000000
- - name: create_index
- index_name: target_index
- index_spec: sample-configs/faiss-sift-ivf/index-spec.json
- - name: ingest
- index_name: target_index
- field_name: target_field
- bulk_size: 500
- dataset_format: hdf5
- dataset_path: ../dataset/sift-128-euclidean.hdf5
- - name: refresh_index
- index_name: target_index
- - name: force_merge
- index_name: target_index
- max_num_segments: 10
- - name: query
- k: 100
- r: 1
- calculate_recall: true
- index_name: target_index
- field_name: target_field
- dataset_format: hdf5
- dataset_path: ../dataset/sift-128-euclidean.hdf5
- neighbors_format: hdf5
- neighbors_path: ../dataset/sift-128-euclidean.hdf5
-cleanup:
- - name: delete_model
- model_id: test-model
- - name: delete_index
- index_name: target_index
diff --git a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/train-index-spec.json b/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/train-index-spec.json
deleted file mode 100644
index 00a418e4f..000000000
--- a/benchmarks/perf-tool/sample-configs/faiss-sift-ivf/train-index-spec.json
+++ /dev/null
@@ -1,16 +0,0 @@
-{
- "settings": {
- "index": {
- "number_of_shards": 3,
- "number_of_replicas": 0
- }
- },
- "mappings": {
- "properties": {
- "train_field": {
- "type": "knn_vector",
- "dimension": 128
- }
- }
- }
-}
diff --git a/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/index-spec.json b/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/index-spec.json
deleted file mode 100644
index 75abe7baa..000000000
--- a/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/index-spec.json
+++ /dev/null
@@ -1,28 +0,0 @@
-{
- "settings": {
- "index": {
- "knn": true,
- "knn.algo_param.ef_search": 512,
- "refresh_interval": "10s",
- "number_of_shards": 1,
- "number_of_replicas": 0
- }
- },
- "mappings": {
- "properties": {
- "target_field": {
- "type": "knn_vector",
- "dimension": 128,
- "method": {
- "name": "hnsw",
- "space_type": "l2",
- "engine": "nmslib",
- "parameters": {
- "ef_construction": 512,
- "m": 16
- }
- }
- }
- }
- }
-}
diff --git a/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/test.yml b/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/test.yml
deleted file mode 100644
index deea1ad47..000000000
--- a/benchmarks/perf-tool/sample-configs/nmslib-sift-hnsw/test.yml
+++ /dev/null
@@ -1,36 +0,0 @@
-endpoint: localhost
-test_name: nmslib_sift_hnsw
-test_id: "Test workflow for nmslib hnsw"
-num_runs: 2
-show_runs: false
-setup:
- - name: delete_index
- index_name: target_index
-steps:
- - name: create_index
- index_name: target_index
- index_spec: sample-configs/nmslib-sift-hnsw/index-spec.json
- - name: ingest
- index_name: target_index
- field_name: target_field
- bulk_size: 500
- dataset_format: hdf5
- dataset_path: ../dataset/sift-128-euclidean.hdf5
- - name: refresh_index
- index_name: target_index
- - name: force_merge
- index_name: target_index
- max_num_segments: 10
- - name: query
- k: 100
- r: 1
- calculate_recall: true
- index_name: target_index
- field_name: target_field
- dataset_format: hdf5
- dataset_path: ../dataset/sift-128-euclidean.hdf5
- neighbors_format: hdf5
- neighbors_path: ../dataset/sift-128-euclidean.hdf5
-cleanup:
- - name: delete_index
- index_name: target_index
diff --git a/build-tools/knnplugin-coverage.gradle b/build-tools/knnplugin-coverage.gradle
index 57a0eeeec..eb3582dab 100644
--- a/build-tools/knnplugin-coverage.gradle
+++ b/build-tools/knnplugin-coverage.gradle
@@ -3,12 +3,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
+apply plugin: 'jacoco'
+
+jacoco {
+ toolVersion = "0.8.10"
+}
+
/**
- * ES Plugin build tools don't work with the Gradle Jacoco Plugin to report coverage out of the box.
- * https://github.com/elastic/elasticsearch/issues/28867.
- *
- * This code sets up coverage reporting manually for ES plugin tests. This is complicated because:
- * 1. The ES integTest Task doesn't implement Gradle's JavaForkOptions so we have to manually start the jacoco agent with the test JVM
+ * This code sets up coverage reporting manually for the k-NN plugin tests. This is complicated because:
+ * 1. The OS integTest Task doesn't implement Gradle's JavaForkOptions so we have to manually start the jacoco agent with the test JVM
* 2. The cluster nodes are stopped using 'kill -9' which means jacoco can't dump it's execution output to a file on VM shutdown
* 3. The Java Security Manager prevents JMX from writing execution output to the file.
*
@@ -16,58 +19,31 @@
* cluster is stopped and dump it to a file. Luckily our current security policy seems to allow this. This will also probably
* break if there are multiple nodes in the integTestCluster. But for now... it sorta works.
*/
-apply plugin: 'jacoco'
-
-// Get gradle to generate the required jvm agent arg for us using a dummy tasks of type Test. Unfortunately Elastic's
-// testing tasks don't derive from Test so the jacoco plugin can't do this automatically.
-def jacocoDir = "${buildDir}/jacoco"
-task dummyTest(type: Test) {
- enabled = false
- workingDir = file("/") // Force absolute path to jacoco agent jar
- jacoco {
- destinationFile = file("${jacocoDir}/test.exec")
- destinationFile.parentFile.mkdirs()
- jmx = true
- }
-}
-
-task dummyIntegTest(type: Test) {
- enabled = false
- workingDir = file("/") // Force absolute path to jacoco agent jar
+integTest {
jacoco {
- destinationFile = file("${jacocoDir}/integTest.exec")
- destinationFile.parentFile.mkdirs()
jmx = true
}
-}
-integTest {
- systemProperty 'jacoco.dir', "${jacocoDir}"
+ systemProperty 'jacoco.dir', project.layout.buildDirectory.get().file("jacoco").asFile.absolutePath
systemProperty 'jmx.serviceUrl', "service:jmx:rmi:///jndi/rmi://127.0.0.1:7777/jmxrmi"
}
jacocoTestReport {
dependsOn integTest, test
- executionData.from = [dummyTest.jacoco.destinationFile, dummyIntegTest.jacoco.destinationFile]
- sourceDirectories.from = sourceSets.main.java.sourceDirectories
- classDirectories.from = files(sourceSets.main.java.outputDir)
-
+ executionData.from = [integTest.jacoco.destinationFile, test.jacoco.destinationFile]
reports {
- html.enabled = true // human readable
- csv.enabled = true
- xml.enabled = true // for coverlay
+ html.getRequired().set(true) // human readable
+ csv.getRequired().set(true)
+ xml.getRequired().set(true) // for coverlay
}
}
-afterEvaluate {
- jacocoTestReport.dependsOn integTest
+testClusters.integTest {
+ jvmArgs " ${integTest.jacoco.getAsJvmArg()}"
- testClusters.integTest {
- jvmArgs " ${dummyIntegTest.jacoco.getAsJvmArg()}".replace('javaagent:','javaagent:/')
- systemProperty 'com.sun.management.jmxremote', "true"
- systemProperty 'com.sun.management.jmxremote.authenticate', "false"
- systemProperty 'com.sun.management.jmxremote.port', "7777"
- systemProperty 'com.sun.management.jmxremote.ssl', "false"
- systemProperty 'java.rmi.server.hostname', "127.0.0.1"
- }
+ systemProperty 'com.sun.management.jmxremote', "true"
+ systemProperty 'com.sun.management.jmxremote.authenticate', "false"
+ systemProperty 'com.sun.management.jmxremote.port', "7777"
+ systemProperty 'com.sun.management.jmxremote.ssl', "false"
+ systemProperty 'java.rmi.server.hostname', "127.0.0.1"
}
diff --git a/build.gradle b/build.gradle
index 63a299666..ead0cba08 100644
--- a/build.gradle
+++ b/build.gradle
@@ -4,14 +4,44 @@
*/
import org.opensearch.gradle.test.RestIntegTestTask
+import org.opensearch.gradle.testclusters.OpenSearchCluster
+import org.apache.tools.ant.taskdefs.condition.Os
+import java.nio.file.Paths
+import java.util.concurrent.Callable
buildscript {
ext {
// build.version_qualifier parameter applies to knn plugin artifacts only. OpenSearch version must be set
// explicitly as 'opensearch.version' property, for instance opensearch.version=2.0.0-rc1-SNAPSHOT
- opensearch_version = System.getProperty("opensearch.version", "2.1.0-SNAPSHOT")
+ opensearch_version = System.getProperty("opensearch.version", "2.19.0-SNAPSHOT")
version_qualifier = System.getProperty("build.version_qualifier", "")
opensearch_group = "org.opensearch"
+ isSnapshot = "true" == System.getProperty("build.snapshot", "true")
+ avx2_enabled = System.getProperty("avx2.enabled", "true")
+ nproc_count = System.getProperty("nproc.count", "1")
+ avx512_enabled = System.getProperty("avx512.enabled", "true")
+ // This flag determines whether the CMake build system should apply a custom patch. It prevents build failures
+ // when the cmakeJniLib task is run multiple times. If the build.lib.commit_patches is true, the CMake build
+ // system skips applying the patch if the patches have been applied already. If build.lib.commit_patches is
+ // false, the patches are always applied. To avoid patch conflicts, disable this flag manually after the first
+ // run of buildJniLib
+ apply_lib_patches = System.getProperty("build.lib.apply_patches", "true")
+ // Flag to determine whether cmake build system should commit the patch or not. In automated build environments
+ // set this to false. In dev environments, set to true. If false, repetitive execution of cmakeJniLib may fail.
+ // To prevent this, set build.lib.apply_patches to false after the first cmakeJniLib run.
+ commit_lib_patches = System.getProperty("build.lib.commit_patches", "true")
+
+ version_tokens = opensearch_version.tokenize('-')
+ opensearch_build = version_tokens[0] + '.0'
+ plugin_no_snapshot = opensearch_build
+ if (version_qualifier) {
+ opensearch_build += "-${version_qualifier}"
+ plugin_no_snapshot += "-${version_qualifier}"
+ }
+ if (isSnapshot) {
+ opensearch_build += "-SNAPSHOT"
+ }
+ opensearch_no_snapshot = opensearch_build.replace("-SNAPSHOT","")
}
// This isn't applying from repositories.gradle so repeating git diff it here
@@ -24,6 +54,12 @@ buildscript {
dependencies {
classpath "${opensearch_group}.gradle:build-tools:${opensearch_version}"
+ configurations.all {
+ resolutionStrategy {
+ force("org.eclipse.platform:org.eclipse.core.runtime:4.29.0") // CVE for < 4.29
+ force("org.eclipse.platform:org.eclipse.core.resources:4.20.0") // CVE for < 4.20
+ }
+ }
}
}
@@ -35,9 +71,9 @@ plugins {
id 'java-library'
id 'java-test-fixtures'
id 'idea'
- id 'jacoco'
- id "com.diffplug.spotless" version "6.3.0" apply false
- id 'io.freefair.lombok' version '6.4.3'
+ id "com.diffplug.spotless" version "6.25.0" apply false
+ id 'io.freefair.lombok' version '8.4'
+ id "de.undercouch.download" version "5.3.0"
}
apply from: 'gradle/formatting.gradle'
@@ -46,9 +82,91 @@ apply plugin: 'opensearch.rest-test'
apply plugin: 'opensearch.pluginzip'
apply plugin: 'opensearch.repositories'
+
+def opensearch_tmp_dir = rootProject.file('build/private/opensearch_tmp').absoluteFile
+opensearch_tmp_dir.mkdirs()
+
ext {
- isSnapshot = "true" == System.getProperty("build.snapshot", "true")
projectSubstitutions = [:]
+
+ configureSecurityPlugin = { OpenSearchCluster cluster ->
+ configurations.zipArchive.asFileTree.each {
+ cluster.plugin(provider(new Callable() {
+ @Override
+ RegularFile call() throws Exception {
+ return new RegularFile() {
+ @Override
+ File getAsFile() {
+ return it
+ }
+ }
+ }
+ }))
+ }
+
+ cluster.getNodes().forEach { node ->
+ var creds = node.getCredentials()
+ if (creds.isEmpty()) {
+ creds.add(Map.of('username', 'admin', 'password', 'admin'))
+ } else {
+ creds.get(0).putAll(Map.of('username', 'admin', 'password', 'admin'))
+ }
+ }
+
+ // Config below including files are copied from security demo configuration
+ ['esnode.pem', 'esnode-key.pem', 'root-ca.pem'].forEach { file ->
+ File local = Paths.get(opensearch_tmp_dir.absolutePath, file).toFile()
+ download.run {
+ src "https://raw.githubusercontent.com/opensearch-project/security/main/bwc-test/src/test/resources/security/" + file
+ dest local
+ overwrite false
+ }
+ cluster.extraConfigFile(file, local)
+ }
+
+ // This configuration is copied from the security plugins demo install:
+ // https://github.com/opensearch-project/security/blob/2.11.1.0/tools/install_demo_configuration.sh#L365-L388
+ cluster.setting("plugins.security.ssl.transport.pemcert_filepath", "esnode.pem")
+ cluster.setting("plugins.security.ssl.transport.pemkey_filepath", "esnode-key.pem")
+ cluster.setting("plugins.security.ssl.transport.pemtrustedcas_filepath", "root-ca.pem")
+ cluster.setting("plugins.security.ssl.transport.enforce_hostname_verification", "false")
+ cluster.setting("plugins.security.ssl.http.enabled", "true")
+ cluster.setting("plugins.security.ssl.http.pemcert_filepath", "esnode.pem")
+ cluster.setting("plugins.security.ssl.http.pemkey_filepath", "esnode-key.pem")
+ cluster.setting("plugins.security.ssl.http.pemtrustedcas_filepath", "root-ca.pem")
+ cluster.setting("plugins.security.allow_unsafe_democertificates", "true")
+ cluster.setting("plugins.security.allow_default_init_securityindex", "true")
+ cluster.setting("plugins.security.unsupported.inject_user.enabled", "true")
+
+ cluster.setting("plugins.security.authcz.admin_dn", "\n- CN=kirk,OU=client,O=client,L=test, C=de")
+ cluster.setting('plugins.security.restapi.roles_enabled', '["all_access", "security_rest_api_access"]')
+ cluster.setting('plugins.security.system_indices.enabled', "true")
+ cluster.setting('plugins.security.system_indices.indices', '[' +
+ '".plugins-ml-config", ' +
+ '".plugins-ml-connector", ' +
+ '".plugins-ml-model-group", ' +
+ '".plugins-ml-model", ".plugins-ml-task", ' +
+ '".plugins-ml-conversation-meta", ' +
+ '".plugins-ml-conversation-interactions", ' +
+ '".opendistro-alerting-config", ' +
+ '".opendistro-alerting-alert*", ' +
+ '".opendistro-anomaly-results*", ' +
+ '".opendistro-anomaly-detector*", ' +
+ '".opendistro-anomaly-checkpoints", ' +
+ '".opendistro-anomaly-detection-state", ' +
+ '".opendistro-reports-*", ' +
+ '".opensearch-notifications-*", ' +
+ '".opensearch-notebooks", ' +
+ '".opensearch-observability", ' +
+ '".ql-datasources", ' +
+ '".opendistro-asynchronous-search-response*", ' +
+ '".replication-metadata-store", ' +
+ '".opensearch-knn-models", ' +
+ '".geospatial-ip2geo-data*"' +
+ ']'
+ )
+ cluster.setSecure(true)
+ }
}
allprojects {
@@ -82,12 +200,27 @@ allprojects {
}
}
+configurations {
+ zipArchive
+}
+
publishing {
+ repositories {
+ maven {
+ name = "Snapshots"
+ url = "https://aws.oss.sonatype.org/content/repositories/snapshots"
+ credentials {
+ username "$System.env.SONATYPE_USERNAME"
+ password "$System.env.SONATYPE_PASSWORD"
+ }
+ }
+ }
publications {
pluginZip(MavenPublication) { publication ->
pom {
name = "opensearch-knn"
description = "OpenSearch k-NN plugin"
+ groupId = "org.opensearch.plugin"
licenses {
license {
name = "The Apache License, Version 2.0"
@@ -111,6 +244,9 @@ compileJava {
compileTestJava {
options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor'])
}
+compileTestFixturesJava {
+ options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor'])
+}
def usingRemoteCluster = System.properties.containsKey('tests.rest.cluster') || System.properties.containsKey('tests.cluster')
def usingMultiNode = project.properties.containsKey('numNodes')
@@ -123,10 +259,6 @@ if (!usingRemoteCluster) {
}
}
-jacoco {
- toolVersion = "0.8.7"
-}
-
check.dependsOn spotlessCheck
check.dependsOn jacocoTestReport
@@ -160,30 +292,61 @@ dependencies {
api "org.opensearch:opensearch:${opensearch_version}"
compileOnly "org.opensearch.plugin:opensearch-scripting-painless-spi:${versions.opensearch}"
api group: 'com.google.guava', name: 'failureaccess', version:'1.0.1'
- api group: 'com.google.guava', name: 'guava', version:'30.0-jre'
+ api group: 'com.google.guava', name: 'guava', version:'32.1.3-jre'
api group: 'commons-lang', name: 'commons-lang', version: '2.6'
testFixturesImplementation "org.opensearch.test:framework:${opensearch_version}"
+ testImplementation group: 'net.bytebuddy', name: 'byte-buddy', version: '1.15.10'
+ testImplementation group: 'org.objenesis', name: 'objenesis', version: '3.3'
+ testImplementation group: 'net.bytebuddy', name: 'byte-buddy-agent', version: '1.15.4'
+ testFixturesImplementation "org.opensearch:common-utils:${version}"
+ implementation 'com.github.oshi:oshi-core:6.4.13'
+ api "net.java.dev.jna:jna:5.13.0"
+ api "net.java.dev.jna:jna-platform:5.13.0"
+ implementation 'org.slf4j:slf4j-api:1.7.36'
+
+ zipArchive group: 'org.opensearch.plugin', name:'opensearch-security', version: "${opensearch_build}"
}
-
-def opensearch_tmp_dir = rootProject.file('build/private/opensearch_tmp').absoluteFile
-opensearch_tmp_dir.mkdirs()
+task windowsPatches(type:Exec) {
+ commandLine 'cmd', '/c', "Powershell -File $rootDir\\scripts\\windowsScript.ps1"
+}
task cmakeJniLib(type:Exec) {
workingDir 'jni'
- commandLine 'cmake', '.', "-DKNN_PLUGIN_VERSION=${opensearch_version}"
+ def args = []
+ args.add("cmake")
+ args.add(".")
+ args.add("-DKNN_PLUGIN_VERSION=${opensearch_version}")
+ args.add("-DAVX2_ENABLED=${avx2_enabled}")
+ args.add("-DAVX512_ENABLED=${avx512_enabled}")
+ args.add("-DCOMMIT_LIB_PATCHES=${commit_lib_patches}")
+ args.add("-DAPPLY_LIB_PATCHES=${apply_lib_patches}")
+ if (Os.isFamily(Os.FAMILY_WINDOWS)) {
+ dependsOn windowsPatches
+ args.add("-G")
+ args.add("Unix Makefiles")
+ args.add("-DBLAS_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll")
+ args.add("-DLAPACK_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll")
+ }
+ commandLine args
}
task buildJniLib(type:Exec) {
dependsOn cmakeJniLib
workingDir 'jni'
- commandLine 'make', 'opensearchknn_nmslib', 'opensearchknn_faiss'
+ commandLine 'make', 'opensearchknn_nmslib', 'opensearchknn_faiss', 'opensearchknn_common', '-j', "${nproc_count}"
}
test {
dependsOn buildJniLib
systemProperty 'tests.security.manager', 'false'
systemProperty "java.library.path", "$rootDir/jni/release"
+ //this change enables mockito-inline that supports mocking of static classes/calls
+ systemProperty "jdk.attach.allowAttachSelf", true
+ if (Os.isFamily(Os.FAMILY_WINDOWS)) {
+ // Add the paths of built JNI libraries and its dependent libraries to PATH variable in System variables
+ environment('PATH', System.getenv('PATH') + ";$rootDir/jni/release" + ";$rootDir/src/main/resources/windowsDependencies")
+ }
}
def _numNodes = findProperty('numNodes') as Integer ?: 1
@@ -197,9 +360,20 @@ integTest {
// allows integration test classes to access test resource from project root path
systemProperty('project.root', project.rootDir.absolutePath)
- systemProperty "https", System.getProperty("https")
- systemProperty "user", System.getProperty("user")
- systemProperty "password", System.getProperty("password")
+ var is_https = System.getProperty("https")
+ var user = System.getProperty("user")
+ var password = System.getProperty("password")
+
+ if (System.getProperty("security.enabled") != null) {
+ // If security is enabled, set is_https/user/password defaults
+ is_https = is_https == null ? "true" : is_https
+ user = user == null ? "admin" : user
+ password = password == null ? "admin" : password
+ }
+
+ systemProperty("https", is_https)
+ systemProperty("user", user)
+ systemProperty("password", password)
doFirst {
// Tell the test JVM if the cluster JVM is running under a debugger so that tests can
@@ -223,7 +397,18 @@ integTest {
testClusters.integTest {
testDistribution = "ARCHIVE"
+
+ // Optionally install security
+ if (System.getProperty("security.enabled") != null) {
+ configureSecurityPlugin(testClusters.integTest)
+ }
+
plugin(project.tasks.bundlePlugin.archiveFile)
+ if (Os.isFamily(Os.FAMILY_WINDOWS)) {
+ // Add the paths of built JNI libraries and its dependent libraries to PATH variable in System variables
+ environment('PATH', System.getenv('PATH') + ";$rootDir/jni/release" + ";$rootDir/src/main/resources/windowsDependencies")
+ }
+
// Cluster shrink exception thrown if we try to set numberOfNodes to 1, so only apply if > 1
if (_numNodes > 1) numberOfNodes = _numNodes
// When running integration tests it doesn't forward the --debug-jvm to the cluster anymore
@@ -270,3 +455,33 @@ run {
}
}
}
+
+// updateVersion: Task to auto increment to the next development iteration
+task updateVersion {
+ onlyIf { System.getProperty('newVersion') }
+ doLast {
+ ext.newVersion = System.getProperty('newVersion')
+ println "Setting version to ${newVersion}."
+ // String tokenization to support -SNAPSHOT
+ // Include the required files that needs to be updated with new Version
+ ant.replaceregexp(match: opensearch_version.tokenize('-')[0], replace: newVersion.tokenize('-')[0], flags:'g', byline:true) {
+ fileset(dir: projectDir) {
+ // Include the required files that needs to be updated with new Version
+ include(name: ".github/workflows/backwards_compatibility_tests_workflow.yml")
+ }
+ }
+ ant.replaceregexp(file:'build.gradle', match: '"opensearch.version", "\\d.*"', replace: '"opensearch.version", "' + newVersion.tokenize('-')[0] + '-SNAPSHOT"', flags:'g', byline:true)
+
+ ext.os_version_without_snapshot = opensearch_version.tokenize('-')[0]
+ ext.os_version_major = os_version_without_snapshot.tokenize('.')[0]
+ ext.os_version_minor = os_version_without_snapshot.tokenize('.')[1]
+ ext.os_version_patch = os_version_without_snapshot.tokenize('.')[2]
+ // This condition will check if the BWC workflow is already updated or not and will run next steps if not updated
+ if (!fileTree(".github/workflows/backwards_compatibility_tests_workflow.yml").getSingleFile().text.contains(os_version_without_snapshot)) {
+ // Extract the oldBWCVersion from the existing OpenSearch Version (oldBWCVersion = major . (minor-1) . patch)
+ ext.oldBWCVersion = os_version_major + '.' + Integer.toString(Integer.valueOf(os_version_minor) - 1) + '.' + os_version_patch
+ // Include the current OpenSearch Version before version bump to the bwc_version matrix
+ ant.replaceregexp(file:".github/workflows/backwards_compatibility_tests_workflow.yml", match: oldBWCVersion, replace: oldBWCVersion + '", "' + opensearch_version.tokenize('-')[0], flags:'g', byline:true)
+ }
+ }
+}
diff --git a/gradle.properties b/gradle.properties
index 162d212e5..c36e5cca0 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -4,7 +4,7 @@
#
version=1.0.0
-systemProp.bwc.version=1.3.2
+systemProp.bwc.version=1.3.4
org.gradle.jvmargs=--add-exports jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED \
--add-exports jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED \
diff --git a/gradle/wrapper/gradle-wrapper.jar b/gradle/wrapper/gradle-wrapper.jar
index 7454180f2..7f93135c4 100644
Binary files a/gradle/wrapper/gradle-wrapper.jar and b/gradle/wrapper/gradle-wrapper.jar differ
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 97060749f..d54311bfa 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,12 +1,7 @@
-#
-# Copyright OpenSearch Contributors
-# SPDX-License-Identifier: Apache-2.0
-#
-
distributionBase=GRADLE_USER_HOME
distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.2-all.zip
+distributionSha256Sum=f2b9ed0faf8472cbe469255ae6c86eddb77076c75191741b4a462f33128dd419
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.4-all.zip
+networkTimeout=10000
zipStoreBase=GRADLE_USER_HOME
zipStorePath=wrapper/dists
-distributionSha256Sum=e6d864e3b5bc05cc62041842b306383fc1fefcec359e70cebb1d470a6094ca82
-
diff --git a/gradlew b/gradlew
index 8250acdd7..1aa94a426 100755
--- a/gradlew
+++ b/gradlew
@@ -1,17 +1,7 @@
-#!/usr/bin/env sh
+#!/bin/sh
#
-# SPDX-License-Identifier: Apache-2.0
-#
-# The OpenSearch Contributors require contributions made to
-# this file be licensed under the Apache-2.0 license or a
-# compatible open source license.
-#
-# Modifications Copyright OpenSearch Contributors. See
-# GitHub history for details.
-#
-#
-# Copyright 2015 the original author or authors.
+# Copyright © 2015-2021 the original authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -27,78 +17,111 @@
#
##############################################################################
-##
-## Gradle start up script for UN*X
-##
+#
+# Gradle start up script for POSIX generated by Gradle.
+#
+# Important for running:
+#
+# (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is
+# noncompliant, but you have some other compliant shell such as ksh or
+# bash, then to run this script, type that shell name before the whole
+# command line, like:
+#
+# ksh Gradle
+#
+# Busybox and similar reduced shells will NOT work, because this script
+# requires all of these POSIX shell features:
+# * functions;
+# * expansions «$var», «${var}», «${var:-default}», «${var+SET}»,
+# «${var#prefix}», «${var%suffix}», and «$( cmd )»;
+# * compound commands having a testable exit status, especially «case»;
+# * various built-in commands including «command», «set», and «ulimit».
+#
+# Important for patching:
+#
+# (2) This script targets any POSIX shell, so it avoids extensions provided
+# by Bash, Ksh, etc; in particular arrays are avoided.
+#
+# The "traditional" practice of packing multiple parameters into a
+# space-separated string is a well documented source of bugs and security
+# problems, so this is (mostly) avoided, by progressively accumulating
+# options in "$@", and eventually passing that to Java.
+#
+# Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS,
+# and GRADLE_OPTS) rely on word-splitting, this is performed explicitly;
+# see the in-line comments for details.
+#
+# There are tweaks for specific operating systems such as AIX, CygWin,
+# Darwin, MinGW, and NonStop.
+#
+# (3) This script is generated from the Groovy template
+# https://github.com/gradle/gradle/blob/HEAD/subprojects/plugins/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt
+# within the Gradle project.
+#
+# You can find Gradle at https://github.com/gradle/gradle/.
+#
##############################################################################
# Attempt to set APP_HOME
+
# Resolve links: $0 may be a link
-PRG="$0"
-# Need this for relative symlinks.
-while [ -h "$PRG" ] ; do
- ls=`ls -ld "$PRG"`
- link=`expr "$ls" : '.*-> \(.*\)$'`
- if expr "$link" : '/.*' > /dev/null; then
- PRG="$link"
- else
- PRG=`dirname "$PRG"`"/$link"
- fi
+app_path=$0
+
+# Need this for daisy-chained symlinks.
+while
+ APP_HOME=${app_path%"${app_path##*/}"} # leaves a trailing /; empty if no leading path
+ [ -h "$app_path" ]
+do
+ ls=$( ls -ld "$app_path" )
+ link=${ls#*' -> '}
+ case $link in #(
+ /*) app_path=$link ;; #(
+ *) app_path=$APP_HOME$link ;;
+ esac
done
-SAVED="`pwd`"
-cd "`dirname \"$PRG\"`/" >/dev/null
-APP_HOME="`pwd -P`"
-cd "$SAVED" >/dev/null
-
-APP_NAME="Gradle"
-APP_BASE_NAME=`basename "$0"`
-# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
-DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"'
+# This is normally unused
+# shellcheck disable=SC2034
+APP_BASE_NAME=${0##*/}
+# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036)
+APP_HOME=$( cd "${APP_HOME:-./}" > /dev/null && pwd -P ) || exit
# Use the maximum available, or set MAX_FD != -1 to use that value.
-MAX_FD="maximum"
+MAX_FD=maximum
warn () {
echo "$*"
-}
+} >&2
die () {
echo
echo "$*"
echo
exit 1
-}
+} >&2
# OS specific support (must be 'true' or 'false').
cygwin=false
msys=false
darwin=false
nonstop=false
-case "`uname`" in
- CYGWIN* )
- cygwin=true
- ;;
- Darwin* )
- darwin=true
- ;;
- MINGW* )
- msys=true
- ;;
- NONSTOP* )
- nonstop=true
- ;;
+case "$( uname )" in #(
+ CYGWIN* ) cygwin=true ;; #(
+ Darwin* ) darwin=true ;; #(
+ MSYS* | MINGW* ) msys=true ;; #(
+ NONSTOP* ) nonstop=true ;;
esac
CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar
+
# Determine the Java command to use to start the JVM.
if [ -n "$JAVA_HOME" ] ; then
if [ -x "$JAVA_HOME/jre/sh/java" ] ; then
# IBM's JDK on AIX uses strange locations for the executables
- JAVACMD="$JAVA_HOME/jre/sh/java"
+ JAVACMD=$JAVA_HOME/jre/sh/java
else
- JAVACMD="$JAVA_HOME/bin/java"
+ JAVACMD=$JAVA_HOME/bin/java
fi
if [ ! -x "$JAVACMD" ] ; then
die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME
@@ -107,87 +130,120 @@ Please set the JAVA_HOME variable in your environment to match the
location of your Java installation."
fi
else
- JAVACMD="java"
- which java >/dev/null 2>&1 || die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+ JAVACMD=java
+ if ! command -v java >/dev/null 2>&1
+ then
+ die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
Please set the JAVA_HOME variable in your environment to match the
location of your Java installation."
+ fi
fi
# Increase the maximum file descriptors if we can.
-if [ "$cygwin" = "false" -a "$darwin" = "false" -a "$nonstop" = "false" ] ; then
- MAX_FD_LIMIT=`ulimit -H -n`
- if [ $? -eq 0 ] ; then
- if [ "$MAX_FD" = "maximum" -o "$MAX_FD" = "max" ] ; then
- MAX_FD="$MAX_FD_LIMIT"
- fi
- ulimit -n $MAX_FD
- if [ $? -ne 0 ] ; then
- warn "Could not set maximum file descriptor limit: $MAX_FD"
- fi
- else
- warn "Could not query maximum file descriptor limit: $MAX_FD_LIMIT"
- fi
+if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then
+ case $MAX_FD in #(
+ max*)
+ # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked.
+ # shellcheck disable=SC2039,SC3045
+ MAX_FD=$( ulimit -H -n ) ||
+ warn "Could not query maximum file descriptor limit"
+ esac
+ case $MAX_FD in #(
+ '' | soft) :;; #(
+ *)
+ # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked.
+ # shellcheck disable=SC2039,SC3045
+ ulimit -n "$MAX_FD" ||
+ warn "Could not set maximum file descriptor limit to $MAX_FD"
+ esac
fi
-# For Darwin, add options to specify how the application appears in the dock
-if $darwin; then
- GRADLE_OPTS="$GRADLE_OPTS \"-Xdock:name=$APP_NAME\" \"-Xdock:icon=$APP_HOME/media/gradle.icns\""
-fi
+# Collect all arguments for the java command, stacking in reverse order:
+# * args from the command line
+# * the main class name
+# * -classpath
+# * -D...appname settings
+# * --module-path (only if needed)
+# * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables.
# For Cygwin or MSYS, switch paths to Windows format before running java
-if [ "$cygwin" = "true" -o "$msys" = "true" ] ; then
- APP_HOME=`cygpath --path --mixed "$APP_HOME"`
- CLASSPATH=`cygpath --path --mixed "$CLASSPATH"`
- JAVACMD=`cygpath --unix "$JAVACMD"`
-
- # We build the pattern for arguments to be converted via cygpath
- ROOTDIRSRAW=`find -L / -maxdepth 1 -mindepth 1 -type d 2>/dev/null`
- SEP=""
- for dir in $ROOTDIRSRAW ; do
- ROOTDIRS="$ROOTDIRS$SEP$dir"
- SEP="|"
- done
- OURCYGPATTERN="(^($ROOTDIRS))"
- # Add a user-defined pattern to the cygpath arguments
- if [ "$GRADLE_CYGPATTERN" != "" ] ; then
- OURCYGPATTERN="$OURCYGPATTERN|($GRADLE_CYGPATTERN)"
- fi
+if "$cygwin" || "$msys" ; then
+ APP_HOME=$( cygpath --path --mixed "$APP_HOME" )
+ CLASSPATH=$( cygpath --path --mixed "$CLASSPATH" )
+
+ JAVACMD=$( cygpath --unix "$JAVACMD" )
+
# Now convert the arguments - kludge to limit ourselves to /bin/sh
- i=0
- for arg in "$@" ; do
- CHECK=`echo "$arg"|egrep -c "$OURCYGPATTERN" -`
- CHECK2=`echo "$arg"|egrep -c "^-"` ### Determine if an option
-
- if [ $CHECK -ne 0 ] && [ $CHECK2 -eq 0 ] ; then ### Added a condition
- eval `echo args$i`=`cygpath --path --ignore --mixed "$arg"`
- else
- eval `echo args$i`="\"$arg\""
+ for arg do
+ if
+ case $arg in #(
+ -*) false ;; # don't mess with options #(
+ /?*) t=${arg#/} t=/${t%%/*} # looks like a POSIX filepath
+ [ -e "$t" ] ;; #(
+ *) false ;;
+ esac
+ then
+ arg=$( cygpath --path --ignore --mixed "$arg" )
fi
- i=`expr $i + 1`
+ # Roll the args list around exactly as many times as the number of
+ # args, so each arg winds up back in the position where it started, but
+ # possibly modified.
+ #
+ # NB: a `for` loop captures its iteration list before it begins, so
+ # changing the positional parameters here affects neither the number of
+ # iterations, nor the values presented in `arg`.
+ shift # remove old arg
+ set -- "$@" "$arg" # push replacement arg
done
- case $i in
- 0) set -- ;;
- 1) set -- "$args0" ;;
- 2) set -- "$args0" "$args1" ;;
- 3) set -- "$args0" "$args1" "$args2" ;;
- 4) set -- "$args0" "$args1" "$args2" "$args3" ;;
- 5) set -- "$args0" "$args1" "$args2" "$args3" "$args4" ;;
- 6) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" ;;
- 7) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" ;;
- 8) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" ;;
- 9) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" "$args8" ;;
- esac
fi
-# Escape application args
-save () {
- for i do printf %s\\n "$i" | sed "s/'/'\\\\''/g;1s/^/'/;\$s/\$/' \\\\/" ; done
- echo " "
-}
-APP_ARGS=`save "$@"`
-# Collect all arguments for the java command, following the shell quoting and substitution rules
-eval set -- $DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS "\"-Dorg.gradle.appname=$APP_BASE_NAME\"" -classpath "\"$CLASSPATH\"" org.gradle.wrapper.GradleWrapperMain "$APP_ARGS"
+# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"'
+
+# Collect all arguments for the java command:
+# * DEFAULT_JVM_OPTS, JAVA_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments,
+# and any embedded shellness will be escaped.
+# * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be
+# treated as '${Hostname}' itself on the command line.
+
+set -- \
+ "-Dorg.gradle.appname=$APP_BASE_NAME" \
+ -classpath "$CLASSPATH" \
+ org.gradle.wrapper.GradleWrapperMain \
+ "$@"
+
+# Stop when "xargs" is not available.
+if ! command -v xargs >/dev/null 2>&1
+then
+ die "xargs is not available"
+fi
+
+# Use "xargs" to parse quoted args.
+#
+# With -n1 it outputs one arg per line, with the quotes and backslashes removed.
+#
+# In Bash we could simply go:
+#
+# readarray ARGS < <( xargs -n1 <<<"$var" ) &&
+# set -- "${ARGS[@]}" "$@"
+#
+# but POSIX shell has neither arrays nor command substitution, so instead we
+# post-process each arg (as a line of input to sed) to backslash-escape any
+# character that might be a shell metacharacter, then use eval to reverse
+# that process (while maintaining the separation between arguments), and wrap
+# the whole thing up as a single "set" statement.
+#
+# This will of course break if any of these variables contains a newline or
+# an unmatched quote.
+#
+
+eval "set -- $(
+ printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" |
+ xargs -n1 |
+ sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' |
+ tr '\n' ' '
+ )" '"$@"'
exec "$JAVACMD" "$@"
diff --git a/gradlew.bat b/gradlew.bat
index 9109989e3..6689b85be 100644
--- a/gradlew.bat
+++ b/gradlew.bat
@@ -14,7 +14,7 @@
@rem limitations under the License.
@rem
-@if "%DEBUG%" == "" @echo off
+@if "%DEBUG%"=="" @echo off
@rem ##########################################################################
@rem
@rem Gradle startup script for Windows
@@ -25,7 +25,8 @@
if "%OS%"=="Windows_NT" setlocal
set DIRNAME=%~dp0
-if "%DIRNAME%" == "" set DIRNAME=.
+if "%DIRNAME%"=="" set DIRNAME=.
+@rem This is normally unused
set APP_BASE_NAME=%~n0
set APP_HOME=%DIRNAME%
@@ -40,7 +41,7 @@ if defined JAVA_HOME goto findJavaFromJavaHome
set JAVA_EXE=java.exe
%JAVA_EXE% -version >NUL 2>&1
-if "%ERRORLEVEL%" == "0" goto init
+if %ERRORLEVEL% equ 0 goto execute
echo.
echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
@@ -54,7 +55,7 @@ goto fail
set JAVA_HOME=%JAVA_HOME:"=%
set JAVA_EXE=%JAVA_HOME%/bin/java.exe
-if exist "%JAVA_EXE%" goto init
+if exist "%JAVA_EXE%" goto execute
echo.
echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
@@ -64,38 +65,26 @@ echo location of your Java installation.
goto fail
-:init
-@rem Get command-line arguments, handling Windows variants
-
-if not "%OS%" == "Windows_NT" goto win9xME_args
-
-:win9xME_args
-@rem Slurp the command line arguments.
-set CMD_LINE_ARGS=
-set _SKIP=2
-
-:win9xME_args_slurp
-if "x%~1" == "x" goto execute
-
-set CMD_LINE_ARGS=%*
-
:execute
@rem Setup the command line
set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
@rem Execute Gradle
-"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %*
:end
@rem End local scope for the variables with windows NT shell
-if "%ERRORLEVEL%"=="0" goto mainEnd
+if %ERRORLEVEL% equ 0 goto mainEnd
:fail
rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
rem the _cmd.exe /c_ return code!
-if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
-exit /b 1
+set EXIT_CODE=%ERRORLEVEL%
+if %EXIT_CODE% equ 0 set EXIT_CODE=1
+if not ""=="%GRADLE_EXIT_CONSOLE%" exit %EXIT_CODE%
+exit /b %EXIT_CODE%
:mainEnd
if "%OS%"=="Windows_NT" endlocal
diff --git a/jni/CMakeLists.txt b/jni/CMakeLists.txt
index c6b2f5f22..6253eed15 100644
--- a/jni/CMakeLists.txt
+++ b/jni/CMakeLists.txt
@@ -3,20 +3,23 @@
# SPDX-License-Identifier: Apache-2.0
#
-cmake_minimum_required(VERSION 3.17)
+cmake_minimum_required(VERSION 3.24.0)
project(KNNPlugin_JNI)
+include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/macros.cmake)
+
# ---------------------------------- SETUP ----------------------------------
# Target libraries to be compiled
-set(TARGET_LIB_COMMON opensearchknn_common) # Shared library with common utilities
+# Shared library with common utilities. Not a JNI library. Other JNI libs should depend on this one.
+set(TARGET_LIB_UTIL opensearchknn_util)
+set(TARGET_LIB_COMMON opensearchknn_common) # common lib for JNI
set(TARGET_LIB_NMSLIB opensearchknn_nmslib) # nmslib JNI
set(TARGET_LIB_FAISS opensearchknn_faiss) # faiss JNI
set(TARGET_LIBS "") # Libs to be installed
-set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED True)
-
option(CONFIG_FAISS "Configure faiss library build when this is on")
option(CONFIG_NMSLIB "Configure nmslib library build when this is on")
option(CONFIG_TEST "Configure tests when this is on")
@@ -27,6 +30,18 @@ else()
set(CONFIG_ALL OFF)
endif ()
+# `git am` will create commits from the patches in the native libraries. This is ideal for development envs
+# because it prevents full lib rebuild everytime cmake is run. However, for build systems that will run the
+# build workflow once, it can cause issues because git commits require that the user and the user's email be set.
+# See https://github.com/opensearch-project/k-NN/issues/1651. So, we provide a flag that allows users to select between
+# the two
+if(NOT DEFINED COMMIT_LIB_PATCHES OR "${COMMIT_LIB_PATCHES}" STREQUAL true)
+ set(GIT_PATCH_COMMAND am)
+else()
+ set(GIT_PATCH_COMMAND apply)
+endif()
+message(STATUS "Using the following git patch command: \"${GIT_PATCH_COMMAND}\"")
+
# Set OS specific variables
if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
set(CMAKE_MACOSX_RPATH 1)
@@ -35,6 +50,14 @@ if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux)
set(JVM_OS_TYPE linux)
set(LIB_EXT .so)
+elseif(${CMAKE_SYSTEM_NAME} STREQUAL Windows)
+# Set the CXX_COMPILER_VERSION, CMAKE_CXX_FLAGS, JVM_OS_TYPE, prefix and extension for the target libraries that are built.
+ set(CXX_COMPILER_VERSION ${CMAKE_CXX_COMPILER_VERSION})
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fpermissive")
+ set(JVM_OS_TYPE win32)
+ set(LIB_EXT .dll)
+ set(CMAKE_SHARED_LIBRARY_PREFIX "")
+ set(CMAKE_STATIC_LIBRARY_PREFIX "")
else()
message(FATAL_ERROR "Unable to run on system: ${CMAKE_SYSTEM_NAME}")
endif()
@@ -52,35 +75,28 @@ elseif(${CMAKE_SYSTEM_PROCESSOR} STREQUAL x86_64)
endif()
# ----------------------------------------------------------------------------
+# ---------------------------------- UTIL ----------------------------------
+add_library(${TARGET_LIB_UTIL} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/jni_util.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/commons.cpp)
+target_include_directories(${TARGET_LIB_UTIL} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE})
+opensearch_set_common_properties(${TARGET_LIB_UTIL})
+list(APPEND TARGET_LIBS ${TARGET_LIB_UTIL})
+# ----------------------------------------------------------------------------
+
# ---------------------------------- COMMON ----------------------------------
-add_library(${TARGET_LIB_COMMON} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/jni_util.cpp)
+add_library(${TARGET_LIB_COMMON} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_JNICommons.cpp)
+target_link_libraries(${TARGET_LIB_COMMON} ${TARGET_LIB_UTIL})
target_include_directories(${TARGET_LIB_COMMON} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE})
-set_target_properties(${TARGET_LIB_COMMON} PROPERTIES SUFFIX ${LIB_EXT})
-set_target_properties(${TARGET_LIB_COMMON} PROPERTIES POSITION_INDEPENDENT_CODE ON)
-set_target_properties(${TARGET_LIB_COMMON} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
+opensearch_set_common_properties(${TARGET_LIB_COMMON})
list(APPEND TARGET_LIBS ${TARGET_LIB_COMMON})
# ----------------------------------------------------------------------------
# ---------------------------------- NMSLIB ----------------------------------
if (${CONFIG_NMSLIB} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
- # Check if nmslib exists
- find_path(NMS_REPO_DIR NAMES similarity_search PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib)
-
- # If not, pull the updated submodule
- if (NOT EXISTS ${NMS_REPO_DIR})
- message(STATUS "Could not find nmslib. Pulling updated submodule.")
- execute_process(COMMAND git submodule update --init -- external/nmslib WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
- endif ()
-
- add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search)
-
+ include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/init-nmslib.cmake)
add_library(${TARGET_LIB_NMSLIB} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_NmslibService.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/nmslib_wrapper.cpp)
- target_link_libraries(${TARGET_LIB_NMSLIB} NonMetricSpaceLib ${TARGET_LIB_COMMON})
+ target_link_libraries(${TARGET_LIB_NMSLIB} NonMetricSpaceLib ${TARGET_LIB_UTIL})
target_include_directories(${TARGET_LIB_NMSLIB} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include)
- set_target_properties(${TARGET_LIB_NMSLIB} PROPERTIES SUFFIX ${LIB_EXT})
- set_target_properties(${TARGET_LIB_NMSLIB} PROPERTIES POSITION_INDEPENDENT_CODE ON)
- set_target_properties(${TARGET_LIB_NMSLIB} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
-
+ opensearch_set_common_properties(${TARGET_LIB_NMSLIB})
list(APPEND TARGET_LIBS ${TARGET_LIB_NMSLIB})
endif ()
@@ -88,150 +104,87 @@ endif ()
# ---------------------------------- FAISS ----------------------------------
if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
- set(BUILD_TESTING OFF) # Avoid building faiss tests
- set(BLA_STATIC ON) # Statically link BLAS
- set(FAISS_OPT_LEVEL generic) # Keep optimization level generic
-
- if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
- if(CMAKE_C_COMPILER_ID MATCHES "Clang\$")
- set(OpenMP_C_FLAGS "-Xpreprocessor -fopenmp")
- set(OpenMP_C_LIB_NAMES "omp")
- set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
- endif()
-
- if(CMAKE_CXX_COMPILER_ID MATCHES "Clang\$")
- set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp -I/usr/local/opt/libomp/include")
- set(OpenMP_CXX_LIB_NAMES "omp")
- set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
- endif()
- endif()
-
+ include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/init-faiss.cmake)
find_package(OpenMP REQUIRED)
- find_package(ZLIB REQUIRED)
- find_package(BLAS REQUIRED)
- enable_language(Fortran)
- find_package(LAPACK REQUIRED)
-
- # Check if faiss exists
- find_path(FAISS_REPO_DIR NAMES faiss PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss)
-
- # If not, pull the updated submodule
- if (NOT EXISTS ${FAISS_REPO_DIR})
- message(STATUS "Could not find faiss. Pulling updated submodule.")
- execute_process(COMMAND git submodule update --init -- external/faiss WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
- endif ()
-
- set(FAISS_ENABLE_GPU OFF)
- set(FAISS_ENABLE_PYTHON OFF)
- add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/faiss EXCLUDE_FROM_ALL)
-
- add_library(${TARGET_LIB_FAISS} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_FaissService.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_wrapper.cpp)
- target_link_libraries(${TARGET_LIB_FAISS} faiss ${TARGET_LIB_COMMON} OpenMP::OpenMP_CXX)
- target_include_directories(${TARGET_LIB_FAISS} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss)
- set_target_properties(${TARGET_LIB_FAISS} PROPERTIES SUFFIX ${LIB_EXT})
- set_target_properties(${TARGET_LIB_FAISS} PROPERTIES POSITION_INDEPENDENT_CODE ON)
- set_target_properties(${TARGET_LIB_FAISS} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
-
+ add_library(
+ ${TARGET_LIB_FAISS} SHARED
+ ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_FaissService.cpp
+ ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_wrapper.cpp
+ ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_util.cpp
+ ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_index_service.cpp
+ ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_methods.cpp
+ )
+ target_link_libraries(${TARGET_LIB_FAISS} ${TARGET_LINK_FAISS_LIB} ${TARGET_LIB_UTIL} OpenMP::OpenMP_CXX)
+ target_include_directories(${TARGET_LIB_FAISS} PRIVATE
+ ${CMAKE_CURRENT_SOURCE_DIR}/include
+ $ENV{JAVA_HOME}/include
+ $ENV{JAVA_HOME}/include/${JVM_OS_TYPE}
+ ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss
+ )
+ opensearch_set_common_properties(${TARGET_LIB_FAISS})
list(APPEND TARGET_LIBS ${TARGET_LIB_FAISS})
endif ()
# ---------------------------------------------------------------------------
# --------------------------------- TESTS -----------------------------------
-if (${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
- # Reference - https://crascit.com/2015/07/25/cmake-gtest/
- configure_file(CMakeLists.txt.in googletest-download/CMakeLists.txt)
- execute_process(COMMAND "${CMAKE_COMMAND}" -G "${CMAKE_GENERATOR}" .
- WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
- )
- execute_process(COMMAND "${CMAKE_COMMAND}" --build .
- WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
- )
- set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
-
- add_subdirectory("${CMAKE_BINARY_DIR}/googletest-src"
- "${CMAKE_BINARY_DIR}/googletest-build" EXCLUDE_FROM_ALL
- )
- add_executable(
- jni_test
- tests/faiss_wrapper_test.cpp
- tests/nmslib_wrapper_test.cpp
- tests/test_util.cpp)
-
- target_link_libraries(
- jni_test
- gtest_main
- gmock_main
- faiss
- NonMetricSpaceLib
- OpenMP::OpenMP_CXX
- ${TARGET_LIB_FAISS}
- ${TARGET_LIB_NMSLIB}
- ${TARGET_LIB_COMMON}
- )
-
- target_include_directories(jni_test PRIVATE
- ${CMAKE_CURRENT_SOURCE_DIR}/tests
- ${CMAKE_CURRENT_SOURCE_DIR}/include
- $ENV{JAVA_HOME}/include
- $ENV{JAVA_HOME}/include/${JVM_OS_TYPE}
- ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss
- ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include
- ${gtest_SOURCE_DIR}/include
- ${gmock_SOURCE_DIR}/include)
-
-
- set_target_properties(jni_test PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/bin)
-endif ()
-# ---------------------------------------------------------------------------
+# Windows : Comment the TESTS for now because the tests are failing(failing to build jni_tests.exe) if we are building our target libraries as SHARED libraries.
+# TODO: Fix the failing JNI TESTS on Windows
+if ("${WIN32}" STREQUAL "")
+ if (${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
+ # Reference - https://crascit.com/2015/07/25/cmake-gtest/
+ configure_file(CMakeLists.txt.in googletest-download/CMakeLists.txt)
+ execute_process(COMMAND "${CMAKE_COMMAND}" -G "${CMAKE_GENERATOR}" .
+ WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
+ )
+ execute_process(COMMAND "${CMAKE_COMMAND}" --build .
+ WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
+ )
+ set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
+
+ add_subdirectory("${CMAKE_BINARY_DIR}/googletest-src"
+ "${CMAKE_BINARY_DIR}/googletest-build" EXCLUDE_FROM_ALL
+ )
+ add_executable(
+ jni_test
+ tests/faiss_wrapper_test.cpp
+ tests/faiss_wrapper_unit_test.cpp
+ tests/faiss_util_test.cpp
+ tests/nmslib_wrapper_test.cpp
+ tests/nmslib_wrapper_unit_test.cpp
+ tests/test_util.cpp
+ tests/commons_test.cpp
+ tests/faiss_stream_support_test.cpp
+ tests/faiss_index_service_test.cpp
+ tests/nmslib_stream_support_test.cpp
+ )
+
+ target_link_libraries(
+ jni_test
+ gtest_main
+ gmock_main
+ faiss
+ NonMetricSpaceLib
+ OpenMP::OpenMP_CXX
+ ${TARGET_LIB_FAISS}
+ ${TARGET_LIB_NMSLIB}
+ ${TARGET_LIB_COMMON}
+ ${TARGET_LIB_UTIL}
+ )
+
+ target_include_directories(jni_test PRIVATE
+ ${CMAKE_CURRENT_SOURCE_DIR}/tests
+ ${CMAKE_CURRENT_SOURCE_DIR}/include
+ $ENV{JAVA_HOME}/include
+ $ENV{JAVA_HOME}/include/${JVM_OS_TYPE}
+ ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss
+ ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include
+ ${gtest_SOURCE_DIR}/include
+ ${gmock_SOURCE_DIR}/include)
+
+
+ set_target_properties(jni_test PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/bin)
+ endif ()
+endif()
-# -------------------------------- INSTALL ----------------------------------
-# Installation rules for shared library
-install(TARGETS ${TARGET_LIBS}
- LIBRARY DESTINATION lib
- COMPONENT library)
-
-set(KNN_MAINTAINER "OpenSearch Team ")
-set(OPENSEARCH_DOWNLOAD_URL "https://opensearch.org/downloads.html")
-set(CPACK_PACKAGE_NAME ${KNN_PACKAGE_NAME})
-set(CPACK_PACKAGE_VERSION ${KNN_PLUGIN_VERSION})
-set(CMAKE_INSTALL_PREFIX /usr)
-set(CPACK_GENERATOR "RPM;DEB")
-set(CPACK_OUTPUT_FILE_PREFIX packages)
-set(CPACK_PACKAGE_RELEASE 1)
-set(CPACK_PACKAGE_VENDOR "Amazon")
-set(CPACK_PACKAGE_CONTACT "Maintainer: ${KNN_MAINTAINER}")
-set(CPACK_PACKAGING_INSTALL_PREFIX ${CMAKE_INSTALL_PREFIX})
-set(CPACK_COMPONENTS_GROUPING IGNORE)
-get_cmake_property(CPACK_COMPONENTS_ALL COMPONENTS)
-list(REMOVE_ITEM CPACK_COMPONENTS_ALL "Unspecified")
-
-# Component variable
-set(KNN_PACKAGE_NAME opensearch-knnlib)
-set(KNN_PACKAGE_DESCRIPTION "KNN JNI libraries built off of nmslib and faiss for OpenSearch")
-
-# RPM
-set(CPACK_RPM_PACKAGE_LICENSE "ASL-2.0")
-set(CPACK_RPM_COMPONENT_INSTALL ON)
-set(CPACK_RPM_PACKAGE_URL ${OPENSEARCH_DOWNLOAD_URL})
-set(CPACK_RPM_PACKAGE_RELEASE ${CPACK_PACKAGE_RELEASE})
-
-set(CPACK_RPM_PACKAGE_NAME ${KNN_PACKAGE_NAME})
-set(CPACK_RPM_FILE_NAME "${CPACK_RPM_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${JVM_OS_TYPE}-${MACH_ARCH}.rpm")
-set(CPACK_RPM_PACKAGE_DESCRIPTION ${KNN_PACKAGE_DESCRIPTION})
-set(CPACK_RPM_PACKAGE_SUMMARY "OpenSearch k-NN JNI Library with nmslib and faiss")
-
-# DEB
-set(CPACK_DEBIAN_PACKAGE_HOMEPAGE ${OPENSEARCH_DOWNLOAD_URL})
-set(CPACK_DEBIAN_PACKAGE_MAINTAINER ${KNN_MAINTAINER})
-set(CPACK_DEBIAN_PACKAGE_VERSION ${CPACK_PACKAGE_VERSION})
-set(CPACK_DEBIAN_PACKAGE_SECTION "libs")
-set(CPACK_DEB_COMPONENT_INSTALL ON)
-
-set(CPACK_DEBIAN_PACKAGE_NAME ${KNN_PACKAGE_NAME})
-set(CPACK_DEBIAN_FILE_NAME "${CPACK_DEBIAN_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${JVM_OS_TYPE}-${MACH_ARCH}.deb")
-set(CPACK_DEBIAN_DESCRIPTION ${KNN_PACKAGE_DESCRIPTION})
-set(CPACK_DEBIAN_PACKAGE_SOURCE ${CPACK_DEBIAN_PACKAGE_NAME})
-
-include(CPack)
# ---------------------------------------------------------------------------
diff --git a/jni/cmake/init-faiss.cmake b/jni/cmake/init-faiss.cmake
new file mode 100644
index 000000000..4492d9f45
--- /dev/null
+++ b/jni/cmake/init-faiss.cmake
@@ -0,0 +1,105 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+
+# Check if faiss exists
+find_path(FAISS_REPO_DIR NAMES faiss PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss NO_DEFAULT_PATH)
+
+# If not, pull the updated submodule
+if (NOT EXISTS ${FAISS_REPO_DIR})
+ message(STATUS "Could not find faiss. Pulling updated submodule.")
+ execute_process(COMMAND git submodule update --init -- external/faiss WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+endif ()
+
+# Apply patches
+if(NOT DEFINED APPLY_LIB_PATCHES OR "${APPLY_LIB_PATCHES}" STREQUAL true)
+ # Define list of patch files
+ set(PATCH_FILE_LIST)
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0002-Enable-precomp-table-to-be-shared-ivfpq.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0003-Custom-patch-to-support-range-search-params.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0004-Custom-patch-to-support-binary-vector.patch")
+
+ # Get patch id of the last commit
+ execute_process(COMMAND sh -c "git --no-pager show HEAD | git patch-id --stable" OUTPUT_VARIABLE PATCH_ID_OUTPUT_FROM_COMMIT WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss)
+ string(REPLACE " " ";" PATCH_ID_LIST_FROM_COMMIT ${PATCH_ID_OUTPUT_FROM_COMMIT})
+ list(GET PATCH_ID_LIST_FROM_COMMIT 0 PATCH_ID_FROM_COMMIT)
+
+ # Find all patch files need to apply
+ list(SORT PATCH_FILE_LIST ORDER DESCENDING)
+ set(PATCH_FILES_TO_APPLY)
+ foreach(PATCH_FILE IN LISTS PATCH_FILE_LIST)
+ # Get patch id of a patch file
+ execute_process(COMMAND sh -c "cat ${PATCH_FILE} | git patch-id --stable" OUTPUT_VARIABLE PATCH_ID_OUTPUT)
+ string(REPLACE " " ";" PATCH_ID_LIST ${PATCH_ID_OUTPUT})
+ list(GET PATCH_ID_LIST 0 PATCH_ID)
+
+ # Add the file to patch list if patch id does not match
+ if (${PATCH_ID} STREQUAL ${PATCH_ID_FROM_COMMIT})
+ break()
+ else()
+ list(APPEND PATCH_FILES_TO_APPLY ${PATCH_FILE})
+ endif()
+ endforeach()
+
+ # Apply patch files
+ list(SORT PATCH_FILES_TO_APPLY)
+ foreach(PATCH_FILE IN LISTS PATCH_FILES_TO_APPLY)
+ message(STATUS "Applying patch of ${PATCH_FILE}")
+ execute_process(COMMAND git ${GIT_PATCH_COMMAND} --3way --ignore-space-change --ignore-whitespace ${PATCH_FILE} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+ if(RESULT_CODE)
+ message(FATAL_ERROR "Failed to apply patch:\n${ERROR_MSG}")
+ endif()
+ endforeach()
+endif()
+
+if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
+ if(CMAKE_C_COMPILER_ID MATCHES "Clang\$")
+ set(OpenMP_C_FLAGS "-Xpreprocessor -fopenmp")
+ set(OpenMP_C_LIB_NAMES "omp")
+ set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
+ endif()
+
+ if(CMAKE_CXX_COMPILER_ID MATCHES "Clang\$")
+ set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp -I/usr/local/opt/libomp/include")
+ set(OpenMP_CXX_LIB_NAMES "omp")
+ set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
+ endif()
+endif()
+
+find_package(ZLIB REQUIRED)
+
+# Statically link BLAS - ensure this is before we find the blas package so we dont dynamically link
+set(BLA_STATIC ON)
+find_package(BLAS REQUIRED)
+enable_language(Fortran)
+find_package(LAPACK REQUIRED)
+
+# Set relevant properties
+set(BUILD_TESTING OFF) # Avoid building faiss tests
+set(FAISS_ENABLE_GPU OFF)
+set(FAISS_ENABLE_PYTHON OFF)
+
+if(NOT DEFINED AVX2_ENABLED)
+ set(AVX2_ENABLED true) # set default value as true if the argument is not set
+endif()
+
+if(NOT DEFINED AVX512_ENABLED)
+ set(AVX512_ENABLED true) # set default value as true if the argument is not set
+endif()
+
+if(${CMAKE_SYSTEM_NAME} STREQUAL Windows OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64" OR ( NOT AVX2_ENABLED AND NOT AVX512_ENABLED))
+ set(FAISS_OPT_LEVEL generic) # Keep optimization level as generic on Windows OS as it is not supported due to MINGW64 compiler issue. Also, on aarch64 avx2 is not supported.
+ set(TARGET_LINK_FAISS_LIB faiss)
+elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux AND AVX512_ENABLED)
+ set(FAISS_OPT_LEVEL avx512) # Keep optimization level as avx512 to improve performance on Linux. This is not present on mac systems, and presently not supported on Windows OS.
+ set(TARGET_LINK_FAISS_LIB faiss_avx512)
+ string(PREPEND LIB_EXT "_avx512") # Prepend "_avx512" to lib extension to create the library as "libopensearchknn_faiss_avx512.so" on linux
+else()
+ set(FAISS_OPT_LEVEL avx2) # Keep optimization level as avx2 to improve performance on Linux and Mac.
+ set(TARGET_LINK_FAISS_LIB faiss_avx2)
+ string(PREPEND LIB_EXT "_avx2") # Prepend "_avx2" to lib extension to create the library as "libopensearchknn_faiss_avx2.so" on linux and "libopensearchknn_faiss_avx2.jnilib" on mac
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/faiss EXCLUDE_FROM_ALL)
diff --git a/jni/cmake/init-nmslib.cmake b/jni/cmake/init-nmslib.cmake
new file mode 100644
index 000000000..a7c3f7d93
--- /dev/null
+++ b/jni/cmake/init-nmslib.cmake
@@ -0,0 +1,57 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+
+# Check if nmslib exists
+find_path(NMS_REPO_DIR NAMES similarity_search PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib NO_DEFAULT_PATH)
+
+# If not, pull the updated submodule
+if (NOT EXISTS ${NMS_REPO_DIR})
+ message(STATUS "Could not find nmslib. Pulling updated submodule.")
+ execute_process(COMMAND git submodule update --init -- external/nmslib WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+endif ()
+
+# Apply patches
+if(NOT DEFINED APPLY_LIB_PATCHES OR "${APPLY_LIB_PATCHES}" STREQUAL true)
+ # Define list of patch files
+ set(PATCH_FILE_LIST)
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0002-Adds-ability-to-pass-ef-parameter-in-the-query-for-h.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0003-Added-streaming-apis-for-vector-index-loading-in-Hnsw.patch")
+ list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0004-Added-a-new-save-apis-in-Hnsw-with-streaming-interfa.patch")
+
+ # Get patch id of the last commit
+ execute_process(COMMAND sh -c "git --no-pager show HEAD | git patch-id --stable" OUTPUT_VARIABLE PATCH_ID_OUTPUT_FROM_COMMIT WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib)
+ string(REPLACE " " ";" PATCH_ID_LIST_FROM_COMMIT ${PATCH_ID_OUTPUT_FROM_COMMIT})
+ list(GET PATCH_ID_LIST_FROM_COMMIT 0 PATCH_ID_FROM_COMMIT)
+
+ # Find all patch files need to apply
+ list(SORT PATCH_FILE_LIST ORDER DESCENDING)
+ set(PATCH_FILES_TO_APPLY)
+ foreach(PATCH_FILE IN LISTS PATCH_FILE_LIST)
+ # Get patch id of a patch file
+ execute_process(COMMAND sh -c "cat ${PATCH_FILE} | git patch-id --stable" OUTPUT_VARIABLE PATCH_ID_OUTPUT)
+ string(REPLACE " " ";" PATCH_ID_LIST ${PATCH_ID_OUTPUT})
+ list(GET PATCH_ID_LIST 0 PATCH_ID)
+
+ # Add the file to patch list if patch id does not match
+ if (${PATCH_ID} STREQUAL ${PATCH_ID_FROM_COMMIT})
+ break()
+ else()
+ list(APPEND PATCH_FILES_TO_APPLY ${PATCH_FILE})
+ endif()
+ endforeach()
+endif()
+
+# Apply patch files
+list(SORT PATCH_FILES_TO_APPLY)
+foreach(PATCH_FILE IN LISTS PATCH_FILES_TO_APPLY)
+ message(STATUS "Applying patch of ${PATCH_FILE}")
+ execute_process(COMMAND git ${GIT_PATCH_COMMAND} --3way --ignore-space-change --ignore-whitespace ${PATCH_FILE} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+ if(RESULT_CODE)
+ message(FATAL_ERROR "Failed to apply patch:\n${ERROR_MSG}")
+ endif()
+endforeach()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search)
diff --git a/jni/cmake/macros.cmake b/jni/cmake/macros.cmake
new file mode 100644
index 000000000..773033b7e
--- /dev/null
+++ b/jni/cmake/macros.cmake
@@ -0,0 +1,16 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+
+macro(opensearch_set_common_properties TARGET)
+ set_target_properties(${TARGET} PROPERTIES SUFFIX ${LIB_EXT})
+ set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+ if (NOT "${WIN32}" STREQUAL "")
+ # Use RUNTIME_OUTPUT_DIRECTORY, to build the target library in the specified directory at runtime.
+ set_target_properties(${TARGET} PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
+ else()
+ set_target_properties(${TARGET} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
+ endif()
+endmacro()
diff --git a/jni/external/faiss b/jni/external/faiss
index 88eabe97f..1f42e815d 160000
--- a/jni/external/faiss
+++ b/jni/external/faiss
@@ -1 +1 @@
-Subproject commit 88eabe97f96d0c0964dfa075f74373c64d46da80
+Subproject commit 1f42e815db7754297e3b4467763352b829b6cde0
diff --git a/jni/include/commons.h b/jni/include/commons.h
new file mode 100644
index 000000000..38b00cc5d
--- /dev/null
+++ b/jni/include/commons.h
@@ -0,0 +1,107 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef OPENSEARCH_KNN_COMMONS_H
+#define OPENSEARCH_KNN_COMMONS_H
+
+#include "jni_util.h"
+#include
+namespace knn_jni {
+ namespace commons {
+ /**
+ * This is utility function that can be used to store data in native memory. This function will allocate memory for
+ * the data(rows*columns) with initialCapacity and return the memory address where the data is stored.
+ * If you are using this function for first time use memoryAddress = 0 to ensure that a new memory location is created.
+ * For subsequent calls you can pass the same memoryAddress. If the data cannot be stored in the memory location
+ * will throw Exception.
+ *
+ * append tells the method to keep appending to the existing vector. Passing the value as false will clear the vector
+ * without reallocating new memory. This helps with reducing memory frangmentation and overhead of allocating
+ * and deallocating when the memory address needs to be reused.
+ *
+ * CAUTION: The behavior is undefined if the memory address is deallocated and the method is called
+ *
+ * @param memoryAddress The address of the memory location where data will be stored.
+ * @param data 2D float array containing data to be stored in native memory.
+ * @param initialCapacity The initial capacity of the memory location.
+ * @param append whether to append or start from index 0 when called subsequently with the same address
+ * @return memory address of std::vector where the data is stored.
+ */
+ jlong storeVectorData(knn_jni::JNIUtilInterface *, JNIEnv *, jlong , jobjectArray, jlong, jboolean);
+
+ /**
+ * This is utility function that can be used to store data in native memory. This function will allocate memory for
+ * the data(rows*columns) with initialCapacity and return the memory address where the data is stored.
+ * If you are using this function for first time use memoryAddress = 0 to ensure that a new memory location is created.
+ * For subsequent calls you can pass the same memoryAddress. If the data cannot be stored in the memory location
+ * will throw Exception.
+ *
+ * append tells the method to keep appending to the existing vector. Passing the value as false will clear the vector
+ * without reallocating new memory. This helps with reducing memory frangmentation and overhead of allocating
+ * and deallocating when the memory address needs to be reused.
+ *
+ * CAUTION: The behavior is undefined if the memory address is deallocated and the method is called
+ *
+ * @param memoryAddress The address of the memory location where data will be stored.
+ * @param data 2D byte array containing binary data to be stored in native memory.
+ * @param initialCapacity The initial capacity of the memory location.
+ * @param append whether to append or start from index 0 when called subsequently with the same address
+ * @return memory address of std::vector where the data is stored.
+ */
+ jlong storeBinaryVectorData(knn_jni::JNIUtilInterface *, JNIEnv *, jlong , jobjectArray, jlong, jboolean);
+
+ /**
+ * This is utility function that can be used to store signed int8 data in native memory. This function will allocate memory for
+ * the data(rows*columns) with initialCapacity and return the memory address where the data is stored.
+ * If you are using this function for first time use memoryAddress = 0 to ensure that a new memory location is created.
+ * For subsequent calls you can pass the same memoryAddress. If the data cannot be stored in the memory location
+ * will throw Exception.
+ *
+ * @param memoryAddress The address of the memory location where data will be stored.
+ * @param data 2D byte array containing int8 data to be stored in native memory.
+ * @param initialCapacity The initial capacity of the memory location.
+ * @param append whether to append or start from index 0 when called subsequently with the same address
+ * @return memory address of std::vector where the data is stored.
+ */
+ jlong storeByteVectorData(knn_jni::JNIUtilInterface *, JNIEnv *, jlong , jobjectArray, jlong, jboolean);
+
+ /**
+ * Free up the memory allocated for the data stored in memory address. This function should be used with the memory
+ * address returned by {@link JNICommons#storeVectorData(long, float[][], long, long)}
+ *
+ * @param memoryAddress address to be freed.
+ */
+ void freeVectorData(jlong);
+
+ /**
+ * Free up the memory allocated for the data stored in memory address. This function should be used with the memory
+ * address returned by {@link JNICommons#storeByteVectorData(long, byte[][], long, long)}
+ *
+ * @param memoryAddress address to be freed.
+ */
+ void freeByteVectorData(jlong);
+
+ /**
+ * Free up the memory allocated for the data stored in memory address. This function should be used with the memory
+ * address returned by {@link JNICommons#storeBinaryVectorData(long, byte[][], long, long)}
+ *
+ * @param memoryAddress address to be freed.
+ */
+ void freeBinaryVectorData(jlong);
+
+ /**
+ * Extracts query time efSearch from method parameters
+ **/
+ int getIntegerMethodParameter(JNIEnv *, knn_jni::JNIUtilInterface *, std::unordered_map, std::string, int);
+ }
+}
+
+#endif
\ No newline at end of file
diff --git a/jni/include/faiss_index_service.h b/jni/include/faiss_index_service.h
new file mode 100644
index 000000000..d96c3e755
--- /dev/null
+++ b/jni/include/faiss_index_service.h
@@ -0,0 +1,202 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// The OpenSearch Contributors require contributions made to
+// this file be licensed under the Apache-2.0 license or a
+// compatible open source license.
+//
+// Modifications Copyright OpenSearch Contributors. See
+// GitHub history for details.
+
+/**
+ * This file contains classes for index operations which are free of JNI
+ */
+
+#ifndef OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H
+#define OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H
+
+#include
+#include "faiss/MetricType.h"
+#include "faiss/impl/io.h"
+#include "jni_util.h"
+#include "faiss_methods.h"
+#include "faiss_stream_support.h"
+#include
+
+namespace knn_jni {
+namespace faiss_wrapper {
+
+
+/**
+ * A class to provide operations on index
+ * This class should evolve to have only cpp object but not jni object
+ */
+class IndexService {
+public:
+ explicit IndexService(std::unique_ptr faissMethods);
+
+ /**
+ * Initialize index
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param dim dimension of vectors
+ * @param numVectors number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param parameters parameters to be applied to faiss index
+ * @return memory address of the native index object
+ */
+ virtual jlong initIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, faiss::MetricType metric, std::string indexDescription, int dim, int numVectors, int threadCount, std::unordered_map parameters);
+
+ /**
+ * Add vectors to index
+ *
+ * @param dim dimension of vectors
+ * @param numIds number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param vectorsAddress memory address which is holding vector data
+ * @param idMapAddress memory address of the native index object
+ */
+ virtual void insertToIndex(int dim, int numIds, int threadCount, int64_t vectorsAddress, std::vector &ids, jlong idMapAddress);
+
+ /**
+ * Write index to disk
+ *
+ * @param writer IOWriter implementation doing IO processing.
+ * In most cases, it is expected to have underlying Lucene's IndexOuptut.
+ * @param idMapAddress memory address of the native index object
+ */
+ virtual void writeIndex(faiss::IOWriter* writer, jlong idMapAddress);
+
+ virtual ~IndexService() = default;
+
+protected:
+ virtual void allocIndex(faiss::Index * index, size_t dim, size_t numVectors);
+
+ std::unique_ptr faissMethods;
+}; // class IndexService
+
+/**
+ * A class to provide operations on index
+ * This class should evolve to have only cpp object but not jni object
+ */
+class BinaryIndexService final : public IndexService {
+public:
+ //TODO Remove dependency on JNIUtilInterface and JNIEnv
+ //TODO Reduce the number of parameters
+ explicit BinaryIndexService(std::unique_ptr faissMethods);
+
+ /**
+ * Initialize index
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param dim dimension of vectors
+ * @param numVectors number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param parameters parameters to be applied to faiss index
+ * @return memory address of the native index object
+ */
+ jlong initIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, faiss::MetricType metric, std::string indexDescription, int dim, int numVectors, int threadCount, std::unordered_map parameters) final;
+
+ /**
+ * Add vectors to index
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param dim dimension of vectors
+ * @param numIds number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param vectorsAddress memory address which is holding vector data
+ * @param idMap a map of document id and vector id
+ * @param parameters parameters to be applied to faiss index
+ */
+ void insertToIndex(int dim, int numIds, int threadCount, int64_t vectorsAddress, std::vector &ids, jlong idMapAddress) final;
+
+ /**
+ * Write index to disk
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param threadCount number of thread count to be used while adding data
+ * @param indexPath path to write index
+ * @param idMap a map of document id and vector id
+ * @param parameters parameters to be applied to faiss index
+ */
+ void writeIndex(faiss::IOWriter* writer, jlong idMapAddress) final;
+
+protected:
+ void allocIndex(faiss::Index * index, size_t dim, size_t numVectors) final;
+}; // class BinaryIndexService
+
+/**
+ * A class to provide operations on index
+ * This class should evolve to have only cpp object but not jni object
+ */
+class ByteIndexService final : public IndexService {
+public:
+ //TODO Remove dependency on JNIUtilInterface and JNIEnv
+ //TODO Reduce the number of parameters
+ explicit ByteIndexService(std::unique_ptr faissMethods);
+
+ /**
+ * Initialize index
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param dim dimension of vectors
+ * @param numVectors number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param parameters parameters to be applied to faiss index
+ * @return memory address of the native index object
+ */
+ jlong initIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, faiss::MetricType metric, std::string indexDescription, int dim, int numVectors, int threadCount, std::unordered_map parameters) final;
+
+ /**
+ * Add vectors to index
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param dim dimension of vectors
+ * @param numIds number of vectors
+ * @param threadCount number of thread count to be used while adding data
+ * @param vectorsAddress memory address which is holding vector data
+ * @param idMap a map of document id and vector id
+ * @param parameters parameters to be applied to faiss index
+ */
+ void insertToIndex(int dim, int numIds, int threadCount, int64_t vectorsAddress, std::vector &ids, jlong idMapAddress) final;
+
+ /**
+ * Write index to disk
+ *
+ * @param jniUtil jni util
+ * @param env jni environment
+ * @param metric space type for distance calculation
+ * @param indexDescription index description to be used by faiss index factory
+ * @param threadCount number of thread count to be used while adding data
+ * @param indexPath path to write index
+ * @param idMap a map of document id and vector id
+ * @param parameters parameters to be applied to faiss index
+ */
+ void writeIndex(faiss::IOWriter* writer, jlong idMapAddress) final;
+
+ protected:
+ void allocIndex(faiss::Index * index, size_t dim, size_t numVectors) final;
+}; // class ByteIndexService
+
+}
+}
+
+
+#endif //OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H
diff --git a/jni/include/faiss_methods.h b/jni/include/faiss_methods.h
new file mode 100644
index 000000000..d8f14d03f
--- /dev/null
+++ b/jni/include/faiss_methods.h
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// The OpenSearch Contributors require contributions made to
+// this file be licensed under the Apache-2.0 license or a
+// compatible open source license.
+//
+// Modifications Copyright OpenSearch Contributors. See
+// GitHub history for details.
+
+#ifndef OPENSEARCH_KNN_FAISS_METHODS_H
+#define OPENSEARCH_KNN_FAISS_METHODS_H
+
+#include "faiss/impl/io.h"
+#include "faiss/Index.h"
+#include "faiss/IndexBinary.h"
+#include "faiss/IndexIDMap.h"
+#include "faiss/index_io.h"
+
+namespace knn_jni {
+namespace faiss_wrapper {
+
+/**
+ * A class having wrapped faiss methods
+ *
+ * This class helps to mock faiss methods during unit test
+ */
+class FaissMethods {
+public:
+ FaissMethods() = default;
+
+ virtual faiss::Index* indexFactory(int d, const char* description, faiss::MetricType metric);
+
+ virtual faiss::IndexBinary* indexBinaryFactory(int d, const char* description);
+
+ virtual faiss::IndexIDMapTemplate* indexIdMap(faiss::Index* index);
+
+ virtual faiss::IndexIDMapTemplate* indexBinaryIdMap(faiss::IndexBinary* index);
+
+ virtual void writeIndex(const faiss::Index* idx, faiss::IOWriter* writer);
+
+ virtual void writeIndexBinary(const faiss::IndexBinary* idx, faiss::IOWriter* writer);
+
+ virtual ~FaissMethods() = default;
+}; // class FaissMethods
+
+} //namespace faiss_wrapper
+} //namespace knn_jni
+
+
+#endif //OPENSEARCH_KNN_FAISS_METHODS_H
\ No newline at end of file
diff --git a/jni/include/faiss_stream_support.h b/jni/include/faiss_stream_support.h
new file mode 100644
index 000000000..eb1b2a404
--- /dev/null
+++ b/jni/include/faiss_stream_support.h
@@ -0,0 +1,98 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef OPENSEARCH_KNN_JNI_FAISS_STREAM_SUPPORT_H
+#define OPENSEARCH_KNN_JNI_FAISS_STREAM_SUPPORT_H
+
+#include "faiss/impl/io.h"
+#include "jni_util.h"
+#include "native_engines_stream_support.h"
+#include "parameter_utils.h"
+
+#include
+#include
+#include
+#include
+
+namespace knn_jni {
+namespace stream {
+
+
+
+/**
+ * A glue component inheriting IOReader to be passed down to Faiss library.
+ * This will then indirectly call the mediator component and eventually read required bytes from Lucene's IndexInput.
+ */
+class FaissOpenSearchIOReader final : public faiss::IOReader {
+ public:
+ explicit FaissOpenSearchIOReader(NativeEngineIndexInputMediator *_mediator)
+ : faiss::IOReader(),
+ mediator(knn_jni::util::ParameterCheck::require_non_null(_mediator, "mediator")) {
+ name = "FaissOpenSearchIOReader";
+ }
+
+ size_t operator()(void *ptr, size_t size, size_t nitems) final {
+ const auto readBytes = size * nitems;
+ if (readBytes > 0) {
+ // Mediator calls IndexInput, then copy read bytes to `ptr`.
+ mediator->copyBytes(readBytes, (uint8_t *) ptr);
+ }
+ return nitems;
+ }
+
+ int filedescriptor() final {
+ throw std::runtime_error("filedescriptor() is not supported in FaissOpenSearchIOReader.");
+ }
+
+ private:
+ NativeEngineIndexInputMediator *mediator;
+}; // class FaissOpenSearchIOReader
+
+
+/**
+ * A glue component inheriting IOWriter to delegate IO processing down to the given
+ * mediator. The mediator is expected to do write bytes via the provided Lucene's IndexOutput.
+ */
+class FaissOpenSearchIOWriter final : public faiss::IOWriter {
+ public:
+ explicit FaissOpenSearchIOWriter(NativeEngineIndexOutputMediator *_mediator)
+ : faiss::IOWriter(),
+ mediator(knn_jni::util::ParameterCheck::require_non_null(_mediator, "mediator")) {
+ name = "FaissOpenSearchIOWriter";
+ }
+
+ size_t operator()(const void *ptr, size_t size, size_t nitems) final {
+ const auto writeBytes = size * nitems;
+ if (writeBytes > 0) {
+ mediator->writeBytes(reinterpret_cast(ptr), writeBytes);
+ }
+ return nitems;
+ }
+
+ // return a file number that can be memory-mapped
+ int filedescriptor() final {
+ throw std::runtime_error("filedescriptor() is not supported in FaissOpenSearchIOWriter.");
+ }
+
+ void flush() {
+ mediator->flush();
+ }
+
+ private:
+ NativeEngineIndexOutputMediator *mediator;
+}; // class FaissOpenSearchIOWriter
+
+
+
+}
+}
+
+#endif //OPENSEARCH_KNN_JNI_FAISS_STREAM_SUPPORT_H
diff --git a/jni/include/faiss_util.h b/jni/include/faiss_util.h
new file mode 100644
index 000000000..f23540aef
--- /dev/null
+++ b/jni/include/faiss_util.h
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// The OpenSearch Contributors require contributions made to
+// this file be licensed under the Apache-2.0 license or a
+// compatible open source license.
+//
+// Modifications Copyright OpenSearch Contributors. See
+// GitHub history for details.
+
+/**
+ * This file contains util methods which are free of JNI to be used in faiss_wrapper.cpp
+ */
+
+#ifndef OPENSEARCH_KNN_FAISS_UTIL_H
+#define OPENSEARCH_KNN_FAISS_UTIL_H
+
+#include "faiss/impl/IDGrouper.h"
+#include
+
+namespace faiss_util {
+ std::unique_ptr buildIDGrouperBitmap(int *parentIdsArray, int parentIdsLength, std::vector* bitmap);
+};
+
+
+#endif //OPENSEARCH_KNN_FAISS_UTIL_H
diff --git a/jni/include/faiss_wrapper.h b/jni/include/faiss_wrapper.h
index 6c8a86143..e48e6faa9 100644
--- a/jni/include/faiss_wrapper.h
+++ b/jni/include/faiss_wrapper.h
@@ -13,35 +13,102 @@
#define OPENSEARCH_KNN_FAISS_WRAPPER_H
#include "jni_util.h"
-
+#include "faiss_index_service.h"
+#include "faiss_stream_support.h"
#include
namespace knn_jni {
namespace faiss_wrapper {
- // Create an index with ids and vectors. The configuration is defined by values in the Java map, parametersJ.
- // The index is serialized to indexPathJ.
- void CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ, jobjectArray vectorsJ,
- jstring indexPathJ, jobject parametersJ);
+ jlong InitIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jlong numDocs, jint dimJ, jobject parametersJ, IndexService *indexService);
+
+ void InsertToIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jintArray idsJ, jlong vectorsAddressJ, jint dimJ, jlong indexAddr, jint threadCount, IndexService *indexService);
+
+ void WriteIndex(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jobject output, jlong indexAddr, IndexService *indexService);
// Create an index with ids and vectors. Instead of creating a new index, this function creates the index
// based off of the template index passed in. The index is serialized to indexPathJ.
void CreateIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,
- jobjectArray vectorsJ, jstring indexPathJ, jbyteArray templateIndexJ,
+ jlong vectorsAddressJ, jint dimJ, jobject output, jbyteArray templateIndexJ,
jobject parametersJ);
+ // Create an index with ids and vectors. Instead of creating a new index, this function creates the index
+ // based off of the template index passed in. The index is serialized to indexPathJ.
+ void CreateBinaryIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,
+ jlong vectorsAddressJ, jint dimJ, jobject output, jbyteArray templateIndexJ,
+ jobject parametersJ);
+
+ // Create a index with ids and byte vectors. Instead of creating a new index, this function creates the index
+ // based off of the template index passed in. The index is serialized to indexPathJ.
+ void CreateByteIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,
+ jlong vectorsAddressJ, jint dimJ, jobject output, jbyteArray templateIndexJ,
+ jobject parametersJ);
+
// Load an index from indexPathJ into memory.
//
// Return a pointer to the loaded index
jlong LoadIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jstring indexPathJ);
- // Execute a query against the index located in memory at indexPointerJ.
+ // Loads an index with a reader implemented IOReader
//
- // Return an array of KNNQueryResults
+ // Returns a pointer of the loaded index
+ jlong LoadIndexWithStream(faiss::IOReader* ioReader);
+
+ // Load a binary index from indexPathJ into memory.
+ //
+ // Return a pointer to the loaded index
+ jlong LoadBinaryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jstring indexPathJ);
+
+ // Loads a binary index with a reader implemented IOReader
+ //
+ // Returns a pointer of the loaded index
+ jlong LoadBinaryIndexWithStream(faiss::IOReader* ioReader);
+
+ // Check if a loaded index requires shared state
+ bool IsSharedIndexStateRequired(jlong indexPointerJ);
+
+ // Initializes the shared index state from an index. Note, this will not set the state for
+ // the index pointed to by indexPointerJ. To set it, SetSharedIndexState needs to be called.
+ //
+ // Return a pointer to the shared index state
+ jlong InitSharedIndexState(jlong indexPointerJ);
+
+ // Sets the sharedIndexState for an index
+ void SetSharedIndexState(jlong indexPointerJ, jlong shareIndexStatePointerJ);
+
+ /**
+ * Execute a query against the index located in memory at indexPointerJ
+ *
+ * Parameters:
+ * methodParamsJ: introduces a map to have additional method parameters
+ *
+ * Return an array of KNNQueryResults
+ */
jobjectArray QueryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
- jfloatArray queryVectorJ, jint kJ);
+ jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ, jintArray parentIdsJ);
+
+ /**
+ * Execute a query against the index located in memory at indexPointerJ along with Filters
+ *
+ * Parameters:
+ * methodParamsJ: introduces a map to have additional method parameters
+ *
+ * Return an array of KNNQueryResults
+ */
+ jobjectArray QueryIndex_WithFilter(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
+ jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ, jlongArray filterIdsJ,
+ jint filterIdsTypeJ, jintArray parentIdsJ);
+
+ // Execute a query against the binary index located in memory at indexPointerJ along with Filters
+ //
+ // Return an array of KNNQueryResults
+ jobjectArray QueryBinaryIndex_WithFilter(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
+ jbyteArray queryVectorJ, jint kJ, jobject methodParamsJ, jlongArray filterIdsJ, jint filterIdsTypeJ, jintArray parentIdsJ);
// Free the index located in memory at indexPointerJ
- void Free(jlong indexPointer);
+ void Free(jlong indexPointer, jboolean isBinaryIndexJ);
+
+ // Free shared index state in memory at shareIndexStatePointerJ
+ void FreeSharedIndexState(jlong shareIndexStatePointerJ);
// Perform initilization operations for the library
void InitLibrary();
@@ -52,6 +119,52 @@ namespace knn_jni {
// Return the serialized representation
jbyteArray TrainIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jobject parametersJ, jint dimension,
jlong trainVectorsPointerJ);
+
+ // Create an empty binary index defined by the values in the Java map, parametersJ. Train the index with
+ // the vector of floats located at trainVectorsPointerJ.
+ //
+ // Return the serialized representation
+ jbyteArray TrainBinaryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jobject parametersJ, jint dimension,
+ jlong trainVectorsPointerJ);
+
+ // Create an empty byte index defined by the values in the Java map, parametersJ. Train the index with
+ // the byte vectors located at trainVectorsPointerJ.
+ //
+ // Return the serialized representation
+ jbyteArray TrainByteIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jobject parametersJ, jint dimension,
+ jlong trainVectorsPointerJ);
+
+ /*
+ * Perform a range search with filter against the index located in memory at indexPointerJ.
+ *
+ * @param indexPointerJ - pointer to the index
+ * @param queryVectorJ - the query vector
+ * @param radiusJ - the radius for the range search
+ * @param methodParamsJ - the method parameters
+ * @param maxResultsWindowJ - the maximum number of results to return
+ * @param filterIdsJ - the filter ids
+ * @param filterIdsTypeJ - the filter ids type
+ * @param parentIdsJ - the parent ids
+ *
+ * @return an array of RangeQueryResults
+ */
+ jobjectArray RangeSearchWithFilter(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jlong indexPointerJ, jfloatArray queryVectorJ,
+ jfloat radiusJ, jobject methodParamsJ, jint maxResultWindowJ, jlongArray filterIdsJ, jint filterIdsTypeJ, jintArray parentIdsJ);
+
+ /*
+ * Perform a range search against the index located in memory at indexPointerJ.
+ *
+ * @param indexPointerJ - pointer to the index
+ * @param queryVectorJ - the query vector
+ * @param radiusJ - the radius for the range search
+ * @param methodParamsJ - the method parameters
+ * @param maxResultsWindowJ - the maximum number of results to return
+ * @param parentIdsJ - the parent ids
+ *
+ * @return an array of RangeQueryResults
+ */
+ jobjectArray RangeSearch(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jlong indexPointerJ, jfloatArray queryVectorJ,
+ jfloat radiusJ, jobject methodParamsJ, jint maxResultWindowJ, jintArray parentIdsJ);
}
}
diff --git a/jni/include/jni_util.h b/jni/include/jni_util.h
index 8d6f2b6f6..45068ae1b 100644
--- a/jni/include/jni_util.h
+++ b/jni/include/jni_util.h
@@ -17,12 +17,13 @@
#include
#include
#include
+#include
+#include
namespace knn_jni {
// Interface for making calls to JNI
- class JNIUtilInterface {
- public:
+ struct JNIUtilInterface {
// -------------------------- EXCEPTION HANDLING ----------------------------
// Takes the name of a Java exception type and a message and throws the corresponding exception
// to the JVM
@@ -33,7 +34,7 @@ namespace knn_jni {
virtual void HasExceptionInStack(JNIEnv* env) = 0;
// HasExceptionInStack with ability to specify message
- virtual void HasExceptionInStack(JNIEnv* env, const std::string& message) = 0;
+ virtual void HasExceptionInStack(JNIEnv* env, const char *message) = 0;
// Catches a C++ exception and throws the corresponding exception to the JVM
virtual void CatchCppExceptionAndThrowJava(JNIEnv* env) = 0;
@@ -68,6 +69,13 @@ namespace knn_jni {
virtual std::vector Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ,
int dim) = 0;
+ virtual void Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ,
+ int dim, std::vector *vect ) = 0;
+ virtual void Convert2dJavaObjectArrayAndStoreToBinaryVector(JNIEnv *env, jobjectArray array2dJ,
+ int dim, std::vector *vect ) = 0;
+ virtual void Convert2dJavaObjectArrayAndStoreToByteVector(JNIEnv *env, jobjectArray array2dJ,
+ int dim, std::vector *vect ) = 0;
+
virtual std::vector ConvertJavaIntArrayToCppIntVector(JNIEnv *env, jintArray arrayJ) = 0;
// --------------------------------------------------------------------------
@@ -75,10 +83,14 @@ namespace knn_jni {
// ------------------------------ MISC HELPERS ------------------------------
virtual int GetInnerDimensionOf2dJavaFloatArray(JNIEnv *env, jobjectArray array2dJ) = 0;
+ virtual int GetInnerDimensionOf2dJavaByteArray(JNIEnv *env, jobjectArray array2dJ) = 0;
+
virtual int GetJavaObjectArrayLength(JNIEnv *env, jobjectArray arrayJ) = 0;
virtual int GetJavaIntArrayLength(JNIEnv *env, jintArray arrayJ) = 0;
+ virtual int GetJavaLongArrayLength(JNIEnv *env, jlongArray arrayJ) = 0;
+
virtual int GetJavaBytesArrayLength(JNIEnv *env, jbyteArray arrayJ) = 0;
virtual int GetJavaFloatArrayLength(JNIEnv *env, jfloatArray arrayJ) = 0;
@@ -93,6 +105,8 @@ namespace knn_jni {
virtual jint * GetIntArrayElements(JNIEnv *env, jintArray array, jboolean * isCopy) = 0;
+ virtual jlong * GetLongArrayElements(JNIEnv *env, jlongArray array, jboolean * isCopy) = 0;
+
virtual jobject GetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index) = 0;
virtual jobject NewObject(JNIEnv *env, jclass clazz, jmethodID methodId, int id, float distance) = 0;
@@ -107,58 +121,121 @@ namespace knn_jni {
virtual void ReleaseIntArrayElements(JNIEnv *env, jintArray array, jint *elems, jint mode) = 0;
+ virtual void ReleaseLongArrayElements(JNIEnv *env, jlongArray array, jlong *elems, jint mode) = 0;
+
virtual void SetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index, jobject val) = 0;
virtual void SetByteArrayRegion(JNIEnv *env, jbyteArray array, jsize start, jsize len, const jbyte * buf) = 0;
+ virtual jobject GetObjectField(JNIEnv * env, jobject obj, jfieldID fieldID) = 0;
+
+ virtual jclass FindClassFromJNIEnv(JNIEnv * env, const char *name) = 0;
+
+ virtual jmethodID GetMethodID(JNIEnv * env, jclass clazz, const char *name, const char *sig) = 0;
+
+ virtual jfieldID GetFieldID(JNIEnv * env, jclass clazz, const char *name, const char *sig) = 0;
+
+ virtual void * GetPrimitiveArrayCritical(JNIEnv * env, jarray array, jboolean *isCopy) = 0;
+
+ virtual void ReleasePrimitiveArrayCritical(JNIEnv * env, jarray array, void *carray, jint mode) = 0;
+
+ virtual jint CallNonvirtualIntMethodA(JNIEnv *env, jobject obj, jclass clazz,
+ jmethodID methodID, jvalue *args) = 0;
+
+ virtual jlong CallNonvirtualLongMethodA(JNIEnv * env, jobject obj, jclass clazz,
+ jmethodID methodID, jvalue* args) = 0;
+
+ virtual void CallNonvirtualVoidMethodA(JNIEnv * env, jobject obj, jclass clazz,
+ jmethodID methodID, jvalue* args) = 0;
+
// --------------------------------------------------------------------------
};
jobject GetJObjectFromMapOrThrow(std::unordered_map map, std::string key);
// Class that implements JNIUtilInterface methods
- class JNIUtil: public JNIUtilInterface {
+ class JNIUtil final : public JNIUtilInterface {
public:
// Initialize and Uninitialize methods are used for caching/cleaning up Java classes and methods
void Initialize(JNIEnv* env);
void Uninitialize(JNIEnv* env);
- void ThrowJavaException(JNIEnv* env, const char* type = "", const char* message = "");
- void HasExceptionInStack(JNIEnv* env);
- void HasExceptionInStack(JNIEnv* env, const std::string& message);
- void CatchCppExceptionAndThrowJava(JNIEnv* env);
- jclass FindClass(JNIEnv * env, const std::string& className);
- jmethodID FindMethod(JNIEnv * env, const std::string& className, const std::string& methodName);
- std::string ConvertJavaStringToCppString(JNIEnv * env, jstring javaString);
- std::unordered_map ConvertJavaMapToCppMap(JNIEnv *env, jobject parametersJ);
- std::string ConvertJavaObjectToCppString(JNIEnv *env, jobject objectJ);
- int ConvertJavaObjectToCppInteger(JNIEnv *env, jobject objectJ);
- std::vector Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim);
- std::vector ConvertJavaIntArrayToCppIntVector(JNIEnv *env, jintArray arrayJ);
- int GetInnerDimensionOf2dJavaFloatArray(JNIEnv *env, jobjectArray array2dJ);
- int GetJavaObjectArrayLength(JNIEnv *env, jobjectArray arrayJ);
- int GetJavaIntArrayLength(JNIEnv *env, jintArray arrayJ);
- int GetJavaBytesArrayLength(JNIEnv *env, jbyteArray arrayJ);
- int GetJavaFloatArrayLength(JNIEnv *env, jfloatArray arrayJ);
-
- void DeleteLocalRef(JNIEnv *env, jobject obj);
- jbyte * GetByteArrayElements(JNIEnv *env, jbyteArray array, jboolean * isCopy);
- jfloat * GetFloatArrayElements(JNIEnv *env, jfloatArray array, jboolean * isCopy);
- jint * GetIntArrayElements(JNIEnv *env, jintArray array, jboolean * isCopy);
- jobject GetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index);
- jobject NewObject(JNIEnv *env, jclass clazz, jmethodID methodId, int id, float distance);
- jobjectArray NewObjectArray(JNIEnv *env, jsize len, jclass clazz, jobject init);
- jbyteArray NewByteArray(JNIEnv *env, jsize len);
- void ReleaseByteArrayElements(JNIEnv *env, jbyteArray array, jbyte *elems, int mode);
- void ReleaseFloatArrayElements(JNIEnv *env, jfloatArray array, jfloat *elems, int mode);
- void ReleaseIntArrayElements(JNIEnv *env, jintArray array, jint *elems, jint mode);
- void SetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index, jobject val);
- void SetByteArrayRegion(JNIEnv *env, jbyteArray array, jsize start, jsize len, const jbyte * buf);
+ void ThrowJavaException(JNIEnv* env, const char* type = "", const char* message = "") final;
+ void HasExceptionInStack(JNIEnv* env) final;
+ void HasExceptionInStack(JNIEnv* env, const char* message) final;
+ void CatchCppExceptionAndThrowJava(JNIEnv* env) final;
+ jclass FindClass(JNIEnv * env, const std::string& className) final;
+ jmethodID FindMethod(JNIEnv * env, const std::string& className, const std::string& methodName) final;
+ std::string ConvertJavaStringToCppString(JNIEnv * env, jstring javaString) final;
+ std::unordered_map ConvertJavaMapToCppMap(JNIEnv *env, jobject parametersJ) final;
+ std::string ConvertJavaObjectToCppString(JNIEnv *env, jobject objectJ) final;
+ int ConvertJavaObjectToCppInteger(JNIEnv *env, jobject objectJ) final;
+ std::vector Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim) final;
+ std::vector ConvertJavaIntArrayToCppIntVector(JNIEnv *env, jintArray arrayJ) final;
+ int GetInnerDimensionOf2dJavaFloatArray(JNIEnv *env, jobjectArray array2dJ) final;
+ int GetInnerDimensionOf2dJavaByteArray(JNIEnv *env, jobjectArray array2dJ) final;
+ int GetJavaObjectArrayLength(JNIEnv *env, jobjectArray arrayJ) final;
+ int GetJavaIntArrayLength(JNIEnv *env, jintArray arrayJ) final;
+ int GetJavaLongArrayLength(JNIEnv *env, jlongArray arrayJ) final;
+ int GetJavaBytesArrayLength(JNIEnv *env, jbyteArray arrayJ) final;
+ int GetJavaFloatArrayLength(JNIEnv *env, jfloatArray arrayJ) final;
+
+ void DeleteLocalRef(JNIEnv *env, jobject obj) final;
+ jbyte * GetByteArrayElements(JNIEnv *env, jbyteArray array, jboolean * isCopy) final;
+ jfloat * GetFloatArrayElements(JNIEnv *env, jfloatArray array, jboolean * isCopy) final;
+ jint * GetIntArrayElements(JNIEnv *env, jintArray array, jboolean * isCopy) final;
+ jlong * GetLongArrayElements(JNIEnv *env, jlongArray array, jboolean * isCopy) final;
+ jobject GetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index) final;
+ jobject NewObject(JNIEnv *env, jclass clazz, jmethodID methodId, int id, float distance) final;
+ jobjectArray NewObjectArray(JNIEnv *env, jsize len, jclass clazz, jobject init) final;
+ jbyteArray NewByteArray(JNIEnv *env, jsize len) final;
+ void ReleaseByteArrayElements(JNIEnv *env, jbyteArray array, jbyte *elems, int mode) final;
+ void ReleaseFloatArrayElements(JNIEnv *env, jfloatArray array, jfloat *elems, int mode) final;
+ void ReleaseIntArrayElements(JNIEnv *env, jintArray array, jint *elems, jint mode) final;
+ void ReleaseLongArrayElements(JNIEnv *env, jlongArray array, jlong *elems, jint mode) final;
+ void SetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index, jobject val) final;
+ void SetByteArrayRegion(JNIEnv *env, jbyteArray array, jsize start, jsize len, const jbyte * buf) final;
+ void Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim, std::vector *vect) final;
+ void Convert2dJavaObjectArrayAndStoreToBinaryVector(JNIEnv *env, jobjectArray array2dJ, int dim, std::vector *vect) final;
+ void Convert2dJavaObjectArrayAndStoreToByteVector(JNIEnv *env, jobjectArray array2dJ, int dim, std::vector *vect) final;
+ jobject GetObjectField(JNIEnv * env, jobject obj, jfieldID fieldID) final;
+ jclass FindClassFromJNIEnv(JNIEnv * env, const char *name) final;
+ jmethodID GetMethodID(JNIEnv * env, jclass clazz, const char *name, const char *sig) final;
+ jfieldID GetFieldID(JNIEnv * env, jclass clazz, const char *name, const char *sig) final;
+ jint CallNonvirtualIntMethodA(JNIEnv *env, jobject obj, jclass clazz, jmethodID methodID, jvalue *args) final;
+ jlong CallNonvirtualLongMethodA(JNIEnv * env, jobject obj, jclass clazz, jmethodID methodID, jvalue* args) final;
+ void CallNonvirtualVoidMethodA(JNIEnv * env, jobject obj, jclass clazz, jmethodID methodID, jvalue* args) final;
+ void * GetPrimitiveArrayCritical(JNIEnv * env, jarray array, jboolean *isCopy) final;
+ void ReleasePrimitiveArrayCritical(JNIEnv * env, jarray array, void *carray, jint mode) final;
private:
std::unordered_map cachedClasses;
std::unordered_map cachedMethods;
- };
+ }; // class JNIUtil
+
+ /**
+ * It's common cleaner to release a primitive array within its destructor.
+ * Ex: JNIReleaseElements release_int_array_elements {[=](){
+ * jniUtil->ReleaseIntArrayElements(env, idsJ, idsCpp, JNI_ABORT);
+ * }};
+ */
+ struct JNIReleaseElements {
+ explicit JNIReleaseElements(std::function _release_func)
+ : release_func(std::move(_release_func)) {
+ }
+
+ ~JNIReleaseElements() {
+ try {
+ if (release_func) {
+ release_func();
+ }
+ } catch (...) {
+ // Ignore
+ }
+ }
+
+ std::function release_func;
+ }; // struct ReleaseIntArrayElements
// ------------------------------- CONSTANTS --------------------------------
extern const std::string FAISS_NAME;
@@ -179,6 +256,7 @@ namespace knn_jni {
extern const std::string COSINESIMIL;
extern const std::string INNER_PRODUCT;
extern const std::string NEG_DOT_PRODUCT;
+ extern const std::string HAMMING;
extern const std::string NPROBES;
extern const std::string COARSE_QUANTIZER;
diff --git a/jni/include/memory_util.h b/jni/include/memory_util.h
new file mode 100644
index 000000000..5e1fc13ae
--- /dev/null
+++ b/jni/include/memory_util.h
@@ -0,0 +1,23 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef KNNPLUGIN_JNI_INCLUDE_MEMORY_UTIL_H_
+#define KNNPLUGIN_JNI_INCLUDE_MEMORY_UTIL_H_
+
+#if defined(__GNUC__) || defined(__clang__)
+#define RESTRICT __restrict__
+#elif defined(_MSC_VER)
+#define RESTRICT __declspec(restrict)
+#else
+#define RESTRICT
+#endif
+
+#endif //KNNPLUGIN_JNI_INCLUDE_MEMORY_UTIL_H_
diff --git a/jni/include/native_engines_stream_support.h b/jni/include/native_engines_stream_support.h
new file mode 100644
index 000000000..07f97f3ac
--- /dev/null
+++ b/jni/include/native_engines_stream_support.h
@@ -0,0 +1,234 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H
+#define OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H
+
+#include "jni_util.h"
+#include "parameter_utils.h"
+#include "memory_util.h"
+
+#include
+#include
+#include
+#include
+
+namespace knn_jni {
+namespace stream {
+
+/**
+ * This class contains Java IndexInputWithBuffer reference and calls its API to copy required bytes into a read buffer.
+ */
+class NativeEngineIndexInputMediator {
+ public:
+ // Expect IndexInputWithBuffer is given as `_indexInput`.
+ NativeEngineIndexInputMediator(JNIUtilInterface *_jni_interface,
+ JNIEnv *_env,
+ jobject _indexInput)
+ : jni_interface(knn_jni::util::ParameterCheck::require_non_null(
+ _jni_interface, "jni_interface")),
+ env(knn_jni::util::ParameterCheck::require_non_null(_env, "env")),
+ indexInput(knn_jni::util::ParameterCheck::require_non_null(_indexInput, "indexInput")),
+ bufferArray((jbyteArray) (_jni_interface->GetObjectField(_env,
+ _indexInput,
+ getBufferFieldId(_jni_interface, _env)))),
+ copyBytesMethod(getCopyBytesMethod(_jni_interface, _env)),
+ remainingBytesMethod(getRemainingBytesMethod(_jni_interface, _env)) {
+ }
+
+ void copyBytes(int64_t nbytes, uint8_t * RESTRICT destination) {
+ auto jclazz = getIndexInputWithBufferClass(jni_interface, env);
+
+ while (nbytes > 0) {
+ // Call `copyBytes` to read bytes as many as possible.
+ jvalue args;
+ args.j = nbytes;
+ const auto readBytes =
+ jni_interface->CallNonvirtualIntMethodA(env, indexInput, jclazz, copyBytesMethod, &args);
+ jni_interface->HasExceptionInStack(env, "Reading bytes via IndexInput has failed.");
+
+ // === Critical Section Start ===
+
+ // Get primitive array pointer, no copy is happening in OpenJDK.
+ jbyte * RESTRICT primitiveArray =
+ (jbyte *) jni_interface->GetPrimitiveArrayCritical(env, bufferArray, nullptr);
+
+ // Copy Java bytes to C++ destination address.
+ std::memcpy(destination, primitiveArray, readBytes);
+
+ // Release the acquired primitive array pointer.
+ // JNI_ABORT tells JVM to directly free memory without copying back to Java byte[].
+ // Since we're merely copying data, we don't need to copying back.
+ // Note than when we received an internal primitive array pointer, then the mode will be ignored.
+ jni_interface->ReleasePrimitiveArrayCritical(env, bufferArray, primitiveArray, JNI_ABORT);
+
+ // === Critical Section End ===
+
+ destination += readBytes;
+ nbytes -= readBytes;
+ } // End while
+ }
+
+ int64_t remainingBytes() {
+ auto bytes = jni_interface->CallNonvirtualLongMethodA(env,
+ indexInput,
+ getIndexInputWithBufferClass(jni_interface, env),
+ remainingBytesMethod,
+ nullptr);
+ jni_interface->HasExceptionInStack(env, "Checking remaining bytes has failed.");
+ return bytes;
+ }
+
+ private:
+ static jclass getIndexInputWithBufferClass(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jclass INDEX_INPUT_WITH_BUFFER_CLASS =
+ jni_interface->FindClassFromJNIEnv(env, "org/opensearch/knn/index/store/IndexInputWithBuffer");
+ return INDEX_INPUT_WITH_BUFFER_CLASS;
+ }
+
+ static jmethodID getCopyBytesMethod(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jmethodID COPY_METHOD_ID =
+ jni_interface->GetMethodID(env, getIndexInputWithBufferClass(jni_interface, env), "copyBytes", "(J)I");
+ return COPY_METHOD_ID;
+ }
+
+ static jmethodID getRemainingBytesMethod(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jmethodID COPY_METHOD_ID =
+ jni_interface->GetMethodID(env, getIndexInputWithBufferClass(jni_interface, env), "remainingBytes", "()J");
+ return COPY_METHOD_ID;
+ }
+
+ static jfieldID getBufferFieldId(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jfieldID BUFFER_FIELD_ID =
+ jni_interface->GetFieldID(env, getIndexInputWithBufferClass(jni_interface, env), "buffer", "[B");
+ return BUFFER_FIELD_ID;
+ }
+
+ JNIUtilInterface *jni_interface;
+ JNIEnv *env;
+
+ // `IndexInputWithBuffer` instance having `IndexInput` instance obtained from `Directory` for reading.
+ jobject indexInput;
+ jbyteArray bufferArray;
+ jmethodID copyBytesMethod;
+ jmethodID remainingBytesMethod;
+}; // class NativeEngineIndexInputMediator
+
+
+
+/**
+ * This class delegates the provided index output to do IO processing.
+ * In most cases, it is expected that IndexOutputWithBuffer was passed down to this,
+ * which eventually have Lucene's IndexOutput to write bytes.
+ */
+class NativeEngineIndexOutputMediator {
+ public:
+ NativeEngineIndexOutputMediator(JNIUtilInterface *_jni_interface,
+ JNIEnv *_env,
+ jobject _indexOutput)
+ : jni_interface(knn_jni::util::ParameterCheck::require_non_null(_jni_interface, "jni_interface")),
+ env(knn_jni::util::ParameterCheck::require_non_null(_env, "env")),
+ indexOutput(knn_jni::util::ParameterCheck::require_non_null(_indexOutput, "indexOutput")),
+ bufferArray((jbyteArray) (_jni_interface->GetObjectField(_env,
+ _indexOutput,
+ getBufferFieldId(_jni_interface, _env)))),
+ writeBytesMethod(getWriteBytesMethod(_jni_interface, _env)),
+ bufferLength(jni_interface->GetJavaBytesArrayLength(env, bufferArray)),
+ nextWriteIndex() {
+ }
+
+ void writeBytes(const uint8_t * RESTRICT source, size_t nbytes) {
+ auto left = nbytes;
+ while (left > 0) {
+ const auto writeBytes = std::min(bufferLength - nextWriteIndex, left);
+
+ // === Critical Section Start ===
+
+ // Get primitive array pointer, no copy is happening in OpenJDK.
+ jbyte * RESTRICT primitiveArray =
+ (jbyte *) jni_interface->GetPrimitiveArrayCritical(env, bufferArray, nullptr);
+
+ // Copy the given bytes to Java byte[] address.
+ std::memcpy(primitiveArray + nextWriteIndex, source, writeBytes);
+
+ // Release the acquired primitive array pointer.
+ // 0 tells JVM to copy back the content, and to free the pointer. It will be ignored if we acquired an internal
+ // primitive array pointer instead of a copied version.
+ // From JNI docs:
+ // Mode 0 : copy back the content and free the elems buffer
+ // The mode argument provides information on how the array buffer should be released. mode has no effect if elems
+ // is not a copy of the elements in array.
+ jni_interface->ReleasePrimitiveArrayCritical(env, bufferArray, primitiveArray, 0);
+
+ // === Critical Section End ===
+
+ nextWriteIndex += writeBytes;
+ if (nextWriteIndex >= bufferLength) {
+ callWriteBytesInIndexOutput();
+ }
+
+ source += writeBytes;
+ left -= writeBytes;
+ } // End while
+ }
+
+ void flush() {
+ if (nextWriteIndex > 0) {
+ callWriteBytesInIndexOutput();
+ }
+ }
+
+ private:
+ static jclass getIndexOutputWithBufferClass(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jclass INDEX_OUTPUT_WITH_BUFFER_CLASS =
+ jni_interface->FindClassFromJNIEnv(env, "org/opensearch/knn/index/store/IndexOutputWithBuffer");
+ return INDEX_OUTPUT_WITH_BUFFER_CLASS;
+ }
+
+ static jmethodID getWriteBytesMethod(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jmethodID WRITE_METHOD_ID =
+ jni_interface->GetMethodID(env, getIndexOutputWithBufferClass(jni_interface, env), "writeBytes", "(I)V");
+ return WRITE_METHOD_ID;
+ }
+
+ static jfieldID getBufferFieldId(JNIUtilInterface *jni_interface, JNIEnv *env) {
+ static jfieldID BUFFER_FIELD_ID =
+ jni_interface->GetFieldID(env, getIndexOutputWithBufferClass(jni_interface, env), "buffer", "[B");
+ return BUFFER_FIELD_ID;
+ }
+
+ void callWriteBytesInIndexOutput() {
+ auto jclazz = getIndexOutputWithBufferClass(jni_interface, env);
+ // Initializing the first integer parameter of `writeBytes`.
+ // `i` represents an integer parameter.
+ jvalue args {.i = nextWriteIndex};
+ jni_interface->CallNonvirtualVoidMethodA(env, indexOutput, jclazz, writeBytesMethod, &args);
+ jni_interface->HasExceptionInStack(env, "Writing bytes via IndexOutput has failed.");
+ nextWriteIndex = 0;
+ }
+
+ JNIUtilInterface *jni_interface;
+ JNIEnv *env;
+
+ // `IndexOutputWithBuffer` instance having `IndexOutput` instance obtained from `Directory` for reading.
+ jobject indexOutput;
+ jbyteArray bufferArray;
+ jmethodID writeBytesMethod;
+ size_t bufferLength;
+ int32_t nextWriteIndex;
+}; // NativeEngineIndexOutputMediator
+
+
+
+}
+}
+
+#endif //OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H
diff --git a/jni/include/nmslib_stream_support.h b/jni/include/nmslib_stream_support.h
new file mode 100644
index 000000000..2c410dde6
--- /dev/null
+++ b/jni/include/nmslib_stream_support.h
@@ -0,0 +1,73 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H
+#define OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H
+
+#include "native_engines_stream_support.h"
+#include "utils.h" // This is from NMSLIB
+#include "parameter_utils.h"
+
+namespace knn_jni {
+namespace stream {
+
+/**
+ * NmslibIOReader implementation delegating NativeEngineIndexInputMediator to read bytes.
+ */
+class NmslibOpenSearchIOReader final : public similarity::NmslibIOReader {
+ public:
+ explicit NmslibOpenSearchIOReader(NativeEngineIndexInputMediator *_mediator)
+ : similarity::NmslibIOReader(),
+ mediator(knn_jni::util::ParameterCheck::require_non_null(_mediator, "mediator")) {
+ }
+
+ void read(char *bytes, size_t len) final {
+ if (len > 0) {
+ // Mediator calls IndexInput, then copy read bytes to `ptr`.
+ mediator->copyBytes(len, (uint8_t *) bytes);
+ }
+ }
+
+ size_t remainingBytes() final {
+ return mediator->remainingBytes();
+ }
+
+ private:
+ NativeEngineIndexInputMediator *mediator;
+}; // class NmslibOpenSearchIOReader
+
+
+class NmslibOpenSearchIOWriter final : public similarity::NmslibIOWriter {
+ public:
+ explicit NmslibOpenSearchIOWriter(NativeEngineIndexOutputMediator *_mediator)
+ : similarity::NmslibIOWriter(),
+ mediator(knn_jni::util::ParameterCheck::require_non_null(_mediator, "mediator")) {
+ }
+
+ void write(char *bytes, size_t len) final {
+ if (len > 0) {
+ mediator->writeBytes((uint8_t *) bytes, len);
+ }
+ }
+
+ void flush() final {
+ mediator->flush();
+ }
+
+ private:
+ NativeEngineIndexOutputMediator *mediator;
+}; // class NmslibOpenSearchIOWriter
+
+
+}
+}
+
+#endif //OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H
diff --git a/jni/include/nmslib_wrapper.h b/jni/include/nmslib_wrapper.h
index 6d862048a..687a96d59 100644
--- a/jni/include/nmslib_wrapper.h
+++ b/jni/include/nmslib_wrapper.h
@@ -25,19 +25,27 @@ namespace knn_jni {
namespace nmslib_wrapper {
// Create an index with ids and vectors. The configuration is defined by values in the Java map, parametersJ.
// The index is serialized to indexPathJ.
- void CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ, jobjectArray vectorsJ,
- jstring indexPathJ, jobject parametersJ);
+ void CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ, jlong vectorsAddress, jint dim,
+ jobject output, jobject parametersJ);
// Load an index from indexPathJ into memory. Use parametersJ to set any query time parameters
//
// Return a pointer to the loaded index
jlong LoadIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jstring indexPathJ, jobject parametersJ);
+ // Load an index via an input stream into memory. Use parametersJ to set any query time parameters
+ //
+ // Return a pointer to the loaded index
+ jlong LoadIndexWithStream(knn_jni::JNIUtilInterface * jniUtil,
+ JNIEnv * env,
+ jobject readStream,
+ jobject parametersJ);
+
// Execute a query against the index located in memory at indexPointerJ.
//
// Return an array of KNNQueryResults
jobjectArray QueryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
- jfloatArray queryVectorJ, jint kJ);
+ jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ);
// Free the index located in memory at indexPointerJ
void Free(jlong indexPointer);
@@ -48,10 +56,10 @@ namespace knn_jni {
struct IndexWrapper {
explicit IndexWrapper(const std::string& spaceType) {
// Index gets constructed with a reference to data (see above) but is otherwise unused
- similarity::ObjectVector data;
space.reset(similarity::SpaceFactoryRegistry::Instance().CreateSpace(spaceType, similarity::AnyParams()));
index.reset(similarity::MethodFactoryRegistry::Instance().CreateMethod(false, "hnsw", spaceType, *space, data));
}
+ similarity::ObjectVector data;
std::unique_ptr> space;
std::unique_ptr> index;
};
diff --git a/jni/include/org_opensearch_knn_jni_FaissService.h b/jni/include/org_opensearch_knn_jni_FaissService.h
index 4af9a24bc..dce580138 100644
--- a/jni/include/org_opensearch_knn_jni_FaissService.h
+++ b/jni/include/org_opensearch_knn_jni_FaissService.h
@@ -20,19 +20,101 @@ extern "C" {
#endif
/*
* Class: org_opensearch_knn_jni_FaissService
- * Method: createIndex
- * Signature: ([I[[FLjava/lang/String;Ljava/util/Map;)V
+ * Method: initIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
*/
-JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_createIndex
- (JNIEnv *, jclass, jintArray, jobjectArray, jstring, jobject);
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_initIndex(JNIEnv * env, jclass cls,
+ jlong numDocs, jint dimJ,
+ jobject parametersJ);
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: initBinaryIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_initBinaryIndex(JNIEnv * env, jclass cls,
+ jlong numDocs, jint dimJ,
+ jobject parametersJ);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: initByteIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_initByteIndex(JNIEnv * env, jclass cls,
+ jlong numDocs, jint dimJ,
+ jobject parametersJ);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: insertToIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_insertToIndex(JNIEnv * env, jclass cls, jintArray idsJ,
+ jlong vectorsAddressJ, jint dimJ,
+ jlong indexAddress, jint threadCount);
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: insertToBinaryIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_insertToBinaryIndex(JNIEnv * env, jclass cls, jintArray idsJ,
+ jlong vectorsAddressJ, jint dimJ,
+ jlong indexAddress, jint threadCount);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: insertToByteIndex
+ * Signature: ([IJILjava/lang/String;Ljava/util/Map;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_insertToByteIndex(JNIEnv * env, jclass cls, jintArray idsJ,
+ jlong vectorsAddressJ, jint dimJ,
+ jlong indexAddress, jint threadCount);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: writeIndex
+ * Signature: (JLorg/opensearch/knn/index/store/IndexOutputWithBuffer;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_writeIndex(JNIEnv *, jclass, jlong, jobject);
+
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: writeBinaryIndex
+ * Signature: (JLorg/opensearch/knn/index/store/IndexOutputWithBuffer;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_writeBinaryIndex(JNIEnv *, jclass, jlong, jobject);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: writeByteIndex
+ * Signature: (JLorg/opensearch/knn/index/store/IndexOutputWithBuffer;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_writeByteIndex(JNIEnv *, jclass, jlong, jobject);
/*
* Class: org_opensearch_knn_jni_FaissService
* Method: createIndexFromTemplate
- * Signature: ([I[[FLjava/lang/String;[BLjava/util/Map;)V
+ * Signature: ([IJILorg/opensearch/knn/index/store/IndexOutputWithBuffer;[BLjava/util/Map;)V
*/
JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_createIndexFromTemplate
- (JNIEnv *, jclass, jintArray, jobjectArray, jstring, jbyteArray, jobject);
+ (JNIEnv *, jclass, jintArray, jlong, jint, jobject, jbyteArray, jobject);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: createBinaryIndexFromTemplate
+ * Signature: ([IJILorg/opensearch/knn/index/store/IndexOutputWithBuffer;[BLjava/util/Map;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_createBinaryIndexFromTemplate
+ (JNIEnv *, jclass, jintArray, jlong, jint, jobject, jbyteArray, jobject);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: createByteIndexFromTemplate
+ * Signature: ([IJILorg/opensearch/knn/index/store/IndexOutputWithBuffer;[BLjava/util/Map;)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_createByteIndexFromTemplate
+ (JNIEnv *, jclass, jintArray, jlong, jint, jobject, jbyteArray, jobject);
/*
* Class: org_opensearch_knn_jni_FaissService
@@ -42,20 +124,92 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_createIndexFromT
JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_loadIndex
(JNIEnv *, jclass, jstring);
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: loadIndexWithStream
+ * Signature: (Lorg/opensearch/knn/index/util/IndexInputWithBuffer;)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_loadIndexWithStream
+ (JNIEnv *, jclass, jobject);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: loadBinaryIndex
+ * Signature: (Ljava/lang/String;)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_loadBinaryIndex
+ (JNIEnv *, jclass, jstring);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: loadBinaryIndexWithStream
+ * Signature: (Lorg/opensearch/knn/index/util/IndexInputWithBuffer;)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_loadBinaryIndexWithStream
+ (JNIEnv *, jclass, jobject);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: isSharedIndexStateRequired
+ * Signature: (J)Z
+ */
+JNIEXPORT jboolean JNICALL Java_org_opensearch_knn_jni_FaissService_isSharedIndexStateRequired
+ (JNIEnv *, jclass, jlong);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: initSharedIndexState
+ * Signature: (J)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_initSharedIndexState
+ (JNIEnv *, jclass, jlong);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: setSharedIndexState
+ * Signature: (JJ)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_setSharedIndexState
+ (JNIEnv *, jclass, jlong, jlong);
+
/*
* Class: org_opensearch_knn_jni_FaissService
* Method: queryIndex
- * Signature: (J[FI)[Lorg/opensearch/knn/index/KNNQueryResult;
+ * Signature: (J[FILjava/util/Map[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
*/
JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_queryIndex
- (JNIEnv *, jclass, jlong, jfloatArray, jint);
+ (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject, jintArray);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: queryIndexWithFilter
+ * Signature: (J[FILjava/util/Map[JI[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
+ */
+JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_queryIndexWithFilter
+ (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject, jlongArray, jint, jintArray);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: queryBIndexWithFilter
+ * Signature: (J[BILjava/util/Map[JI[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
+ */
+JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_queryBinaryIndexWithFilter
+ (JNIEnv *, jclass, jlong, jbyteArray, jint, jobject, jlongArray, jint, jintArray);
/*
* Class: org_opensearch_knn_jni_FaissService
* Method: free
- * Signature: (J)V
+ * Signature: (JZ)V
*/
JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_free
+ (JNIEnv *, jclass, jlong, jboolean);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: freeSharedIndexState
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_freeSharedIndexState
(JNIEnv *, jclass, jlong);
/*
@@ -74,6 +228,22 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_initLibrary
JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainIndex
(JNIEnv *, jclass, jobject, jint, jlong);
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: trainBinaryIndex
+ * Signature: (Ljava/util/Map;IJ)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainBinaryIndex
+ (JNIEnv *, jclass, jobject, jint, jlong);
+
+/*
+ * Class: org_opensearch_knn_jni_FaissService
+ * Method: trainByteIndex
+ * Signature: (Ljava/util/Map;IJ)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainByteIndex
+ (JNIEnv *, jclass, jobject, jint, jlong);
+
/*
* Class: org_opensearch_knn_jni_FaissService
* Method: transferVectors
@@ -82,13 +252,21 @@ JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainIndex
JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors
(JNIEnv *, jclass, jlong, jobjectArray);
+/*
+* Class: org_opensearch_knn_jni_FaissService
+* Method: rangeSearchIndexWithFilter
+* Signature: (J[FJLjava/util/MapI[JII)[Lorg/opensearch/knn/index/query/RangeQueryResult;
+*/
+JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_rangeSearchIndexWithFilter
+ (JNIEnv *, jclass, jlong, jfloatArray, jfloat, jobject, jint, jlongArray, jint, jintArray);
+
/*
* Class: org_opensearch_knn_jni_FaissService
- * Method: freeVectors
- * Signature: (J)V
+ * Method: rangeSearchIndex
+ * Signature: (J[FJLjava/util/MapII)[Lorg/opensearch/knn/index/query/RangeQueryResult;
*/
-JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_freeVectors
- (JNIEnv *, jclass, jlong);
+JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_rangeSearchIndex
+ (JNIEnv *, jclass, jlong, jfloatArray, jfloat, jobject, jint, jintArray);
#ifdef __cplusplus
}
diff --git a/jni/include/org_opensearch_knn_jni_JNICommons.h b/jni/include/org_opensearch_knn_jni_JNICommons.h
new file mode 100644
index 000000000..8bfbcc266
--- /dev/null
+++ b/jni/include/org_opensearch_knn_jni_JNICommons.h
@@ -0,0 +1,72 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+/* DO NOT EDIT THIS FILE - it is machine generated */
+#include
+/* Header for class org_opensearch_knn_jni_JNICommons */
+
+#ifndef _Included_org_opensearch_knn_jni_JNICommons
+#define _Included_org_opensearch_knn_jni_JNICommons
+#ifdef __cplusplus
+extern "C" {
+#endif
+/*
+ * Class: org_opensearch_knn_jni_JNICommons
+ * Method: storeVectorData
+ * Signature: (J[[FJJJ)
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_JNICommons_storeVectorData
+ (JNIEnv *, jclass, jlong, jobjectArray, jlong, jboolean);
+
+/*
+ * Class: org_opensearch_knn_jni_JNICommons
+ * Method: storeBinaryVectorData
+ * Signature: (J[[FJJ)
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_JNICommons_storeBinaryVectorData
+ (JNIEnv *, jclass, jlong, jobjectArray, jlong, jboolean);
+
+/*
+ * Class: org_opensearch_knn_jni_JNICommons
+ * Method: storeByteVectorData
+ * Signature: (J[[FJJ)
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_JNICommons_storeByteVectorData
+ (JNIEnv *, jclass, jlong, jobjectArray, jlong, jboolean);
+
+/*
+ * Class: org_opensearch_knn_jni_JNICommons
+ * Method: freeVectorData
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_JNICommons_freeVectorData
+ (JNIEnv *, jclass, jlong);
+
+/*
+* Class: org_opensearch_knn_jni_JNICommons
+* Method: freeBinaryVectorData
+* Signature: (J)V
+*/
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_JNICommons_freeBinaryVectorData
+(JNIEnv *, jclass, jlong);
+
+/*
+* Class: org_opensearch_knn_jni_JNICommons
+* Method: freeByteVectorData
+* Signature: (J)V
+*/
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_JNICommons_freeByteVectorData
+(JNIEnv *, jclass, jlong);
+
+#ifdef __cplusplus
+}
+#endif
+#endif
diff --git a/jni/include/org_opensearch_knn_jni_NmslibService.h b/jni/include/org_opensearch_knn_jni_NmslibService.h
index dd907581d..0e035c3dd 100644
--- a/jni/include/org_opensearch_knn_jni_NmslibService.h
+++ b/jni/include/org_opensearch_knn_jni_NmslibService.h
@@ -21,10 +21,10 @@ extern "C" {
/*
* Class: org_opensearch_knn_jni_NmslibService
* Method: createIndex
- * Signature: ([I[[FLjava/lang/String;Ljava/util/Map;)V
+ * Signature: ([IJILorg/opensearch/knn/index/store/IndexOutputWithBuffer;Ljava/util/Map;)V
*/
JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_NmslibService_createIndex
- (JNIEnv *, jclass, jintArray, jobjectArray, jstring, jobject);
+ (JNIEnv *, jclass, jintArray, jlong, jint, jobject, jobject);
/*
* Class: org_opensearch_knn_jni_NmslibService
@@ -34,13 +34,21 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_NmslibService_createIndex
JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_NmslibService_loadIndex
(JNIEnv *, jclass, jstring, jobject);
+/*
+ * Class: org_opensearch_knn_jni_NmslibService
+ * Method: loadIndexWithStream
+ * Signature: (Lorg/opensearch/knn/index/store/IndexInputWithBuffer;Ljava/util/Map;)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_NmslibService_loadIndexWithStream
+ (JNIEnv *, jclass, jobject, jobject);
+
/*
* Class: org_opensearch_knn_jni_NmslibService
* Method: queryIndex
- * Signature: (J[FI)[Lorg/opensearch/knn/index/KNNQueryResult;
+ * Signature: (J[FI)[Lorg/opensearch/knn/index/query/KNNQueryResult;
*/
JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_NmslibService_queryIndex
- (JNIEnv *, jclass, jlong, jfloatArray, jint);
+ (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject);
/*
* Class: org_opensearch_knn_jni_NmslibService
diff --git a/jni/include/parameter_utils.h b/jni/include/parameter_utils.h
new file mode 100644
index 000000000..aff922324
--- /dev/null
+++ b/jni/include/parameter_utils.h
@@ -0,0 +1,39 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+#ifndef KNNPLUGIN_JNI_INCLUDE_PARAMETER_UTILS_H_
+#define KNNPLUGIN_JNI_INCLUDE_PARAMETER_UTILS_H_
+
+#include
+#include
+
+namespace knn_jni {
+namespace util {
+
+struct ParameterCheck {
+ template
+ static PtrType *require_non_null(PtrType *ptr, const char *parameter_name) {
+ if (ptr == nullptr) {
+ throw std::invalid_argument(std::string("Parameter [") + parameter_name + "] should not be null.");
+ }
+ return ptr;
+ }
+
+ private:
+ ParameterCheck() = default;
+}; // class ParameterCheck
+
+
+
+}
+} // namespace knn_jni
+
+#endif //KNNPLUGIN_JNI_INCLUDE_PARAMETER_UTILS_H_
diff --git a/jni/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch b/jni/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
new file mode 100644
index 000000000..26b143346
--- /dev/null
+++ b/jni/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
@@ -0,0 +1,1091 @@
+From e775a8e65da96232822d5aed77f538592fccffda Mon Sep 17 00:00:00 2001
+From: Heemin Kim
+Date: Tue, 30 Jan 2024 14:43:56 -0800
+Subject: [PATCH] Add IDGrouper for HNSW
+
+Signed-off-by: Heemin Kim
+---
+ faiss/CMakeLists.txt | 3 +
+ faiss/Index.h | 6 +-
+ faiss/IndexHNSW.cpp | 13 +-
+ faiss/IndexIDMap.cpp | 29 +++++
+ faiss/IndexIDMap.h | 22 ++++
+ faiss/impl/HNSW.cpp | 6 +
+ faiss/impl/IDGrouper.cpp | 51 ++++++++
+ faiss/impl/IDGrouper.h | 51 ++++++++
+ faiss/impl/ResultHandler.h | 189 +++++++++++++++++++++++++++++
+ faiss/utils/GroupHeap.h | 182 ++++++++++++++++++++++++++++
+ tests/CMakeLists.txt | 2 +
+ tests/test_group_heap.cpp | 98 +++++++++++++++
+ tests/test_id_grouper.cpp | 241 +++++++++++++++++++++++++++++++++++++
+ 13 files changed, 889 insertions(+), 4 deletions(-)
+ create mode 100644 faiss/impl/IDGrouper.cpp
+ create mode 100644 faiss/impl/IDGrouper.h
+ create mode 100644 faiss/utils/GroupHeap.h
+ create mode 100644 tests/test_group_heap.cpp
+ create mode 100644 tests/test_id_grouper.cpp
+
+diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt
+index 2871d974..d0bcec6a 100644
+--- a/faiss/CMakeLists.txt
++++ b/faiss/CMakeLists.txt
+@@ -55,6 +55,7 @@ set(FAISS_SRC
+ impl/AuxIndexStructures.cpp
+ impl/CodePacker.cpp
+ impl/IDSelector.cpp
++ impl/IDGrouper.cpp
+ impl/FaissException.cpp
+ impl/HNSW.cpp
+ impl/NSG.cpp
+@@ -151,6 +152,7 @@ set(FAISS_HEADERS
+ impl/AuxIndexStructures.h
+ impl/CodePacker.h
+ impl/IDSelector.h
++ impl/IDGrouper.h
+ impl/DistanceComputer.h
+ impl/FaissAssert.h
+ impl/FaissException.h
+@@ -186,6 +188,7 @@ set(FAISS_HEADERS
+ invlists/InvertedListsIOHook.h
+ utils/AlignedTable.h
+ utils/bf16.h
++ utils/GroupHeap.h
+ utils/Heap.h
+ utils/WorkerThread.h
+ utils/distances.h
+diff --git a/faiss/Index.h b/faiss/Index.h
+index f57140ec..8f511e5d 100644
+--- a/faiss/Index.h
++++ b/faiss/Index.h
+@@ -51,8 +51,9 @@
+ namespace faiss {
+
+ /// Forward declarations see impl/AuxIndexStructures.h, impl/IDSelector.h
+-/// and impl/DistanceComputer.h
++/// ,impl/IDGrouper.h and impl/DistanceComputer.h
+ struct IDSelector;
++struct IDGrouper;
+ struct RangeSearchResult;
+ struct DistanceComputer;
+
+@@ -64,6 +65,9 @@ struct DistanceComputer;
+ struct SearchParameters {
+ /// if non-null, only these IDs will be considered during search.
+ IDSelector* sel = nullptr;
++ /// if non-null, only best matched ID per group will be included in the
++ /// result.
++ IDGrouper* grp = nullptr;
+ /// make sure we can dynamic_cast this
+ virtual ~SearchParameters() {}
+ };
+diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp
+index 6a1186ca..9c8a8255 100644
+--- a/faiss/IndexHNSW.cpp
++++ b/faiss/IndexHNSW.cpp
+@@ -301,10 +301,17 @@ void IndexHNSW::search(
+ const SearchParameters* params_in) const {
+ FAISS_THROW_IF_NOT(k > 0);
+
+- using RH = HeapBlockResultHandler;
+- RH bres(n, distances, labels, k);
++ if (params_in && params_in->grp) {
++ using RH = GroupedHeapBlockResultHandler;
++ RH bres(n, distances, labels, k, params_in->grp);
+
+- hnsw_search(this, n, x, bres, params_in);
++ hnsw_search(this, n, x, bres, params_in);
++ } else {
++ using RH = HeapBlockResultHandler;
++ RH bres(n, distances, labels, k);
++
++ hnsw_search(this, n, x, bres, params_in);
++ }
+
+ if (is_similarity_metric(this->metric_type)) {
+ // we need to revert the negated distances
+diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp
+index dc84052b..3f375e7b 100644
+--- a/faiss/IndexIDMap.cpp
++++ b/faiss/IndexIDMap.cpp
+@@ -102,6 +102,23 @@ struct ScopedSelChange {
+ }
+ };
+
++/// RAII object to reset the IDGrouper in the params object
++struct ScopedGrpChange {
++ SearchParameters* params = nullptr;
++ IDGrouper* old_grp = nullptr;
++
++ void set(SearchParameters* params_2, IDGrouper* new_grp) {
++ this->params = params_2;
++ old_grp = params_2->grp;
++ params_2->grp = new_grp;
++ }
++ ~ScopedGrpChange() {
++ if (params) {
++ params->grp = old_grp;
++ }
++ }
++};
++
+ } // namespace
+
+ template
+@@ -114,6 +131,8 @@ void IndexIDMapTemplate::search(
+ const SearchParameters* params) const {
+ IDSelectorTranslated this_idtrans(this->id_map, nullptr);
+ ScopedSelChange sel_change;
++ IDGrouperTranslated this_idgrptrans(this->id_map, nullptr);
++ ScopedGrpChange grp_change;
+
+ if (params && params->sel) {
+ auto idtrans = dynamic_cast(params->sel);
+@@ -131,6 +150,16 @@ void IndexIDMapTemplate::search(
+ sel_change.set(params_non_const, &this_idtrans);
+ }
+ }
++
++ if (params && params->grp) {
++ auto idtrans = dynamic_cast(params->grp);
++
++ if (!idtrans) {
++ auto params_non_const = const_cast(params);
++ this_idgrptrans.grp = params->grp;
++ grp_change.set(params_non_const, &this_idgrptrans);
++ }
++ }
+ index->search(n, x, k, distances, labels, params);
+ idx_t* li = labels;
+ #pragma omp parallel for
+diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h
+index 2d164123..a68887bd 100644
+--- a/faiss/IndexIDMap.h
++++ b/faiss/IndexIDMap.h
+@@ -9,6 +9,7 @@
+
+ #include
+ #include
++#include
+ #include
+
+ #include
+@@ -124,4 +125,25 @@ struct IDSelectorTranslated : IDSelector {
+ }
+ };
+
++// IDGrouper that translates the ids using an IDMap
++struct IDGrouperTranslated : IDGrouper {
++ const std::vector& id_map;
++ const IDGrouper* grp;
++
++ IDGrouperTranslated(
++ const std::vector& id_map,
++ const IDGrouper* grp)
++ : id_map(id_map), grp(grp) {}
++
++ IDGrouperTranslated(IndexBinaryIDMap& index_idmap, const IDGrouper* grp)
++ : id_map(index_idmap.id_map), grp(grp) {}
++
++ IDGrouperTranslated(IndexIDMap& index_idmap, const IDGrouper* grp)
++ : id_map(index_idmap.id_map), grp(grp) {}
++
++ idx_t get_group(idx_t id) const override {
++ return grp->get_group(id_map[id]);
++ }
++};
++
+ } // namespace faiss
+diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp
+index c3693fd9..7ae28062 100644
+--- a/faiss/impl/HNSW.cpp
++++ b/faiss/impl/HNSW.cpp
+@@ -906,6 +906,12 @@ int extract_k_from_ResultHandler(ResultHandler& res) {
+ if (auto hres = dynamic_cast(&res)) {
+ return hres->k;
+ }
++
++ if (auto hres = dynamic_cast<
++ GroupedHeapBlockResultHandler::SingleResultHandler*>(&res)) {
++ return hres->k;
++ }
++
+ return 1;
+ }
+
+diff --git a/faiss/impl/IDGrouper.cpp b/faiss/impl/IDGrouper.cpp
+new file mode 100644
+index 00000000..ca9f5fda
+--- /dev/null
++++ b/faiss/impl/IDGrouper.cpp
+@@ -0,0 +1,51 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++
++#include
++#include
++#include
++
++namespace faiss {
++
++/***********************************************************************
++ * IDGrouperBitmap
++ ***********************************************************************/
++
++IDGrouperBitmap::IDGrouperBitmap(size_t n, uint64_t* bitmap)
++ : n(n), bitmap(bitmap) {}
++
++idx_t IDGrouperBitmap::get_group(idx_t id) const {
++ assert(id >= 0 && "id shouldn't be less than zero");
++ assert(id < this->n * 64 && "is should be less than total number of bits");
++
++ idx_t index = id >> 6; // div by 64
++ uint64_t block = this->bitmap[index] >>
++ (id & 63); // Equivalent of words[i] >> (index % 64)
++ // block is non zero after right shift, it means, next set bit is in current
++ // block The index of set bit is "given index" + "trailing zero in the right
++ // shifted word"
++ if (block != 0) {
++ return id + __builtin_ctzll(block);
++ }
++
++ while (++index < this->n) {
++ block = this->bitmap[index];
++ if (block != 0) {
++ return (index << 6) + __builtin_ctzll(block);
++ }
++ }
++
++ return NO_MORE_DOCS;
++}
++
++void IDGrouperBitmap::set_group(idx_t group_id) {
++ idx_t index = group_id >> 6;
++ this->bitmap[index] |= 1ULL
++ << (group_id & 63); // Equivalent of 1ULL << (value % 64)
++}
++
++} // namespace faiss
+diff --git a/faiss/impl/IDGrouper.h b/faiss/impl/IDGrouper.h
+new file mode 100644
+index 00000000..d56113d9
+--- /dev/null
++++ b/faiss/impl/IDGrouper.h
+@@ -0,0 +1,51 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++
++#pragma once
++
++#include
++#include
++#include
++
++#include
++
++/** IDGrouper is intended to define a group of vectors to include only
++ * the nearest vector of each group during search */
++
++namespace faiss {
++
++/** Encapsulates a group id of ids */
++struct IDGrouper {
++ const idx_t NO_MORE_DOCS = std::numeric_limits::max();
++ virtual idx_t get_group(idx_t id) const = 0;
++ virtual ~IDGrouper() {}
++};
++
++/** One bit per element. Constructed with a bitmap, size ceil(n / 8).
++ */
++struct IDGrouperBitmap : IDGrouper {
++ // length of the bitmap array
++ size_t n;
++
++ // Array of uint64_t holding the bits
++ // Using uint64_t to leverage function __builtin_ctzll which is defined in
++ // faiss/impl/platform_macros.h Group id of a given id is next set bit in
++ // the bitmap
++ uint64_t* bitmap;
++
++ /** Construct with a binary mask
++ *
++ * @param n size of the bitmap array
++ * @param bitmap group id of a given id is next set bit in the bitmap
++ */
++ IDGrouperBitmap(size_t n, uint64_t* bitmap);
++ idx_t get_group(idx_t id) const final;
++ void set_group(idx_t group_id);
++ ~IDGrouperBitmap() override {}
++};
++
++} // namespace faiss
+diff --git a/faiss/impl/ResultHandler.h b/faiss/impl/ResultHandler.h
+index 3116eb24..126ed015 100644
+--- a/faiss/impl/ResultHandler.h
++++ b/faiss/impl/ResultHandler.h
+@@ -14,6 +14,8 @@
+ #include
+ #include
+ #include
++#include
++#include
+ #include
+ #include
+
+@@ -286,6 +288,193 @@ struct HeapBlockResultHandler : BlockResultHandler {
+ }
+ };
+
++/*****************************************************************
++ * Heap based result handler with grouping
++ *****************************************************************/
++
++template
++struct GroupedHeapBlockResultHandler : BlockResultHandler {
++ using T = typename C::T;
++ using TI = typename C::TI;
++ using BlockResultHandler::i0;
++ using BlockResultHandler::i1;
++
++ T* heap_dis_tab;
++ TI* heap_ids_tab;
++ int64_t k; // number of results to keep
++
++ IDGrouper* id_grouper;
++ TI* heap_group_ids_tab;
++ std::unordered_map* group_id_to_index_in_heap_tab;
++
++ GroupedHeapBlockResultHandler(
++ size_t nq,
++ T* heap_dis_tab,
++ TI* heap_ids_tab,
++ size_t k,
++ IDGrouper* id_grouper)
++ : BlockResultHandler(nq),
++ heap_dis_tab(heap_dis_tab),
++ heap_ids_tab(heap_ids_tab),
++ k(k),
++ id_grouper(id_grouper) {}
++
++ /******************************************************
++ * API for 1 result at a time (each SingleResultHandler is
++ * called from 1 thread)
++ */
++
++ struct SingleResultHandler : ResultHandler {
++ GroupedHeapBlockResultHandler& hr;
++ using ResultHandler::threshold;
++ size_t k;
++
++ T* heap_dis;
++ TI* heap_ids;
++ TI* heap_group_ids;
++ std::unordered_map group_id_to_index_in_heap;
++
++ explicit SingleResultHandler(GroupedHeapBlockResultHandler& hr)
++ : hr(hr), k(hr.k) {}
++
++ /// begin results for query # i
++ void begin(size_t i) {
++ heap_dis = hr.heap_dis_tab + i * k;
++ heap_ids = hr.heap_ids_tab + i * k;
++ heap_heapify(k, heap_dis, heap_ids);
++ threshold = heap_dis[0];
++ heap_group_ids = new TI[hr.k];
++ for (size_t i = 0; i < hr.k; i++) {
++ heap_group_ids[i] = -1;
++ }
++ }
++
++ /// add one result for query i
++ bool add_result(T dis, TI idx) final {
++ if (!C::cmp(threshold, dis)) {
++ return false;
++ }
++
++ idx_t group_id = hr.id_grouper->get_group(idx);
++ typename std::unordered_map::const_iterator it_pos =
++ group_id_to_index_in_heap.find(group_id);
++ if (it_pos == group_id_to_index_in_heap.end()) {
++ group_heap_replace_top(
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids,
++ dis,
++ idx,
++ group_id,
++ &group_id_to_index_in_heap);
++ threshold = heap_dis[0];
++ return true;
++ } else {
++ size_t pos = it_pos->second;
++ if (!C::cmp(heap_dis[pos], dis)) {
++ return false;
++ }
++ group_heap_replace_at(
++ pos,
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids,
++ dis,
++ idx,
++ group_id,
++ &group_id_to_index_in_heap);
++ threshold = heap_dis[0];
++ return true;
++ }
++ }
++
++ /// series of results for query i is done
++ void end() {
++ heap_reorder(k, heap_dis, heap_ids);
++ delete[] heap_group_ids;
++ }
++ };
++
++ /******************************************************
++ * API for multiple results (called from 1 thread)
++ */
++
++ /// begin
++ void begin_multiple(size_t i0_2, size_t i1_2) final {
++ this->i0 = i0_2;
++ this->i1 = i1_2;
++ for (size_t i = i0; i < i1; i++) {
++ heap_heapify(k, heap_dis_tab + i * k, heap_ids_tab + i * k);
++ }
++ size_t size = (i1 - i0) * k;
++ heap_group_ids_tab = new TI[size];
++ for (size_t i = 0; i < size; i++) {
++ heap_group_ids_tab[i] = -1;
++ }
++ group_id_to_index_in_heap_tab =
++ new std::unordered_map[i1 - i0];
++ }
++
++ /// add results for query i0..i1 and j0..j1
++ void add_results(size_t j0, size_t j1, const T* dis_tab) final {
++#pragma omp parallel for
++ for (int64_t i = i0; i < i1; i++) {
++ T* heap_dis = heap_dis_tab + i * k;
++ TI* heap_ids = heap_ids_tab + i * k;
++ const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
++ T thresh = heap_dis[0]; // NOLINT(*-use-default-none)
++ for (size_t j = j0; j < j1; j++) {
++ T dis = dis_tab_i[j];
++ if (C::cmp(thresh, dis)) {
++ idx_t group_id = id_grouper->get_group(j);
++ typename std::unordered_map::const_iterator
++ it_pos = group_id_to_index_in_heap_tab[i - i0].find(
++ group_id);
++ if (it_pos == group_id_to_index_in_heap_tab[i - i0].end()) {
++ group_heap_replace_top(
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids_tab + ((i - i0) * k),
++ dis,
++ j,
++ group_id,
++ &group_id_to_index_in_heap_tab[i - i0]);
++ thresh = heap_dis[0];
++ } else {
++ size_t pos = it_pos->first;
++ if (C::cmp(heap_dis[pos], dis)) {
++ group_heap_replace_at(
++ pos,
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids_tab + ((i - i0) * k),
++ dis,
++ j,
++ group_id,
++ &group_id_to_index_in_heap_tab[i - i0]);
++ thresh = heap_dis[0];
++ }
++ }
++ }
++ }
++ }
++ }
++
++ /// series of results for queries i0..i1 is done
++ void end_multiple() final {
++ // maybe parallel for
++ for (size_t i = i0; i < i1; i++) {
++ heap_reorder(k, heap_dis_tab + i * k, heap_ids_tab + i * k);
++ }
++ delete[] group_id_to_index_in_heap_tab;
++ delete[] heap_group_ids_tab;
++ }
++};
++
+ /*****************************************************************
+ * Reservoir result handler
+ *
+diff --git a/faiss/utils/GroupHeap.h b/faiss/utils/GroupHeap.h
+new file mode 100644
+index 00000000..3b7078da
+--- /dev/null
++++ b/faiss/utils/GroupHeap.h
+@@ -0,0 +1,182 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++
++#pragma once
++
++#include
++#include
++#include
++
++#include
++#include
++#include
++
++#include
++#include
++
++#include
++#include
++
++namespace faiss {
++
++/**
++ * From start_index, it compare its value with parent node's and swap if needed.
++ * Continue until either there is no swap or it reaches the top node.
++ */
++template
++static inline void group_up_heap(
++ typename C::T* heap_dis,
++ typename C::TI* heap_ids,
++ typename C::TI* heap_group_ids,
++ std::unordered_map* group_id_to_index_in_heap,
++ size_t start_index) {
++ heap_dis--; /* Use 1-based indexing for easier node->child translation */
++ heap_ids--;
++ heap_group_ids--;
++ size_t i = start_index + 1, i_father;
++ typename C::T target_dis = heap_dis[i];
++ typename C::TI target_id = heap_ids[i];
++ typename C::TI target_group_id = heap_group_ids[i];
++
++ while (i > 1) {
++ i_father = i >> 1;
++ if (!C::cmp2(
++ target_dis,
++ heap_dis[i_father],
++ target_id,
++ heap_ids[i_father])) {
++ /* the heap structure is ok */
++ break;
++ }
++ heap_dis[i] = heap_dis[i_father];
++ heap_ids[i] = heap_ids[i_father];
++ heap_group_ids[i] = heap_group_ids[i_father];
++ (*group_id_to_index_in_heap)[heap_group_ids[i]] = i - 1;
++ i = i_father;
++ }
++ heap_dis[i] = target_dis;
++ heap_ids[i] = target_id;
++ heap_group_ids[i] = target_group_id;
++ (*group_id_to_index_in_heap)[heap_group_ids[i]] = i - 1;
++}
++
++/**
++ * From start_index, it compare its value with child node's and swap if needed.
++ * Continue until either there is no swap or it reaches the leaf node.
++ */
++template
++static inline void group_down_heap(
++ size_t k,
++ typename C::T* heap_dis,
++ typename C::TI* heap_ids,
++ typename C::TI* heap_group_ids,
++ std::unordered_map* group_id_to_index_in_heap,
++ size_t start_index) {
++ heap_dis--; /* Use 1-based indexing for easier node->child translation */
++ heap_ids--;
++ heap_group_ids--;
++ size_t i = start_index + 1, i1, i2;
++ typename C::T target_dis = heap_dis[i];
++ typename C::TI target_id = heap_ids[i];
++ typename C::TI target_group_id = heap_group_ids[i];
++
++ while (1) {
++ i1 = i << 1;
++ i2 = i1 + 1;
++ if (i1 > k) {
++ break;
++ }
++
++ // Note that C::cmp2() is a bool function answering
++ // `(a1 > b1) || ((a1 == b1) && (a2 > b2))` for max
++ // heap and same with the `<` sign for min heap.
++ if ((i2 == k + 1) ||
++ C::cmp2(heap_dis[i1], heap_dis[i2], heap_ids[i1], heap_ids[i2])) {
++ if (C::cmp2(target_dis, heap_dis[i1], target_id, heap_ids[i1])) {
++ break;
++ }
++ heap_dis[i] = heap_dis[i1];
++ heap_ids[i] = heap_ids[i1];
++ heap_group_ids[i] = heap_group_ids[i1];
++ (*group_id_to_index_in_heap)[heap_group_ids[i]] = i - 1;
++ i = i1;
++ } else {
++ if (C::cmp2(target_dis, heap_dis[i2], target_id, heap_ids[i2])) {
++ break;
++ }
++ heap_dis[i] = heap_dis[i2];
++ heap_ids[i] = heap_ids[i2];
++ heap_group_ids[i] = heap_group_ids[i2];
++ (*group_id_to_index_in_heap)[heap_group_ids[i]] = i - 1;
++ i = i2;
++ }
++ }
++ heap_dis[i] = target_dis;
++ heap_ids[i] = target_id;
++ heap_group_ids[i] = target_group_id;
++ (*group_id_to_index_in_heap)[heap_group_ids[i]] = i - 1;
++}
++
++template
++static inline void group_heap_replace_top(
++ size_t k,
++ typename C::T* heap_dis,
++ typename C::TI* heap_ids,
++ typename C::TI* heap_group_ids,
++ typename C::T dis,
++ typename C::TI id,
++ typename C::TI group_id,
++ std::unordered_map* group_id_to_index_in_heap) {
++ assert(group_id_to_index_in_heap->find(group_id) ==
++ group_id_to_index_in_heap->end() &&
++ "group id should not exist in the binary heap");
++
++ group_id_to_index_in_heap->erase(heap_group_ids[0]);
++ heap_group_ids[0] = group_id;
++ heap_dis[0] = dis;
++ heap_ids[0] = id;
++ (*group_id_to_index_in_heap)[group_id] = 0;
++ group_down_heap(
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids,
++ group_id_to_index_in_heap,
++ 0);
++}
++
++template
++static inline void group_heap_replace_at(
++ size_t pos,
++ size_t k,
++ typename C::T* heap_dis,
++ typename C::TI* heap_ids,
++ typename C::TI* heap_group_ids,
++ typename C::T dis,
++ typename C::TI id,
++ typename C::TI group_id,
++ std::unordered_map* group_id_to_index_in_heap) {
++ assert(group_id_to_index_in_heap->find(group_id) !=
++ group_id_to_index_in_heap->end() &&
++ "group id should exist in the binary heap");
++ assert(group_id_to_index_in_heap->find(group_id)->second == pos &&
++ "index of group id in the heap should be same as pos");
++
++ heap_dis[pos] = dis;
++ heap_ids[pos] = id;
++ group_up_heap(
++ heap_dis, heap_ids, heap_group_ids, group_id_to_index_in_heap, pos);
++ group_down_heap(
++ k,
++ heap_dis,
++ heap_ids,
++ heap_group_ids,
++ group_id_to_index_in_heap,
++ pos);
++}
++
++} // namespace faiss
+\ No newline at end of file
+diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
+index c41edf0c..87ab2020 100644
+--- a/tests/CMakeLists.txt
++++ b/tests/CMakeLists.txt
+@@ -27,6 +27,8 @@ set(FAISS_TEST_SRC
+ test_approx_topk.cpp
+ test_RCQ_cropping.cpp
+ test_distances_simd.cpp
++ test_id_grouper.cpp
++ test_group_heap.cpp
+ test_heap.cpp
+ test_code_distance.cpp
+ test_hnsw.cpp
+diff --git a/tests/test_group_heap.cpp b/tests/test_group_heap.cpp
+new file mode 100644
+index 00000000..0e8fe7a7
+--- /dev/null
++++ b/tests/test_group_heap.cpp
+@@ -0,0 +1,98 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++#include
++#include
++#include
++#include
++
++using namespace faiss;
++
++TEST(GroupHeap, group_heap_replace_top) {
++ using C = CMax;
++ const int k = 100;
++ float binary_heap_values[k];
++ int64_t binary_heap_ids[k];
++ heap_heapify(k, binary_heap_values, binary_heap_ids);
++ int64_t binary_heap_group_ids[k];
++ for (size_t i = 0; i < k; i++) {
++ binary_heap_group_ids[i] = -1;
++ }
++ std::unordered_map group_id_to_index_in_heap;
++ for (int i = 1000; i > 0; i--) {
++ group_heap_replace_top(
++ k,
++ binary_heap_values,
++ binary_heap_ids,
++ binary_heap_group_ids,
++ i * 10.0,
++ i,
++ i,
++ &group_id_to_index_in_heap);
++ }
++
++ heap_reorder(k, binary_heap_values, binary_heap_ids);
++
++ for (int i = 0; i < k; i++) {
++ ASSERT_EQ((i + 1) * 10.0, binary_heap_values[i]);
++ ASSERT_EQ(i + 1, binary_heap_ids[i]);
++ }
++}
++
++TEST(GroupHeap, group_heap_replace_at) {
++ using C = CMax;
++ const int k = 10;
++ float binary_heap_values[k];
++ int64_t binary_heap_ids[k];
++ heap_heapify(k, binary_heap_values, binary_heap_ids);
++ int64_t binary_heap_group_ids[k];
++ for (size_t i = 0; i < k; i++) {
++ binary_heap_group_ids[i] = -1;
++ }
++ std::unordered_map group_id_to_index_in_heap;
++
++ std::unordered_map group_id_to_id;
++ for (int i = 1000; i > 0; i--) {
++ int64_t group_id = rand() % 100;
++ group_id_to_id[group_id] = i;
++ if (group_id_to_index_in_heap.find(group_id) ==
++ group_id_to_index_in_heap.end()) {
++ group_heap_replace_top(
++ k,
++ binary_heap_values,
++ binary_heap_ids,
++ binary_heap_group_ids,
++ i * 10.0,
++ i,
++ group_id,
++ &group_id_to_index_in_heap);
++ } else {
++ group_heap_replace_at(
++ group_id_to_index_in_heap.at(group_id),
++ k,
++ binary_heap_values,
++ binary_heap_ids,
++ binary_heap_group_ids,
++ i * 10.0,
++ i,
++ group_id,
++ &group_id_to_index_in_heap);
++ }
++ }
++
++ heap_reorder(k, binary_heap_values, binary_heap_ids);
++
++ std::vector sorted_ids;
++ for (const auto& pair : group_id_to_id) {
++ sorted_ids.push_back(pair.second);
++ }
++ std::sort(sorted_ids.begin(), sorted_ids.end());
++
++ for (int i = 0; i < k && binary_heap_ids[i] != -1; i++) {
++ ASSERT_EQ(sorted_ids[i] * 10.0, binary_heap_values[i]);
++ ASSERT_EQ(sorted_ids[i], binary_heap_ids[i]);
++ }
++}
+diff --git a/tests/test_id_grouper.cpp b/tests/test_id_grouper.cpp
+new file mode 100644
+index 00000000..6601795b
+--- /dev/null
++++ b/tests/test_id_grouper.cpp
+@@ -0,0 +1,241 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++#include
++#include
++#include
++#include
++
++#include
++#include
++#include
++#include
++#include
++
++// 64-bit int
++using idx_t = faiss::idx_t;
++
++using namespace faiss;
++
++TEST(IdGrouper, get_group) {
++ uint64_t ids1[1] = {0b1000100010001000};
++ IDGrouperBitmap bitmap(1, ids1);
++
++ ASSERT_EQ(3, bitmap.get_group(0));
++ ASSERT_EQ(3, bitmap.get_group(1));
++ ASSERT_EQ(3, bitmap.get_group(2));
++ ASSERT_EQ(3, bitmap.get_group(3));
++ ASSERT_EQ(7, bitmap.get_group(4));
++ ASSERT_EQ(7, bitmap.get_group(5));
++ ASSERT_EQ(7, bitmap.get_group(6));
++ ASSERT_EQ(7, bitmap.get_group(7));
++ ASSERT_EQ(11, bitmap.get_group(8));
++ ASSERT_EQ(11, bitmap.get_group(9));
++ ASSERT_EQ(11, bitmap.get_group(10));
++ ASSERT_EQ(11, bitmap.get_group(11));
++ ASSERT_EQ(15, bitmap.get_group(12));
++ ASSERT_EQ(15, bitmap.get_group(13));
++ ASSERT_EQ(15, bitmap.get_group(14));
++ ASSERT_EQ(15, bitmap.get_group(15));
++ ASSERT_EQ(bitmap.NO_MORE_DOCS, bitmap.get_group(16));
++}
++
++TEST(IdGrouper, set_group) {
++ idx_t group_ids[] = {64, 127, 128, 1022};
++ uint64_t ids[16] = {}; // 1023 / 64 + 1
++ IDGrouperBitmap bitmap(16, ids);
++
++ for (int i = 0; i < 4; i++) {
++ bitmap.set_group(group_ids[i]);
++ }
++
++ int group_id_index = 0;
++ for (int i = 0; i <= group_ids[3]; i++) {
++ ASSERT_EQ(group_ids[group_id_index], bitmap.get_group(i));
++ if (group_ids[group_id_index] == i) {
++ group_id_index++;
++ }
++ }
++ ASSERT_EQ(bitmap.NO_MORE_DOCS, bitmap.get_group(group_ids[3] + 1));
++}
++
++TEST(IdGrouper, sanity_test) {
++ int d = 1; // dimension
++ int nb = 10; // database size
++
++ std::mt19937 rng;
++ std::uniform_real_distribution<> distrib;
++
++ float* xb = new float[d * nb];
++
++ for (int i = 0; i < nb; i++) {
++ for (int j = 0; j < d; j++)
++ xb[d * i + j] = distrib(rng);
++ xb[d * i] += i / 1000.;
++ }
++
++ uint64_t bitmap[1] = {};
++ faiss::IDGrouperBitmap id_grouper(1, bitmap);
++ for (int i = 0; i < nb; i++) {
++ id_grouper.set_group(i);
++ }
++
++ int k = 5;
++ int m = 8;
++ faiss::Index* index =
++ new faiss::IndexHNSWFlat(d, m, faiss::MetricType::METRIC_L2);
++ index->add(nb, xb); // add vectors to the index
++
++ // search
++ auto pSearchParameters = new faiss::SearchParametersHNSW();
++
++ idx_t* expectedI = new idx_t[k];
++ float* expectedD = new float[k];
++ index->search(1, xb, k, expectedD, expectedI, pSearchParameters);
++
++ idx_t* I = new idx_t[k];
++ float* D = new float[k];
++ pSearchParameters->grp = &id_grouper;
++ index->search(1, xb, k, D, I, pSearchParameters);
++
++ // compare
++ for (int j = 0; j < k; j++) {
++ ASSERT_EQ(expectedI[j], I[j]);
++ ASSERT_EQ(expectedD[j], D[j]);
++ }
++
++ delete[] expectedI;
++ delete[] expectedD;
++ delete[] I;
++ delete[] D;
++ delete[] xb;
++}
++
++TEST(IdGrouper, bitmap_with_hnsw) {
++ int d = 1; // dimension
++ int nb = 10; // database size
++
++ std::mt19937 rng;
++ std::uniform_real_distribution<> distrib;
++
++ float* xb = new float[d * nb];
++
++ for (int i = 0; i < nb; i++) {
++ for (int j = 0; j < d; j++)
++ xb[d * i + j] = distrib(rng);
++ xb[d * i] += i / 1000.;
++ }
++
++ uint64_t bitmap[1] = {};
++ faiss::IDGrouperBitmap id_grouper(1, bitmap);
++ for (int i = 0; i < nb; i++) {
++ if (i % 2 == 1) {
++ id_grouper.set_group(i);
++ }
++ }
++
++ int k = 10;
++ int m = 8;
++ faiss::Index* index =
++ new faiss::IndexHNSWFlat(d, m, faiss::MetricType::METRIC_L2);
++ index->add(nb, xb); // add vectors to the index
++
++ // search
++ idx_t* I = new idx_t[k];
++ float* D = new float[k];
++
++ auto pSearchParameters = new faiss::SearchParametersHNSW();
++ pSearchParameters->grp = &id_grouper;
++
++ index->search(1, xb, k, D, I, pSearchParameters);
++
++ std::unordered_set group_ids;
++ ASSERT_EQ(0, I[0]);
++ ASSERT_EQ(0, D[0]);
++ group_ids.insert(id_grouper.get_group(I[0]));
++ for (int j = 1; j < 5; j++) {
++ ASSERT_NE(-1, I[j]);
++ ASSERT_NE(std::numeric_limits::max(), D[j]);
++ group_ids.insert(id_grouper.get_group(I[j]));
++ }
++ for (int j = 5; j < k; j++) {
++ ASSERT_EQ(-1, I[j]);
++ ASSERT_EQ(std::numeric_limits::max(), D[j]);
++ }
++ ASSERT_EQ(5, group_ids.size());
++
++ delete[] I;
++ delete[] D;
++ delete[] xb;
++}
++
++TEST(IdGrouper, bitmap_with_hnswn_idmap) {
++ int d = 1; // dimension
++ int nb = 10; // database size
++
++ std::mt19937 rng;
++ std::uniform_real_distribution<> distrib;
++
++ float* xb = new float[d * nb];
++ idx_t* xids = new idx_t[d * nb];
++
++ for (int i = 0; i < nb; i++) {
++ for (int j = 0; j < d; j++)
++ xb[d * i + j] = distrib(rng);
++ xb[d * i] += i / 1000.;
++ }
++
++ uint64_t bitmap[1] = {};
++ faiss::IDGrouperBitmap id_grouper(1, bitmap);
++ int num_grp = 0;
++ int grp_size = 2;
++ int id_in_grp = 0;
++ for (int i = 0; i < nb; i++) {
++ xids[i] = i + num_grp;
++ id_in_grp++;
++ if (id_in_grp == grp_size) {
++ id_grouper.set_group(i + num_grp + 1);
++ num_grp++;
++ id_in_grp = 0;
++ }
++ }
++
++ int k = 10;
++ int m = 8;
++ faiss::Index* index =
++ new faiss::IndexHNSWFlat(d, m, faiss::MetricType::METRIC_L2);
++ faiss::IndexIDMap id_map =
++ faiss::IndexIDMap(index); // add vectors to the index
++ id_map.add_with_ids(nb, xb, xids);
++
++ // search
++ idx_t* I = new idx_t[k];
++ float* D = new float[k];
++
++ auto pSearchParameters = new faiss::SearchParametersHNSW();
++ pSearchParameters->grp = &id_grouper;
++
++ id_map.search(1, xb, k, D, I, pSearchParameters);
++
++ std::unordered_set group_ids;
++ ASSERT_EQ(0, I[0]);
++ ASSERT_EQ(0, D[0]);
++ group_ids.insert(id_grouper.get_group(I[0]));
++ for (int j = 1; j < 5; j++) {
++ ASSERT_NE(-1, I[j]);
++ ASSERT_NE(std::numeric_limits::max(), D[j]);
++ group_ids.insert(id_grouper.get_group(I[j]));
++ }
++ for (int j = 5; j < k; j++) {
++ ASSERT_EQ(-1, I[j]);
++ ASSERT_EQ(std::numeric_limits::max(), D[j]);
++ }
++ ASSERT_EQ(5, group_ids.size());
++
++ delete[] I;
++ delete[] D;
++ delete[] xb;
++}
+--
+2.37.0
+
diff --git a/jni/patches/faiss/0002-Enable-precomp-table-to-be-shared-ivfpq.patch b/jni/patches/faiss/0002-Enable-precomp-table-to-be-shared-ivfpq.patch
new file mode 100644
index 000000000..88e6a8106
--- /dev/null
+++ b/jni/patches/faiss/0002-Enable-precomp-table-to-be-shared-ivfpq.patch
@@ -0,0 +1,512 @@
+From 9b33874562c9e62abf4a863657c54f0d349b0f67 Mon Sep 17 00:00:00 2001
+From: John Mazanec
+Date: Wed, 21 Feb 2024 15:34:15 -0800
+Subject: [PATCH] Enable precomp table to be shared ivfpq
+
+Changes IVFPQ and IVFPQFastScan indices to be able to share the
+precomputed table amongst other instances. Switches var to a pointer and
+add necessary functions to set them correctly.
+
+Adds a tests to validate the behavior.
+
+Signed-off-by: John Mazanec
+---
+ faiss/IndexIVFPQ.cpp | 47 +++++++-
+ faiss/IndexIVFPQ.h | 16 ++-
+ faiss/IndexIVFPQFastScan.cpp | 47 ++++++--
+ faiss/IndexIVFPQFastScan.h | 11 +-
+ tests/CMakeLists.txt | 1 +
+ tests/test_disable_pq_sdc_tables.cpp | 4 +-
+ tests/test_ivfpq_share_table.cpp | 173 +++++++++++++++++++++++++++
+ 7 files changed, 284 insertions(+), 15 deletions(-)
+ create mode 100644 tests/test_ivfpq_share_table.cpp
+
+diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
+index 100f499c..09508890 100644
+--- a/faiss/IndexIVFPQ.cpp
++++ b/faiss/IndexIVFPQ.cpp
+@@ -59,6 +59,29 @@ IndexIVFPQ::IndexIVFPQ(
+ polysemous_training = nullptr;
+ do_polysemous_training = false;
+ polysemous_ht = 0;
++ precomputed_table = new AlignedTable();
++ owns_precomputed_table = true;
++}
++
++IndexIVFPQ::IndexIVFPQ(const IndexIVFPQ& orig) : IndexIVF(orig), pq(orig.pq) {
++ code_size = orig.pq.code_size;
++ invlists->code_size = code_size;
++ is_trained = orig.is_trained;
++ by_residual = orig.by_residual;
++ use_precomputed_table = orig.use_precomputed_table;
++ scan_table_threshold = orig.scan_table_threshold;
++
++ polysemous_training = orig.polysemous_training;
++ do_polysemous_training = orig.do_polysemous_training;
++ polysemous_ht = orig.polysemous_ht;
++ precomputed_table = new AlignedTable(*orig.precomputed_table);
++ owns_precomputed_table = true;
++}
++
++IndexIVFPQ::~IndexIVFPQ() {
++ if (owns_precomputed_table) {
++ delete precomputed_table;
++ }
+ }
+
+ /****************************************************************
+@@ -464,11 +487,23 @@ void IndexIVFPQ::precompute_table() {
+ use_precomputed_table,
+ quantizer,
+ pq,
+- precomputed_table,
++ *precomputed_table,
+ by_residual,
+ verbose);
+ }
+
++void IndexIVFPQ::set_precomputed_table(
++ AlignedTable* _precompute_table,
++ int _use_precomputed_table) {
++ // Clean up old pre-computed table
++ if (owns_precomputed_table) {
++ delete precomputed_table;
++ }
++ owns_precomputed_table = false;
++ precomputed_table = _precompute_table;
++ use_precomputed_table = _use_precomputed_table;
++}
++
+ namespace {
+
+ #define TIC t0 = get_cycles()
+@@ -648,7 +683,7 @@ struct QueryTables {
+
+ fvec_madd(
+ pq.M * pq.ksub,
+- ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
++ ivfpq.precomputed_table->data() + key * pq.ksub * pq.M,
+ -2.0,
+ sim_table_2,
+ sim_table);
+@@ -677,7 +712,7 @@ struct QueryTables {
+ k >>= cpq.nbits;
+
+ // get corresponding table
+- const float* pc = ivfpq.precomputed_table.data() +
++ const float* pc = ivfpq.precomputed_table->data() +
+ (ki * pq.M + cm * Mf) * pq.ksub;
+
+ if (polysemous_ht == 0) {
+@@ -707,7 +742,7 @@ struct QueryTables {
+ dis0 = coarse_dis;
+
+ const float* s =
+- ivfpq.precomputed_table.data() + key * pq.ksub * pq.M;
++ ivfpq.precomputed_table->data() + key * pq.ksub * pq.M;
+ for (int m = 0; m < pq.M; m++) {
+ sim_table_ptrs[m] = s;
+ s += pq.ksub;
+@@ -727,7 +762,7 @@ struct QueryTables {
+ int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
+ k >>= cpq.nbits;
+
+- const float* pc = ivfpq.precomputed_table.data() +
++ const float* pc = ivfpq.precomputed_table->data() +
+ (ki * pq.M + cm * Mf) * pq.ksub;
+
+ for (int m = m0; m < m0 + Mf; m++) {
+@@ -1344,6 +1379,8 @@ IndexIVFPQ::IndexIVFPQ() {
+ do_polysemous_training = false;
+ polysemous_ht = 0;
+ polysemous_training = nullptr;
++ precomputed_table = new AlignedTable();
++ owns_precomputed_table = true;
+ }
+
+ struct CodeCmp {
+diff --git a/faiss/IndexIVFPQ.h b/faiss/IndexIVFPQ.h
+index d5d21da4..850bbe44 100644
+--- a/faiss/IndexIVFPQ.h
++++ b/faiss/IndexIVFPQ.h
+@@ -48,7 +48,8 @@ struct IndexIVFPQ : IndexIVF {
+
+ /// if use_precompute_table
+ /// size nlist * pq.M * pq.ksub
+- AlignedTable precomputed_table;
++ bool owns_precomputed_table;
++ AlignedTable* precomputed_table;
+
+ IndexIVFPQ(
+ Index* quantizer,
+@@ -58,6 +59,10 @@ struct IndexIVFPQ : IndexIVF {
+ size_t nbits_per_idx,
+ MetricType metric = METRIC_L2);
+
++ IndexIVFPQ(const IndexIVFPQ& orig);
++
++ ~IndexIVFPQ();
++
+ void encode_vectors(
+ idx_t n,
+ const float* x,
+@@ -139,6 +144,15 @@ struct IndexIVFPQ : IndexIVF {
+ /// build precomputed table
+ void precompute_table();
+
++ /**
++ * Initialize the precomputed table
++ * @param precompute_table
++ * @param _use_precomputed_table
++ */
++ void set_precomputed_table(
++ AlignedTable* precompute_table,
++ int _use_precomputed_table);
++
+ IndexIVFPQ();
+ };
+
+diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp
+index 2844ae49..895df342 100644
+--- a/faiss/IndexIVFPQFastScan.cpp
++++ b/faiss/IndexIVFPQFastScan.cpp
+@@ -46,6 +46,8 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(
+ : IndexIVFFastScan(quantizer, d, nlist, 0, metric), pq(d, M, nbits) {
+ by_residual = false; // set to false by default because it's faster
+
++ precomputed_table = new AlignedTable();
++ owns_precomputed_table = true;
+ init_fastscan(M, nbits, nlist, metric, bbs);
+ }
+
+@@ -53,6 +55,17 @@ IndexIVFPQFastScan::IndexIVFPQFastScan() {
+ by_residual = false;
+ bbs = 0;
+ M2 = 0;
++ precomputed_table = new AlignedTable();
++ owns_precomputed_table = true;
++}
++
++IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQFastScan& orig)
++ : IndexIVFFastScan(orig), pq(orig.pq) {
++ by_residual = orig.by_residual;
++ bbs = orig.bbs;
++ M2 = orig.M2;
++ precomputed_table = new AlignedTable(*orig.precomputed_table);
++ owns_precomputed_table = true;
+ }
+
+ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
+@@ -71,13 +84,15 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
+ ntotal = orig.ntotal;
+ is_trained = orig.is_trained;
+ nprobe = orig.nprobe;
++ precomputed_table = new AlignedTable();
++ owns_precomputed_table = true;
+
+- precomputed_table.resize(orig.precomputed_table.size());
++ precomputed_table->resize(orig.precomputed_table->size());
+
+- if (precomputed_table.nbytes() > 0) {
+- memcpy(precomputed_table.get(),
+- orig.precomputed_table.data(),
+- precomputed_table.nbytes());
++ if (precomputed_table->nbytes() > 0) {
++ memcpy(precomputed_table->get(),
++ orig.precomputed_table->data(),
++ precomputed_table->nbytes());
+ }
+
+ for (size_t i = 0; i < nlist; i++) {
+@@ -102,6 +117,12 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
+ orig_invlists = orig.invlists;
+ }
+
++IndexIVFPQFastScan::~IndexIVFPQFastScan() {
++ if (owns_precomputed_table) {
++ delete precomputed_table;
++ }
++}
++
+ /*********************************************************
+ * Training
+ *********************************************************/
+@@ -127,11 +148,23 @@ void IndexIVFPQFastScan::precompute_table() {
+ use_precomputed_table,
+ quantizer,
+ pq,
+- precomputed_table,
++ *precomputed_table,
+ by_residual,
+ verbose);
+ }
+
++void IndexIVFPQFastScan::set_precomputed_table(
++ AlignedTable* _precompute_table,
++ int _use_precomputed_table) {
++ // Clean up old pre-computed table
++ if (owns_precomputed_table) {
++ delete precomputed_table;
++ }
++ owns_precomputed_table = false;
++ precomputed_table = _precompute_table;
++ use_precomputed_table = _use_precomputed_table;
++}
++
+ /*********************************************************
+ * Code management functions
+ *********************************************************/
+@@ -229,7 +262,7 @@ void IndexIVFPQFastScan::compute_LUT(
+ if (cij >= 0) {
+ fvec_madd_simd(
+ dim12,
+- precomputed_table.get() + cij * dim12,
++ precomputed_table->get() + cij * dim12,
+ -2,
+ ip_table.get() + i * dim12,
+ tab);
+diff --git a/faiss/IndexIVFPQFastScan.h b/faiss/IndexIVFPQFastScan.h
+index 00dd2f11..91f35a6e 100644
+--- a/faiss/IndexIVFPQFastScan.h
++++ b/faiss/IndexIVFPQFastScan.h
+@@ -38,7 +38,8 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
+ /// precomputed tables management
+ int use_precomputed_table = 0;
+ /// if use_precompute_table size (nlist, pq.M, pq.ksub)
+- AlignedTable precomputed_table;
++ bool owns_precomputed_table;
++ AlignedTable* precomputed_table;
+
+ IndexIVFPQFastScan(
+ Index* quantizer,
+@@ -51,6 +52,10 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
+
+ IndexIVFPQFastScan();
+
++ IndexIVFPQFastScan(const IndexIVFPQFastScan& orig);
++
++ ~IndexIVFPQFastScan();
++
+ // built from an IndexIVFPQ
+ explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
+
+@@ -60,6 +65,10 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
+
+ /// build precomputed table, possibly updating use_precomputed_table
+ void precompute_table();
++ /// Pass in externally a precomputed
++ void set_precomputed_table(
++ AlignedTable* precompute_table,
++ int _use_precomputed_table);
+
+ /// same as the regular IVFPQ encoder. The codes are not reorganized by
+ /// blocks a that point
+diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
+index 87ab2020..a859516c 100644
+--- a/tests/CMakeLists.txt
++++ b/tests/CMakeLists.txt
+@@ -38,6 +38,7 @@ set(FAISS_TEST_SRC
+ test_common_ivf_empty_index.cpp
+ test_callback.cpp
+ test_utils.cpp
++ test_ivfpq_share_table.cpp
+ )
+
+ add_executable(faiss_test ${FAISS_TEST_SRC})
+diff --git a/tests/test_disable_pq_sdc_tables.cpp b/tests/test_disable_pq_sdc_tables.cpp
+index b211a5c4..a27973d5 100644
+--- a/tests/test_disable_pq_sdc_tables.cpp
++++ b/tests/test_disable_pq_sdc_tables.cpp
+@@ -15,7 +15,9 @@
+ #include "faiss/index_io.h"
+ #include "test_util.h"
+
+-pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
++namespace {
++ pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
++}
+
+ TEST(IO, TestReadHNSWPQ_whenSDCDisabledFlagPassed_thenDisableSDCTable) {
+ Tempfilename index_filename(&temp_file_mutex, "/tmp/faiss_TestReadHNSWPQ");
+diff --git a/tests/test_ivfpq_share_table.cpp b/tests/test_ivfpq_share_table.cpp
+new file mode 100644
+index 00000000..f827315d
+--- /dev/null
++++ b/tests/test_ivfpq_share_table.cpp
+@@ -0,0 +1,173 @@
++/**
++ * Copyright (c) Facebook, Inc. and its affiliates.
++ *
++ * This source code is licensed under the MIT license found in the
++ * LICENSE file in the root directory of this source tree.
++ */
++
++#include
++
++#include
++
++#include "faiss/Index.h"
++#include "faiss/IndexHNSW.h"
++#include "faiss/IndexIVFPQFastScan.h"
++#include "faiss/index_factory.h"
++#include "faiss/index_io.h"
++#include "test_util.h"
++
++namespace {
++ pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
++}
++
++std::vector generate_data(
++ int d,
++ int n,
++ std::default_random_engine rng,
++ std::uniform_real_distribution u) {
++ std::vector vectors(n * d);
++ for (size_t i = 0; i < n * d; i++) {
++ vectors[i] = u(rng);
++ }
++ return vectors;
++}
++
++void assert_float_vectors_almost_equal(
++ std::vector a,
++ std::vector b) {
++ float margin = 0.000001;
++ ASSERT_EQ(a.size(), b.size());
++ for (int i = 0; i < a.size(); i++) {
++ ASSERT_NEAR(a[i], b[i], margin);
++ }
++}
++
++/// Test case test precomputed table sharing for IVFPQ indices.
++template /// T represents class cast to use for index
++void test_ivfpq_table_sharing(
++ const std::string& index_description,
++ const std::string& filename,
++ faiss::MetricType metric) {
++ // Setup the index:
++ // 1. Build an index
++ // 2. ingest random data
++ // 3. serialize to disk
++ int d = 32, n = 1000;
++ std::default_random_engine rng(
++ std::chrono::system_clock::now().time_since_epoch().count());
++ std::uniform_real_distribution u(0, 100);
++
++ std::vector index_vectors = generate_data(d, n, rng, u);
++ std::vector query_vectors = generate_data(d, n, rng, u);
++
++ Tempfilename index_filename(&temp_file_mutex, filename);
++ {
++ std::unique_ptr index_writer(
++ faiss::index_factory(d, index_description.c_str(), metric));
++
++ index_writer->train(n, index_vectors.data());
++ index_writer->add(n, index_vectors.data());
++ faiss::write_index(index_writer.get(), index_filename.c_str());
++ }
++
++ // Load index from disk. Confirm that the sdc table is equal to 0 when
++ // disable sdc is set
++ std::unique_ptr> sharedAlignedTable(
++ new faiss::AlignedTable());
++ int shared_use_precomputed_table = 0;
++ int k = 10;
++ std::vector distances_test_a(k * n);
++ std::vector labels_test_a(k * n);
++ {
++ std::vector distances_baseline(k * n);
++ std::vector labels_baseline(k * n);
++
++ std::unique_ptr index_read_pq_table_enabled(
++ dynamic_cast(faiss::read_index(
++ index_filename.c_str(), faiss::IO_FLAG_READ_ONLY)));
++ std::unique_ptr index_read_pq_table_disabled(
++ dynamic_cast(faiss::read_index(
++ index_filename.c_str(),
++ faiss::IO_FLAG_READ_ONLY |
++ faiss::IO_FLAG_SKIP_PRECOMPUTE_TABLE)));
++ faiss::initialize_IVFPQ_precomputed_table(
++ shared_use_precomputed_table,
++ index_read_pq_table_disabled->quantizer,
++ index_read_pq_table_disabled->pq,
++ *sharedAlignedTable,
++ index_read_pq_table_disabled->by_residual,
++ index_read_pq_table_disabled->verbose);
++ index_read_pq_table_disabled->set_precomputed_table(
++ sharedAlignedTable.get(), shared_use_precomputed_table);
++
++ ASSERT_TRUE(index_read_pq_table_enabled->owns_precomputed_table);
++ ASSERT_FALSE(index_read_pq_table_disabled->owns_precomputed_table);
++ index_read_pq_table_enabled->search(
++ n,
++ query_vectors.data(),
++ k,
++ distances_baseline.data(),
++ labels_baseline.data());
++ index_read_pq_table_disabled->search(
++ n,
++ query_vectors.data(),
++ k,
++ distances_test_a.data(),
++ labels_test_a.data());
++
++ assert_float_vectors_almost_equal(distances_baseline, distances_test_a);
++ ASSERT_EQ(labels_baseline, labels_test_a);
++ }
++
++ // The precomputed table should only be set for L2 metric type
++ if (metric == faiss::METRIC_L2) {
++ ASSERT_EQ(shared_use_precomputed_table, 1);
++ } else {
++ ASSERT_EQ(shared_use_precomputed_table, 0);
++ }
++
++ // At this point, the original has gone out of scope, the destructor has
++ // been called. Confirm that initializing a new index from the table
++ // preserves the functionality.
++ {
++ std::vector distances_test_b(k * n);
++ std::vector labels_test_b(k * n);
++
++ std::unique_ptr index_read_pq_table_disabled(
++ dynamic_cast(faiss::read_index(
++ index_filename.c_str(),
++ faiss::IO_FLAG_READ_ONLY |
++ faiss::IO_FLAG_SKIP_PRECOMPUTE_TABLE)));
++ index_read_pq_table_disabled->set_precomputed_table(
++ sharedAlignedTable.get(), shared_use_precomputed_table);
++ ASSERT_FALSE(index_read_pq_table_disabled->owns_precomputed_table);
++ index_read_pq_table_disabled->search(
++ n,
++ query_vectors.data(),
++ k,
++ distances_test_b.data(),
++ labels_test_b.data());
++ assert_float_vectors_almost_equal(distances_test_a, distances_test_b);
++ ASSERT_EQ(labels_test_a, labels_test_b);
++ }
++}
++
++TEST(TestIVFPQTableSharing, L2) {
++ test_ivfpq_table_sharing(
++ "IVF16,PQ8x4", "/tmp/ivfpql2", faiss::METRIC_L2);
++}
++
++TEST(TestIVFPQTableSharing, IP) {
++ test_ivfpq_table_sharing(
++ "IVF16,PQ8x4", "/tmp/ivfpqip", faiss::METRIC_INNER_PRODUCT);
++}
++
++TEST(TestIVFPQTableSharing, FastScanL2) {
++ test_ivfpq_table_sharing(
++ "IVF16,PQ8x4fsr", "/tmp/ivfpqfsl2", faiss::METRIC_L2);
++}
++
++TEST(TestIVFPQTableSharing, FastScanIP) {
++ test_ivfpq_table_sharing(
++ "IVF16,PQ8x4fsr", "/tmp/ivfpqfsip", faiss::METRIC_INNER_PRODUCT);
++}
+--
+2.37.0
+
diff --git a/jni/patches/faiss/0003-Custom-patch-to-support-range-search-params.patch b/jni/patches/faiss/0003-Custom-patch-to-support-range-search-params.patch
new file mode 100644
index 000000000..bdc202bf6
--- /dev/null
+++ b/jni/patches/faiss/0003-Custom-patch-to-support-range-search-params.patch
@@ -0,0 +1,53 @@
+From af6770b505a32b2c4eab2036d2509dec4b137f28 Mon Sep 17 00:00:00 2001
+From: Junqiu Lei
+Date: Tue, 23 Apr 2024 17:18:56 -0700
+Subject: [PATCH] Custom patch to support range search params
+
+Signed-off-by: Junqiu Lei
+---
+ faiss/IndexIDMap.cpp | 28 ++++++++++++++++++++++++----
+ 1 file changed, 24 insertions(+), 4 deletions(-)
+
+diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp
+index 3f375e7b..11f3a847 100644
+--- a/faiss/IndexIDMap.cpp
++++ b/faiss/IndexIDMap.cpp
+@@ -176,11 +176,31 @@ void IndexIDMapTemplate::range_search(
+ RangeSearchResult* result,
+ const SearchParameters* params) const {
+ if (params) {
+- SearchParameters internal_search_parameters;
+- IDSelectorTranslated id_selector_translated(id_map, params->sel);
+- internal_search_parameters.sel = &id_selector_translated;
++ IDSelectorTranslated this_idtrans(this->id_map, nullptr);
++ ScopedSelChange sel_change;
++ IDGrouperTranslated this_idgrptrans(this->id_map, nullptr);
++ ScopedGrpChange grp_change;
++
++ if (params->sel) {
++ auto idtrans = dynamic_cast(params->sel);
++
++ if (!idtrans) {
++ auto params_non_const = const_cast(params);
++ this_idtrans.sel = params->sel;
++ sel_change.set(params_non_const, &this_idtrans);
++ }
++ }
++
++ if (params->grp) {
++ auto idtrans = dynamic_cast