diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index c4db6490..00000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,6 +0,0 @@ -[build] -rustdocflags = ["--document-private-items"] - -[target.'cfg(target_os="macos")'] -# Postgres symbols won't be available until runtime -rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup"] diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..4623fd97 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,15 @@ +.git +.github +build +target +**/target +TODO* +docs +tests +*.md +*.zip +*.deb +__pycache__ +**/__pycache__ +.pytest_cache +**/.pytest_cache diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..9cbb9b7c --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,8 @@ +version: 2 +updates: + - package-ecosystem: github-actions + directory: / + schedule: + interval: weekly + cooldown: + default-days: 7 diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml new file mode 100644 index 00000000..d8c794a4 --- /dev/null +++ b/.github/workflows/check.yml @@ -0,0 +1,1154 @@ +name: Check + +on: + push: + pull_request: + workflow_dispatch: + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + style: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -style')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -style')) || + github.event_name == 'workflow_dispatch' + + runs-on: "ubuntu-24.04" + + steps: + - name: Set up Environment + run: | + curl -fsSL https://github.com/tamasfe/taplo/releases/download/0.10.0/taplo-linux-$(uname -m).gz | gzip -d - | install -m 755 /dev/stdin /usr/local/bin/taplo + + curl -fsSL https://github.com/EmbarkStudios/cargo-deny/releases/download/0.18.9/cargo-deny-0.18.9-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - cargo-deny-0.18.9-$(uname -m)-unknown-linux-musl/cargo-deny | install -m 755 /dev/stdin /usr/local/bin/cargo-deny + + - name: Install Rust + run: rustup update + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Typos + uses: crate-ci/typos@3bc303c295c081add5df4a4d52a1e117f2fb2dce + + - name: Taplo + run: taplo fmt --check + + - name: Rustfmt + run: cargo fmt --check -- --config-path /dev/null --config imports_granularity=Module + + - name: Deny + run: | + cargo deny check + cargo deny --manifest-path services/tilemaxsimd/Cargo.toml check + + - name: License Header + run: | + UPSTREAM_HEADER=$(cat <> $GITHUB_ENV + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Cargo Test (Miri) + env: + RUSTFLAGS: "-Ctarget-cpu=sapphirerapids" + MIRIFLAGS: "-Zmiri-strict-provenance" + run: | + cargo miri test --locked --target x86_64-unknown-linux-gnu \ + --workspace --exclude vchord --no-fail-fast \ + -- --no-capture --test-threads=1 + + lint: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -lint')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -lint')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-24.04' || 'ubuntu-24.04-arm' }} + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo apt-get update + + if [ "$(uname -m)" == "x86_64" ]; then + wget https://downloadmirror.intel.com/859732/sde-external-9.58.0-2025-06-16-lin.tar.xz -O /tmp/sde-external.tar.xz + sudo tar -xf /tmp/sde-external.tar.xz -C /opt + sudo mv /opt/sde-external-9.58.0-2025-06-16-lin /opt/sde + fi + + if [ "$(uname -m)" == "aarch64" ]; then + sudo apt-get install -y qemu-user-static + fi + + - name: Install Rust + run: rustup update + + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Cargo Test + run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture + + - name: Cargo Test (QEMU) + run: | + if [ "$(uname -m)" == "x86_64" ]; then + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["/opt/sde/sde64", "-spr", "--"]' \ + test --locked -p simd -- --no-capture + fi + if [ "$(uname -m)" == "aarch64" ]; then + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=16"]' \ + test --locked -p simd -- --no-capture + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=32"]' \ + test --locked -p simd -- --no-capture + cargo \ + --config 'target.'\''cfg(all())'\''.runner = ["qemu-aarch64-static", "-cpu", "max,sve-default-vector-length=64"]' \ + test --locked -p simd -- --no-capture + fi + + - name: Clippy + run: cargo clippy --locked --workspace --exclude vchord + + psql: + if: | + (github.event_name == 'push' && !contains(github.event.head_commit.message, 'job: -psql')) || + (github.event_name == 'pull_request' && !contains(github.event.pull_request.body, 'job: -psql')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["14", "15", "16", "17", "18"] + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo apt-get update + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: rustup update + + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Install PostgreSQL & pgvector + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + + sudo apt-get install -y --no-install-recommends postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} + sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }} + + echo "local all all trust" | sudo tee /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all 127.0.0.1/32 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all ::1/128 trust" | sudo tee -a /etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + + sudo mkdir -p /etc/systemd/system/postgresql@.service.d/ + echo "[Service]" | sudo tee /etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitNOFILE=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf + echo "LimitMEMLOCK=infinity" | sudo tee -a /etc/systemd/system/postgresql@.service.d/limit.conf + sudo systemctl daemon-reload + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) + if [ "${{ matrix.version }}" -ge "18" ]; then + psql -c 'ALTER SYSTEM SET io_method = io_uring' + fi + psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo systemctl stop postgresql + + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + + sudo apt-get install -y --no-install-recommends postgresql-${{ matrix.version }}-pgvector + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Clippy + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build + + - name: Install + run: sudo make PG_CONFIG=$PG_CONFIG install + + - name: Upload Artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-gnu + path: ./build + compression-level: 9 + retention-days: 14 + + - name: Service + run: | + sudo systemctl start postgresql + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Release schema install and upgrade + if: matrix.version == '14' && matrix.arch == 'x86_64' + run: | + EXTENSION_DIR="$(pg_config --sharedir)/extension" + cmp sql/install/vchord--1.2.0.sql build/sharedir/extension/vchord--1.2.0.sql + sudo install -m 0644 sql/install/vchord--1.1.1.sql "$EXTENSION_DIR/" + sudo install -m 0644 sql/install/vchord--1.2.0.sql "$EXTENSION_DIR/" + sudo install -m 0644 sql/upgrade/vchord--1.1.1--1.2.0.sql "$EXTENSION_DIR/" + + dropdb --if-exists vchord_upgrade_test + createdb vchord_upgrade_test + psql -v ON_ERROR_STOP=1 -d vchord_upgrade_test <<'SQL' + CREATE EXTENSION vector; + CREATE EXTENSION vchord VERSION '1.1.1'; + CREATE TABLE items(id bigserial PRIMARY KEY, embedding vector(3)); + INSERT INTO items(embedding) + VALUES (ARRAY[1, 0, 0]::vector), (ARRAY[0, 1, 0]::vector); + CREATE INDEX items_vchord_idx + ON items USING vchordrq (embedding vector_l2_ops); + ALTER EXTENSION vchord UPDATE TO '1.2.0'; + SET enable_seqscan = off; + SELECT id FROM items + ORDER BY embedding <-> ARRAY[1, 0, 0]::vector + LIMIT 1; + SQL + test "$(psql -At -d vchord_upgrade_test -c \ + "SELECT extversion FROM pg_extension WHERE extname = 'vchord'")" = "1.2.0" + test "$(psql -At -d vchord_upgrade_test -c \ + "SELECT to_regprocedure('vchordrq_tilemaxsim_rerank(regclass,vector[],bigint[],integer)') IS NOT NULL")" = "t" + + dropdb --if-exists vchord_fresh_install_test + createdb vchord_fresh_install_test + psql -v ON_ERROR_STOP=1 -d vchord_fresh_install_test <<'SQL' + CREATE EXTENSION vector; + CREATE EXTENSION vchord VERSION '1.2.0'; + SQL + test "$(psql -At -d vchord_fresh_install_test -c \ + "SELECT extversion FROM pg_extension WHERE extname = 'vchord'")" = "1.2.0" + test "$(psql -At -d vchord_fresh_install_test -c \ + "SELECT to_regprocedure('vchordrq_tilemaxsim_rerank(regclass,vector[],bigint[],integer)') IS NOT NULL")" = "t" + + - name: Logging + if: always() + run: | + cat /var/log/postgresql/postgresql-${{ matrix.version }}-main.log + + psql_macos: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_macos')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_macos')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["14", "15", "16", "17", "18"] + arch: ["aarch64", "x86_64"] + + runs-on: ${{ matrix.arch == 'aarch64' && 'macos-15' || 'macos-15-intel' }} + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + brew update + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-${{ matrix.arch }}-apple-darwin.tar.gz | tar -xOzf - ./sqllogictest | tee /usr/local/bin/sqllogictest > /dev/null && chmod 755 /usr/local/bin/sqllogictest + + - name: Install Rust + run: rustup update + + - name: Install Clang + run: echo CC=$(brew --prefix llvm@18)/bin/clang >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + with: + version: "v0.12.0" + + - name: Install PostgreSQL & pgvector + run: | + brew install postgresql@${{ matrix.version }} + brew services start postgresql@${{ matrix.version }} + for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done + [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" + + $(brew --prefix postgresql@${{ matrix.version }})/bin/createdb -O $(whoami) $(whoami) + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + brew services stop postgresql@${{ matrix.version }} + + $(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config + echo PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config >> $GITHUB_ENV + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=$(brew --prefix postgresql@${{ matrix.version }})/bin/pg_config CC='sccache clang' OPTFLAGS="" install + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Clippy + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build + + - name: Install + run: sudo make PG_CONFIG=$PG_CONFIG install + + - name: Upload Artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-apple-darwin + path: ./build + compression-level: 9 + retention-days: 14 + + - name: Service + run: | + brew services start postgresql@${{ matrix.version }} + for i in {1..60}; do [ -S /tmp/.s.PGSQL.5432 ] && echo "PostgreSQL ready" && break || sleep 1; done + [ -S /tmp/.s.PGSQL.5432 ] || echo "PostgreSQL socket not found after 60 seconds" + $(brew --prefix postgresql@${{ matrix.version }})/bin/psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Logging + if: always() + run: | + cat $(brew --prefix)/var/log/postgresql@${{ matrix.version }}.log + + psql_windows: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_windows')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_windows')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["14", "15", "16", "17", "18"] + arch: ["x86_64"] + + runs-on: "windows-2022" + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + Invoke-WebRequest -Uri https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-${{ matrix.arch }}-pc-windows-msvc.zip -OutFile "$env:TEMP\sqllogictest-install.zip" + Expand-Archive -Path "$env:TEMP\sqllogictest-install.zip" -DestinationPath "D:\sqllogictest-install" -Force + Add-Content -Path $env:GITHUB_PATH -Value "D:\sqllogictest-install" + + - name: Install Rust + run: rustup update + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Install PostgreSQL & pgvector + run: | + 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } + + if ( "${{ matrix.version }}" -eq "14" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-14.20-2-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "15" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-15.15-2-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "16" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-16.11-2-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "17" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-17.7-2-windows-x64-binaries.zip" + } + if ( "${{ matrix.version }}" -eq "18" ) { + $postgresql_url = "https://get.enterprisedb.com/postgresql/postgresql-18.1-2-windows-x64-binaries.zip" + } + Invoke-WebRequest -Uri $postgresql_url -OutFile "$env:TEMP\postgresql-install.zip" + Expand-Archive -Path "$env:TEMP\postgresql-install.zip" -DestinationPath "D:\postgresql-install" -Force + D:\postgresql-install\pgsql\bin\initdb.exe -D D:\postgresql-install\pgsql\data -U postgres + + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data -l D:\postgresql-install\postgresql.log + + D:\postgresql-install\pgsql\bin\createuser.exe -U postgres -s -r $env:USERNAME + D:\postgresql-install\pgsql\bin\createdb.exe -O $env:USERNAME $env:USERNAME + D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET max_worker_processes = 1024' + D:\postgresql-install\pgsql\bin\psql.exe -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + D:\postgresql-install\pgsql\bin\pg_ctl.exe stop -D D:\postgresql-install\pgsql\data + + D:\postgresql-install\pgsql\bin\pg_config.exe + Add-Content -Path $env:GITHUB_ENV -Value "PG_CONFIG=D:\postgresql-install\pgsql\bin\pg_config.exe" + + Invoke-WebRequest -Uri "https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.zip" -OutFile "$env:TEMP\pgvector-install.zip" + Expand-Archive -Path "$env:TEMP\pgvector-install.zip" -DestinationPath "D:\pgvector-install" -Force + & 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Launch-VsDevShell.ps1' -HostArch amd64 -Arch amd64 + Push-Location -Path D:\pgvector-install\pgvector-0.8.1 + nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" + nmake /F Makefile.win PGROOT="D:\postgresql-install\pgsql" install + Pop-Location + + - name: Install Clang + run: Add-Content -Path $env:GITHUB_ENV -Value "CC=clang-cl.exe" + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Clippy + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build + + - name: Install + run: make install + + - name: Upload Artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-pc-windows-msvc + path: ./build + compression-level: 9 + retention-days: 14 + + - name: Service + run: | + 'PGBIN','PGDATA','PGROOT', 'PGUSER', 'PGPASSWORD' | ForEach-Object { Remove-Item "env:$_" } + + D:\postgresql-install\pgsql\bin\pg_ctl.exe start -D D:\postgresql-install\pgsql\data -l D:\postgresql-install\postgresql.log + D:\postgresql-install\pgsql\bin\psql.exe -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if ( "${{ matrix.version }}" -ge "17" ) { + sqllogictest --db $env:USERNAME --user $env:USERNAME './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + } + + - name: Logging + if: always() + run: | + Get-Content -Path 'D:\postgresql-install\postgresql.log' + + psql_alpine: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +psql_alpine')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +psql_alpine')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["15", "16", "17"] + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-24.04' || 'ubuntu-24.04-arm' }} + container: + image: alpine:3.22 + volumes: + - /opt:/opt:rw,rshared + - /opt:/__e/node20:ro,rshared + - /opt:/__e/node24:ro,rshared + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + # RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sed -i "/^ID=/s/alpine/NotpineForGHA/" /etc/os-release + apk add --update-cache curl git make nodejs sudo zip + mkdir /opt/bin + ln -s /usr/bin/node /opt/bin/node + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: | + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo $HOME/.cargo/bin >> $GITHUB_PATH + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + + - name: Install Clang + run: | + sudo apk add --no-cache clang18 clang18-dev + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Install PostgreSQL & pgvector + run: | + sudo apk add --no-cache postgresql${{ matrix.version }} postgresql${{ matrix.version }}-dev + sudo touch /var/log/postgresql.log + sudo chmod 644 /var/log/postgresql.log + sudo chown postgres /var/log/postgresql.log + sudo install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql + sudo install --verbose --directory --owner postgres --group postgres --mode 1777 /var/lib/postgresql/data + sudo install --verbose --directory --owner postgres --group postgres --mode 3777 /run/postgresql + sudo -iu postgres initdb -D /var/lib/postgresql/data + + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) + if [ "${{ matrix.version }}" -ge "18" ]; then + psql -c 'ALTER SYSTEM SET io_method = io_uring' + fi + psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo -iu postgres pg_ctl stop -D /var/lib/postgresql/data + + /usr/libexec/postgresql${{ matrix.version }}/pg_config + echo PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config >> $GITHUB_ENV + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' OPTFLAGS="" + make -C ~/pgvector-install/pgvector-0.8.1 PG_CONFIG=/usr/libexec/postgresql${{ matrix.version }}/pg_config CC='sccache clang-18' OPTFLAGS="" install + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Patch + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + cat << EOF > ./.cargo/config.toml + unstable.host-config = true + unstable.target-applies-to-host = true + host.rustflags = ["-Dwarnings", "-Ctarget-feature=-crt-static"] + build.rustflags = ["-Dwarnings", "-Ctarget-feature=-crt-static"] + EOF + + - name: Clippy + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build + + - name: Install + run: sudo make PG_CONFIG=$PG_CONFIG install + + - name: Upload Artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.arch }}-linux-musl + path: ./build + compression-level: 9 + retention-days: 14 + + - name: Service + run: | + sudo -iu postgres pg_ctl start -D /var/lib/postgresql/data -l /var/log/postgresql.log + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Logging + if: always() + run: | + cat /var/log/postgresql.log + + check_debian: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +check_debian')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +check_debian')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + include: + - version: "15" + platform: "s390x" + clang_triple: "s390x-linux-gnu" + rust_triple: "s390x-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "s390x" + clang_triple: "s390x-linux-gnu" + rust_triple: "s390x-unknown-linux-gnu" + gcc_version: "14" + debian_version: "trixie" + - version: "14" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "15" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "16" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "18" + platform: "ppc64le" + clang_triple: "powerpc64le-linux-gnu" + rust_triple: "powerpc64le-unknown-linux-gnu" + gcc_version: "12" + debian_version: "bookworm" + - version: "17" + platform: "riscv64" + clang_triple: "riscv64-linux-gnu" + rust_triple: "riscv64gc-unknown-linux-gnu" + gcc_version: "14" + debian_version: "trixie" + runs-on: "ubuntu-24.04" + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + # RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo apt-get update + sudo apt-get install -y qemu-user-static + sudo systemctl enable --now systemd-binfmt + sudo touch /etc/sudoers.d/qemu + echo "Defaults env_keep += \"QEMU_CPU\"" | sudo tee -a /etc/sudoers.d/qemu + sudo mkdir /sysroot + curl -fsSL https://github.com/debuerreotype/docker-debian-artifacts/raw/refs/heads/dist-${{ matrix.platform }}/${{ matrix.debian_version }}/oci/blobs/rootfs.tar.gz | sudo tar -xz -C /sysroot + sudo mount --bind /dev /sysroot/dev + sudo mount --bind /dev/pts /sysroot/dev/pts + sudo mount --bind /etc/resolv.conf /sysroot/etc/resolv.conf + sudo mount --bind /proc /sysroot/proc + sudo mount --bind /sys /sysroot/sys + sudo mount --bind /tmp /sysroot/tmp + sudo chroot /sysroot apt-get update + sudo chroot /sysroot apt-get install --no-install-recommends -y libc6-dev libgcc-${{ matrix.gcc_version }}-dev + sudo chroot /sysroot apt-get install --no-install-recommends -y ca-certificates sudo + QEMU_LD_PREFIX=/sysroot + QEMU_CPU=max + if [ "${{ matrix.platform }}" = "ppc64le" ]; then + sudo rm -f /sysroot/lib64/ld64.so.2 + sudo ln /sysroot/usr/lib/powerpc64le-linux-gnu/ld64.so.2 /sysroot/lib64/ld64.so.2 + sudo rm -f /sysroot/usr/lib/powerpc64le-linux-gnu/libm.so + sudo ln /sysroot/lib/powerpc64le-linux-gnu/libm.so.6 /sysroot/usr/lib/powerpc64le-linux-gnu/libm.so + QEMU_CPU=power10 + fi + export QEMU_LD_PREFIX + export QEMU_CPU + echo QEMU_LD_PREFIX=$QEMU_LD_PREFIX >> $GITHUB_ENV + echo QEMU_CPU=$QEMU_CPU >> $GITHUB_ENV + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: | + rustup default beta + rustup component add rustfmt clippy + rustup target add ${{ matrix.rust_triple }} + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Install PostgreSQL & pgvector + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-common + sudo chroot /sysroot /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} + sudo chroot /sysroot apt-get install -y --no-install-recommends postgresql-${{ matrix.version }} + + echo "local all all trust" | sudo tee /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all 127.0.0.1/32 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + echo "host all all ::1/128 trust" | sudo tee -a /sysroot/etc/postgresql/${{ matrix.version }}/main/pg_hba.conf + + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start + + sudo chroot /sysroot sudo -iu postgres createuser -s -r $(whoami) + sudo chroot /sysroot sudo -iu postgres createdb -O $(whoami) $(whoami) + if [ "${{ matrix.version }}" -ge "18" ]; then + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = worker' + fi + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + sudo chroot /sysroot sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main stop + + sudo touch /usr/bin/pg_config + echo "#!/usr/bin/env bash" | sudo tee -a /usr/bin/pg_config + echo "sudo chroot /sysroot pg_config \"\$@\" \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^/(.*)$|/sysroot/\1|' \\" | sudo tee -a /usr/bin/pg_config + echo " | sed -E 's|^([A-Z]+) = /(.*)$|\1 = /sysroot/\2|'" | sudo tee -a /usr/bin/pg_config + sudo chmod 755 /usr/bin/pg_config + + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" + sudo make -C ~/pgvector-install/pgvector-0.8.1 with_llvm=no CC='sccache clang' CFLAGS="-fuse-ld=lld --target=${{ matrix.clang_triple }} --sysroot=/sysroot" OPTFLAGS="" install + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Patch + run: | + mkdir ./.cargo && touch ./.cargo/config.toml + cat << EOF > ./.cargo/config.toml + [target.${{ matrix.rust_triple }}] + runner = ["qemu-${{ matrix.platform }}-static"] + linker = "clang" + rustflags = [ + "-Clink-arg=-fuse-ld=lld", + "-Clink-arg=--target=${{ matrix.clang_triple }}", + "-Clink-arg=--sysroot=/sysroot", + "-Dwarnings", + ] + [env] + CFLAGS_${{ matrix.rust_triple }} = "--target=${{ matrix.clang_triple }} --sysroot=/sysroot" + BINDGEN_EXTRA_CLANG_ARGS_${{ matrix.rust_triple }} = "--sysroot=/sysroot" + EOF + + - name: Cargo Test + run: cargo test --locked --target ${{ matrix.rust_triple }} --workspace --exclude vchord --no-fail-fast -- --no-capture + + - name: Clippy + run: make TARGET=${{ matrix.rust_triple }} PROFILE=dev RUNNER='qemu-${{ matrix.platform }}-static' clippy + + - name: Build + run: make TARGET=${{ matrix.rust_triple }} PROFILE=dev RUNNER='qemu-${{ matrix.platform }}-static' build + + - name: Install + run: sudo make PG_CONFIG=$PG_CONFIG install + + - name: Upload Artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: postgresql-${{ matrix.version }}-vchord_0.0.0_${{ matrix.clang_triple }} + path: ./build + compression-level: 9 + retention-days: 14 + + - name: Service + run: | + sudo chroot /sysroot pg_ctlcluster ${{ matrix.version }} main start + sudo chroot /sysroot psql -d $(whoami) -U $(whoami) -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Logging + if: always() + run: | + cat /sysroot/var/log/postgresql/postgresql-${{ matrix.version }}-main.log + + check_arch: + if: | + (github.event_name == 'push' && contains(github.event.head_commit.message, 'job: +check_arch')) || + (github.event_name == 'pull_request' && contains(github.event.pull_request.body, 'job: +check_arch')) || + github.event_name == 'workflow_dispatch' + + strategy: + matrix: + version: ["17", "18"] + + runs-on: "ubuntu-24.04" + + container: + image: archlinux/archlinux:base-devel + volumes: + - /etc/passwd:/etc/passwd:rw,rshared + - /etc/group:/etc/group:rw,rshared + - /etc/shadow:/etc/shadow:rw,rshared + - /etc/gshadow:/etc/gshadow:rw,rshared + - /etc/sudoers.d/runner:/etc/sudoers.d/runner:rw,rshared + options: --user 1001 --privileged + + env: + SCCACHE_GHA_ENABLED: "true" + RUSTUP_AUTO_INSTALL: "0" + RUSTC_WRAPPER: "sccache" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo sed -i 's/^DownloadUser = alpm$/DownloadUser = runner/' /etc/pacman.conf + sudo pacman -Syu --ignore systemd,systemd-libs,systemd-sysvcompat --noconfirm git nodejs + sudo mkdir -p /var/lib/postgresql + sudo chmod 700 /var/lib/postgresql + sudo chown postgres:postgres /var/lib/postgresql + + sudo pacman -S --noconfirm liburing numactl python valgrind + + curl -fsSL https://github.com/risinglightdb/sqllogictest-rs/releases/download/v0.29.0/sqllogictest-bin-v0.29.0-$(uname -m)-unknown-linux-musl.tar.gz | tar -xOzf - ./sqllogictest | sudo install -m 755 /dev/stdin /usr/local/bin/sqllogictest + + - name: Install Rust + run: | + curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain beta + echo $HOME/.cargo/bin >> $GITHUB_PATH + echo RUSTC_BOOTSTRAP=1 >> $GITHUB_ENV + + - name: Install Clang + run: | + sudo pacman -S --noconfirm clang + echo CC=clang >> $GITHUB_ENV + + - name: Set up Sccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + + - name: Install PostgreSQL & pgvector + run: | + mkdir ~/postgresql-install + if [ "${{ matrix.version }}" = "17" ]; then + curl -fsSL https://ftp.postgresql.org/pub/source/v17.6/postgresql-17.6.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-17.6 + fi + if [ "${{ matrix.version }}" = "18" ]; then + curl -fsSL https://ftp.postgresql.org/pub/source/v18.0/postgresql-18.0.tar.bz2 | tar -xj -C ~/postgresql-install + pushd ~/postgresql-install/postgresql-18.0 + fi + ./configure \ + --prefix=/usr \ + --sysconfdir=/etc \ + --mandir=/usr/share/man \ + --datadir=/usr/share/postgresql \ + --disable-rpath \ + --enable-nls \ + --with-gssapi \ + --with-icu \ + --with-ldap \ + --with-lz4 \ + --with-openssl \ + --with-python \ + --with-readline \ + --with-system-tzdata=/usr/share/zoneinfo \ + --with-uuid=e2fs \ + --with-zstd \ + ${{ matrix.version >= 18 && '--with-libcurl' || '' }} \ + ${{ matrix.version >= 18 && '--with-libnuma' || '' }} \ + ${{ matrix.version >= 18 && '--with-liburing' || '' }} \ + --enable-debug \ + --enable-cassert \ + CC='sccache clang' \ + CFLAGS='-O0 -g -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -fasynchronous-unwind-tables' \ + CPPFLAGS='-DUSE_ASSERT_CHECKING=1 -DRANDOMIZE_ALLOCATED_MEMORY=1 -DUSE_VALGRIND=1' + make all -j$(nproc) + sudo make install + popd + + sudo touch /usr/bin/postmaster + echo '#!/usr/bin/sh' | sudo tee -a /usr/bin/postmaster + echo 'export DEBUGINFOD_URLS=https://debuginfod.archlinux.org' | sudo tee -a /usr/bin/postmaster + echo 'exec \' | sudo tee -a /usr/bin/postmaster + echo ' valgrind \' | sudo tee -a /usr/bin/postmaster + echo ' --leak-check=no \' | sudo tee -a /usr/bin/postmaster + echo ' --gen-suppressions=all \' | sudo tee -a /usr/bin/postmaster + echo ' --time-stamp=yes \' | sudo tee -a /usr/bin/postmaster + echo ' --error-markers=VALGRINDERROR-BEGIN,VALGRINDERROR-END \' | sudo tee -a /usr/bin/postmaster + echo ' --trace-children=yes \' | sudo tee -a /usr/bin/postmaster + echo ' /usr/bin/postgres "$@"' | sudo tee -a /usr/bin/postmaster + sudo chmod 755 /usr/bin/postmaster + sudo mkdir -p /var/lib/postgres + sudo chmod 700 /var/lib/postgres + sudo chown postgres:postgres /var/lib/postgres + sudo touch /var/log/postgresql.log + sudo chmod 700 /var/log/postgresql.log + sudo chown postgres:postgres /var/log/postgresql.log + sudo -iu postgres initdb -D /var/lib/postgres/data + sudo -iu postgres pg_ctl start -D /var/lib/postgres/data -l /var/log/postgresql.log -p /usr/bin/postmaster + + sudo -iu postgres createuser -s -r $(whoami) + sudo -iu postgres createdb -O $(whoami) $(whoami) + if [ "${{ matrix.version }}" -ge "18" ]; then + # valgrind works not well + sudo -iu postgres psql -c 'ALTER SYSTEM SET io_method = worker' + fi + sudo -iu postgres psql -c 'ALTER SYSTEM SET max_worker_processes = 1024' + sudo -iu postgres psql -c 'ALTER SYSTEM SET shared_preload_libraries = "vchord"' + sudo -iu postgres pg_ctl stop -D /var/lib/postgres/data + + pg_config + echo PG_CONFIG=pg_config >> $GITHUB_ENV + + mkdir ~/pgvector-install + curl -fsSL https://github.com/pgvector/pgvector/archive/refs/tags/v0.8.1.tar.gz | tar -xz -C ~/pgvector-install + pushd ~/pgvector-install/pgvector-0.8.1 + make OPTFLAGS="" -j$(nproc) + sudo make OPTFLAGS="" install + popd + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Cargo Test + run: cargo test --locked --workspace --exclude vchord --no-fail-fast -- --no-capture + + - name: Clippy + run: make PROFILE=dev clippy + + - name: Build + run: make PROFILE=dev build + + - name: Install + run: sudo make PG_CONFIG=$PG_CONFIG install + + - name: Service + run: | + sudo -iu postgres pg_ctl start -D /var/lib/postgres/data -l /var/log/postgresql.log -p /usr/bin/postmaster + psql -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;' + + - name: Sqllogictest + run: | + sqllogictest --db $(whoami) --user $(whoami) './tests/general/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordg/*.slt' --label pg${{ matrix.version }} + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/*.slt' --label pg${{ matrix.version }} + if [ "${{ matrix.version }}" -ge "17" ]; then + sqllogictest --db $(whoami) --user $(whoami) './tests/vchordrq/pg17/*.slt' --label pg${{ matrix.version }} + fi + + - name: Logging + if: always() + run: | + sudo cat /var/log/postgresql.log diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 00000000..f701f7b1 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,241 @@ +name: Publish + +on: + release: + types: [created] + workflow_dispatch: + inputs: + tag: + description: "tag name (semver without v-prefix)" + required: true + type: string + +concurrency: + group: ${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +permissions: + contents: write + packages: write + +jobs: + semver: + runs-on: "ubuntu-latest" + + steps: + - name: Semver + id: semver + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + with: + script: | + const tag = "${{ github.event.inputs.tag }}" || "${{ github.event.release.tag_name }}"; + console.log(`Tag: ${tag}`); + const r = /^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$/; + if (!r.test(tag)) { + core.setFailed(`Action failed with an invalid semver.`); + } + core.setOutput('SEMVER', tag); + + outputs: + SEMVER: ${{ steps.semver.outputs.SEMVER }} + + build: + needs: ["semver"] + strategy: + matrix: + version: ["14", "15", "16", "17", "18"] + arch: ["x86_64", "aarch64"] + runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} + + env: + RUSTUP_AUTO_INSTALL: "0" + RUSTFLAGS: "-Dwarnings" + CARGO_TERM_COLOR: "always" + RUST_BACKTRACE: "1" + + steps: + - name: Set up Environment + run: | + sudo apt-get update + + - name: Install Rust + run: rustup update + + - name: Install Clang + run: | + curl --proto '=https' --tlsv1.2 -sSf https://apt.llvm.org/llvm.sh | sudo bash -s -- 18 + echo CC=clang-18 >> $GITHUB_ENV + + - name: Install PostgreSQL + run: | + sudo apt-get remove -y '^postgres.*' '^libpq.*' + sudo apt-get purge -y '^postgres.*' '^libpq.*' + + sudo apt-get install -y --no-install-recommends postgresql-common + sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y + sudo apt-get install -y --no-install-recommends postgresql-server-dev-${{ matrix.version }} + + pg_config + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Build + run: | + grep -q "default_version = '${{ needs.semver.outputs.SEMVER }}'" vchord.control || exit 1 + make build + + - name: Package (binary package) + env: + GH_TOKEN: ${{ github.token }} + SEMVER: ${{ needs.semver.outputs.SEMVER }} + VERSION: ${{ matrix.version }} + run: | + ARCH=$(uname -m) + (cd ./build && zip -r ~/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip .) + gh release upload --clobber $SEMVER ~/postgresql-${VERSION}-vchord_${SEMVER}_${ARCH}-linux-gnu.zip + + - name: Package (debian package) + env: + GH_TOKEN: ${{ github.token }} + SEMVER: ${{ needs.semver.outputs.SEMVER }} + VERSION: ${{ matrix.version }} + run: | + make DESTDIR="~/debian" install + mkdir -p ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/ + echo "Format: https://www.debian.org/doc/packaging-manuals/copyright-format/1.0/ + Upstream-Name: VectorChord + Source: https://github.com/HuXinjing/VectorChord + + Files: * + Copyright: Copyright (c) 2025-2026 TensorChord Inc. + Copyright (c) 2026 Hu Xinjing + License: AGPL-3.0-only OR Elastic-2.0 + + License: AGPL-3.0-only + $(cat ./licenses/LICENSE.AGPLv3 | sed 's/^$/./' | sed 's/^/ /') + + License: Elastic-2.0 + $(cat ./licenses/LICENSE.ELv2 | sed 's/^$/./' | sed 's/^/ /') + " \ + >> ~/debian/usr/share/doc/postgresql-${VERSION}-vchord/copyright + ARCH=$(uname -m) + PLATFORM=$(dpkg --print-architecture) + mkdir -p ~/debian/DEBIAN + echo "Package: postgresql-${VERSION}-vchord + Version: ${SEMVER}-1 + Depends: postgresql-${VERSION}, libgcc-s1, libc6 (>= 2.35) + Section: database + Priority: optional + Architecture: ${PLATFORM} + Maintainer: Hu Xinjing + Description: Scalable, fast, and disk-friendly vector search in Postgres + VectorChord is a PostgreSQL extension designed for scalable, + high-performance, and disk-efficient vector similarity search. It + supports L2 distance, inner product, cosine distance and maxsim. + Homepage: https://github.com/HuXinjing/VectorChord" \ + > ~/debian/DEBIAN/control + (cd ~/debian && find usr -type f -print0 | xargs -0 md5sum) > ~/debian/DEBIAN/md5sums + dpkg-deb --root-owner-group -Zxz --build ~/debian ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + gh release upload --clobber $SEMVER ~/postgresql-${VERSION}-vchord_${SEMVER}-1_${PLATFORM}.deb + + tilemaxsimd: + runs-on: ubuntu-24.04 + needs: ["semver", "build"] + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Free space for the CUDA build image + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /usr/local/.ghcup + docker builder prune --all --force + + - name: Log in to the fork container registry + env: + GH_TOKEN: ${{ github.token }} + run: echo "$GH_TOKEN" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Build and publish native TileMaxSim image + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + run: | + IMAGE=$(echo "ghcr.io/${GITHUB_REPOSITORY_OWNER}/vectorchord-tilemaxsimd" | tr '[:upper:]' '[:lower:]') + docker build --target test -f services/Dockerfile.tilemaxsimd . + docker build -f services/Dockerfile.tilemaxsimd -t "${IMAGE}:${SEMVER}" . + docker run --rm -v /var/run/docker.sock:/var/run/docker.sock \ + aquasec/trivy:0.72.0@sha256:cffe3f5161a47a6823fbd23d985795b3ed72a4c806da4c4df16266c02accdd6f \ + image --scanners vuln --severity HIGH,CRITICAL --ignore-unfixed \ + --exit-code 1 "${IMAGE}:${SEMVER}" + docker push "${IMAGE}:${SEMVER}" + + postgres-image: + runs-on: ubuntu-24.04 + needs: ["semver", "build"] + strategy: + matrix: + version: ["16"] + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Log in to the fork container registry + env: + GH_TOKEN: ${{ github.token }} + run: echo "$GH_TOKEN" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Build and publish VectorChord PostgreSQL image + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + run: | + IMAGE=$(echo "ghcr.io/${GITHUB_REPOSITORY_OWNER}/vectorchord-postgres" | tr '[:upper:]' '[:lower:]') + docker build \ + --build-arg POSTGRES_VERSION=${{ matrix.version }} \ + -f services/Dockerfile.postgres \ + -t "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" \ + -t "${IMAGE}:${{ matrix.version }}-${SEMVER}" . + CID=$(docker run -d \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" \ + -c shared_preload_libraries=vchord) + trap 'docker rm -f "$CID" >/dev/null 2>&1 || true' EXIT + for attempt in $(seq 1 60); do + if docker exec "$CID" pg_isready -U postgres >/dev/null 2>&1; then break; fi + sleep 1 + done + docker exec "$CID" psql -v ON_ERROR_STOP=1 -U postgres -d postgres \ + -c 'CREATE EXTENSION vchord CASCADE' \ + -c "SELECT extversion FROM pg_extension WHERE extname = 'vchord'" + docker run --rm -v /var/run/docker.sock:/var/run/docker.sock \ + aquasec/trivy:0.72.0@sha256:cffe3f5161a47a6823fbd23d985795b3ed72a4c806da4c4df16266c02accdd6f \ + image --scanners vuln --severity HIGH,CRITICAL --ignore-unfixed \ + --exit-code 1 "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" + docker push "${IMAGE}:pg${{ matrix.version }}-${SEMVER}" + docker push "${IMAGE}:${{ matrix.version }}-${SEMVER}" + + pgxn: + if: ${{ vars.PUBLISH_PGXN == 'true' }} + runs-on: "ubuntu-latest" + needs: ["semver", "build"] + + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Upload + env: + SEMVER: ${{ needs.semver.outputs.SEMVER }} + PGXN_USERNAME: ${{ secrets.PGXN_USERNAME }} + PGXN_PASSWORD: ${{ secrets.PGXN_PASSWORD }} + run: | + mkdir -p ./build + + sed -e "s/@DISTVERSION@/${SEMVER}/g" META.json.in > META.json + git archive --format zip --prefix "vchord-${SEMVER}/" --add-file META.json -o "./build/vchord--${SEMVER}.zip" HEAD + + curl --fail -sS \ + --user "${PGXN_USERNAME}:${PGXN_PASSWORD}" \ + -F "submit=Release It!" \ + -F "archive=@./build/vchord--${SEMVER}.zip" \ + -H "X-Requested-With: XMLHttpRequest" \ + https://manager.pgxn.org/upload diff --git a/.gitignore b/.gitignore index 85d1619b..d2f22b6e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,8 @@ /target +/services/tilemaxsimd/target +/benchmarks/ +**/__pycache__/ +.ruff_cache/ **/*.rs.bk .vscode .ignore diff --git a/Cargo.lock b/Cargo.lock index 97654851..a2e604b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,101 +1,124 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] -name = "ahash" -version = "0.7.8" +name = "adler2" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "891477e0c6a8957309ee5c45a6368af3ae14bb510732d2684ffa19af310920f9" -dependencies = [ - "getrandom", - "once_cell", - "version_check", -] +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ "memchr", ] [[package]] -name = "annotate-snippets" -version = "0.9.2" +name = "alloca" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccaf7e9dfbb6ab22c82e473cd1a8a7bd313c19a5b7e40970f3d89ef5a5c9e81e" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" dependencies = [ - "unicode-width", - "yansi-term", + "cc", ] [[package]] -name = "anyhow" -version = "1.0.86" +name = "always_equal" +version = "0.0.0" + +[[package]] +name = "anes" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" [[package]] -name = "approx" -version = "0.5.1" +name = "annotate-snippets" +version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +checksum = "710e8eae58854cdc1790fcb56cca04d712a17be849eeb81da2a724bf4bae2bc4" dependencies = [ - "num-traits", + "anstyle", + "unicode-width", ] [[package]] -name = "atomic-traits" -version = "0.3.0" +name = "anstream" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b29ec3788e96fb4fdb275ccb9d62811f2fa903d76c5eb4dd6fe7d09a7ed5871f" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" dependencies = [ - "cfg-if", - "rustc_version", + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", ] [[package]] -name = "autocfg" -version = "1.3.0" +name = "anstyle" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" [[package]] -name = "base" -version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" dependencies = [ - "base_macros", - "detect", - "half 2.4.1", - "libc", - "log", - "rand", - "rayon", - "serde", - "thiserror", - "toml", - "validator", + "utf8parse", ] [[package]] -name = "base_macros" -version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.77", + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", ] +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "bindgen" -version = "0.70.1" +version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f" +checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ "annotate-snippets", "bitflags", @@ -107,14 +130,14 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.77", + "syn", ] [[package]] name = "bitflags" -version = "2.6.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "bitvec" @@ -129,61 +152,34 @@ dependencies = [ ] [[package]] -name = "bytecheck" -version = "0.6.12" +name = "bumpalo" +version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23cdc57ce23ac53c931e88a43d06d070a6fd142f2617be5855eb75efc9beb1c2" -dependencies = [ - "bytecheck_derive", - "ptr_meta", - "simdutf8", -] +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" [[package]] -name = "bytecheck_derive" -version = "0.6.12" +name = "cargo_toml" +version = "0.22.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3db406d29fbcd95542e92559bed4d8ad92636d1ca8b3b72ede10b4bcc010e659" +checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77" dependencies = [ - "proc-macro2", - "quote", - "syn 1.0.109", + "serde", + "toml 0.9.12+spec-1.1.0", ] [[package]] -name = "bytemuck" -version = "1.17.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773d90827bc3feecfb67fab12e24de0749aad83c74b9504ecde46237b5cd24e2" - -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - -[[package]] -name = "bytes" -version = "1.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50" - -[[package]] -name = "cargo_toml" -version = "0.19.2" +name = "cast" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a98356df42a2eb1bd8f1793ae4ee4de48e384dd974ce5eac8eee802edb7492be" -dependencies = [ - "serde", - "toml", -] +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.1.15" +version = "1.2.60" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57b6a275aa2903740dc87da01c62040406b8812552e97129a63ea8850a17c6e6" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" dependencies = [ + "find-msvc-tools", "shlex", ] @@ -208,9 +204,47 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.1", +] + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half 2.7.1", +] [[package]] name = "clang-sys" @@ -224,33 +258,136 @@ dependencies = [ ] [[package]] -name = "common" -version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" dependencies = [ - "base", - "log", - "memmap2", - "rand", - "rustix", - "serde", - "serde_json", + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "codepage" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48f68d061bc2828ae826206326e61251aca94c1e4a5305cf52d9138639c918b4" +dependencies = [ + "encoding_rs", ] +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "convert_case" -version = "0.6.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +checksum = "baaaa0ecca5b51987b9423ccdc971514dd8b0bb7b4060b983d3664dad3f1f89f" dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" -version = "0.8.5" +version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" dependencies = [ "crossbeam-epoch", "crossbeam-utils", @@ -258,30 +395,30 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] [[package]] name = "crossbeam-utils" -version = "0.8.20" +version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "darling" -version = "0.20.10" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f63b86c8a8826a49b8c21f08a2d07338eec8d900540f8630dc76284be802989" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ "darling_core", "darling_macro", @@ -289,52 +426,66 @@ dependencies = [ [[package]] name = "darling_core" -version = "0.20.10" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95133861a8032aaea082871032f5815eb9e98cef03fa916ab4500513994df9e5" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" dependencies = [ - "fnv", "ident_case", "proc-macro2", "quote", "strsim", - "syn 2.0.77", + "syn", ] [[package]] name = "darling_macro" -version = "0.20.10" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core", "quote", - "syn 2.0.77", + "syn", ] [[package]] -name = "detect" -version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ - "detect_macros", + "proc-macro2", + "quote", + "syn", ] [[package]] -name = "detect_macros" +name = "distance" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.77", + "zerocopy", ] [[package]] name = "either" -version = "1.13.0" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "encoding_rs" +version = "0.8.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] [[package]] name = "enum-map" @@ -353,23 +504,23 @@ checksum = "f282cfdfe92516eb26c2af8589c274c7c17681f5ecc03c18255fe741c6aa64eb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] [[package]] name = "equivalent" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" -version = "0.3.9" +version = "0.3.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys", ] [[package]] @@ -382,23 +533,71 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "feistel" +version = "0.0.0" +dependencies = [ + "rand", + "wyhash", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + [[package]] name = "fixedbitset" -version = "0.4.2" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "foldhash" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" [[package]] -name = "fnv" -version = "1.0.7" +name = "foldhash" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" [[package]] name = "form_urlencoded" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" dependencies = [ "percent-encoding", ] @@ -411,20 +610,23 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "getrandom" -version = "0.2.15" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", "libc", - "wasi", + "r-efi", + "rand_core 0.10.1", + "wasip2", + "wasip3", ] [[package]] name = "glob" -version = "0.3.1" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "half" @@ -434,107 +636,215 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.4.1" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", - "rand", - "rand_distr", - "serde", + "zerocopy", ] [[package]] -name = "hash32" -version = "0.3.1" +name = "hashbrown" +version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "byteorder", + "foldhash 0.1.5", ] [[package]] name = "hashbrown" -version = "0.12.3" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" dependencies = [ - "ahash", + "foldhash 0.2.0", ] [[package]] name = "hashbrown" -version = "0.14.5" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" [[package]] -name = "heapless" -version = "0.8.0" +name = "hashlink" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +checksum = "ea0b22561a9c04a7cb1a302c013e0259cd3b4bb619f145b32f72b8b4bcbed230" dependencies = [ - "hash32", - "stable_deref_trait", + "hashbrown 0.16.1", ] [[package]] -name = "hermit-abi" -version = "0.4.0" +name = "heck" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] -name = "home" -version = "0.5.9" +name = "humansize" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +checksum = "6cb51c9a029ddc91b07a787f1d86b53ccfa49b0e86688c946ebe8d3555685dd7" dependencies = [ - "windows-sys 0.52.0", + "libm", ] [[package]] -name = "ident_case" -version = "1.0.1" +name = "icu_collections" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] [[package]] -name = "idna" -version = "0.5.0" +name = "icu_locale_core" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" dependencies = [ - "unicode-bidi", - "unicode-normalization", + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", ] [[package]] -name = "indenter" -version = "0.3.3" +name = "icu_normalizer" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] [[package]] -name = "indexmap" -version = "2.5.0" +name = "icu_normalizer_data" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68b900aa2f7301e21c36462b170ee99994de34dff39a4a6a528e80e7376d07e5" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" dependencies = [ - "equivalent", - "hashbrown 0.14.5", + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", ] [[package]] -name = "is-terminal" -version = "0.4.13" +name = "icu_properties_data" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "261f68e344040fbd0edea105bef17c66edf46f984ddb1115b775ce31be948f4b" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" dependencies = [ - "hermit-abi", - "libc", - "windows-sys 0.52.0", + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indenter" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" + +[[package]] +name = "index" +version = "0.0.0" +dependencies = [ + "always_equal", + "bumpalo", + "dary_heap", + "small_iter", + "zerocopy", +] + +[[package]] +name = "index_accessor" +version = "0.0.0" +dependencies = [ + "distance", + "rabitq", + "simd", + "vector", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.0", + "serde", + "serde_core", ] [[package]] @@ -543,123 +853,150 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" -version = "0.12.1" +version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" dependencies = [ "either", ] [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" +dependencies = [ + "once_cell", + "wasm-bindgen", +] [[package]] name = "k_means" version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" dependencies = [ - "base", - "common", - "half 2.4.1", + "always_equal", + "distance", + "rabitq", "rand", - "smawk", + "rayon", + "simd", ] +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "libc" -version = "0.2.158" +version = "0.2.185" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" [[package]] name = "libloading" -version = "0.8.5" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" dependencies = [ "cfg-if", - "windows-targets", + "windows-link", ] [[package]] name = "libm" -version = "0.2.8" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "libmimalloc-sys" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc89deee4af0429081d2a518c0431ae068222a5a262a3bc6ff4d8535ec2e02fe" +dependencies = [ + "cc", +] + +[[package]] +name = "libsqlite3-sys" +version = "0.37.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" +checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] [[package]] name = "linux-raw-sys" -version = "0.4.14" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] -name = "log" -version = "0.4.22" +name = "litemap" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] -name = "matrixmultiply" -version = "0.3.9" +name = "log" +version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9380b911e3e96d10c1f415da0876389aaf1b56759054eeb0de7df940c456ba1a" -dependencies = [ - "autocfg", - "rawpointer", -] +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] name = "memchr" -version = "2.7.4" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] -name = "memmap2" -version = "0.9.4" +name = "mimalloc" +version = "0.1.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322" +checksum = "aca3c01a711f395b4257b81674c0e90e8dd1f1e62c4b7db45f684cc7a4fcb18a" dependencies = [ - "libc", + "libmimalloc-sys", ] [[package]] -name = "minimal-lexical" -version = "0.2.1" +name = "min-max-heap" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +checksum = "2687e6cf9c00f48e9284cf9fd15f2ef341d03cc7743abf9df4c5f07fdee50b18" [[package]] -name = "nalgebra" -version = "0.33.0" +name = "minimal-lexical" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c4b5f057b303842cf3262c27e465f4c303572e7f6b0648f60e16248ac3397f4" -dependencies = [ - "approx", - "matrixmultiply", - "nalgebra-macros", - "num-complex", - "num-rational", - "num-traits", - "simba", - "typenum", -] +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] -name = "nalgebra-macros" -version = "0.2.2" +name = "miniz_oxide" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.77", + "adler2", + "simd-adler32", ] [[package]] @@ -673,67 +1010,61 @@ dependencies = [ ] [[package]] -name = "num-bigint" -version = "0.4.6" +name = "num-traits" +version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ - "num-integer", - "num-traits", + "autocfg", ] [[package]] -name = "num-complex" -version = "0.4.6" +name = "object" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +checksum = "63944c133d03f44e75866bbd160b95af0ec3f6a13d936d69d31c81078cbc5baf" dependencies = [ - "num-traits", + "flate2", + "memchr", + "ruzstd", + "wasmparser 0.245.1", ] [[package]] -name = "num-integer" -version = "0.1.46" +name = "once_cell" +version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" -dependencies = [ - "num-traits", -] +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] -name = "num-rational" -version = "0.4.2" +name = "once_cell_polyfill" +version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] -name = "num-traits" -version = "0.2.19" +name = "oorandom" +version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" -dependencies = [ - "autocfg", - "libm", -] +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" [[package]] -name = "once_cell" -version = "1.19.0" +name = "owo-colors" +version = "4.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "d211803b9b6b570f68772237e415a029d5a50c65d382910b879fb19d3271f94d" +dependencies = [ + "supports-color", +] [[package]] -name = "owo-colors" -version = "4.0.0" +name = "page_size" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caff54706df99d2a78a5a4e3455ff45448d81ef1bb63c22cd14052ca0e993a3f" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" dependencies = [ - "supports-color", + "libc", + "winapi", ] [[package]] @@ -754,44 +1085,32 @@ dependencies = [ [[package]] name = "percent-encoding" -version = "2.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" - -[[package]] -name = "pest" -version = "2.7.11" +version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd53dff83f26735fdc1ca837098ccf133605d794cdae66acfc2bfac3ec809d95" -dependencies = [ - "memchr", - "thiserror", - "ucd-trie", -] +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" [[package]] name = "petgraph" -version = "0.6.5" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", + "hashbrown 0.15.5", "indexmap", + "serde", ] [[package]] name = "pgrx" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd42eed31a68e3e6d7cec284cdaacd8ffa9b2b3ebb3f1df9ec2fd3558834c6e7" dependencies = [ - "atomic-traits", "bitflags", "bitvec", "enum-map", - "heapless", "libc", - "once_cell", - "paste", "pgrx-macros", "pgrx-pg-sys", "pgrx-sql-entity-graph", @@ -805,8 +1124,9 @@ dependencies = [ [[package]] name = "pgrx-bindgen" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d30fa3688ab1cb6429ca46725eb32d5d9f186806c0f0aa6419f86e418b8a81ab" dependencies = [ "bindgen", "cc", @@ -815,43 +1135,59 @@ dependencies = [ "pgrx-pg-config", "proc-macro2", "quote", + "regex", "shlex", - "syn 2.0.77", + "syn", "walkdir", ] +[[package]] +name = "pgrx-catalog" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7b8320f441bd6aa7460a397d1a16dfb915979a1f665c85d9a7bc10d57e38cb8" +dependencies = [ + "paste", + "pgrx", +] + [[package]] name = "pgrx-macros" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ca7cc537daa4cc004f12ad1d31756b506358aa900251a82c3280e719da88fa4" dependencies = [ "pgrx-sql-entity-graph", "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] [[package]] name = "pgrx-pg-config" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cd8670ad408093a83615fc68071929d7677b40021a475649a626319c464f0a3" dependencies = [ "cargo_toml", + "codepage", + "encoding_rs", "eyre", - "home", "owo-colors", "pathsearch", "serde", "serde_json", "thiserror", - "toml", + "toml 0.9.12+spec-1.1.0", "url", + "winapi", ] [[package]] name = "pgrx-pg-sys" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9f81faae87d533b573156261b17521c0c7f66f5ca5dab2b9e6ddbf5a63dbd14" dependencies = [ "cee-scape", "libc", @@ -859,133 +1195,160 @@ dependencies = [ "pgrx-macros", "pgrx-sql-entity-graph", "serde", - "sptr", ] [[package]] name = "pgrx-sql-entity-graph" -version = "0.12.5" -source = "git+https://github.com/tensorchord/pgrx.git?branch=v0.12.5-patch#1fd3d1544c2f9ac68ec0c6e293fd2de398b6d3d7" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4800fa91a481a83d3c1aaffd1324fcd2e587df0ed342c71d6a66d0899027e505" dependencies = [ "convert_case", "eyre", "petgraph", "proc-macro2", "quote", - "syn 2.0.77", + "syn", "thiserror", "unescape", ] [[package]] -name = "ppv-lite86" -version = "0.2.20" +name = "pin-project" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" dependencies = [ - "zerocopy", + "pin-project-internal", ] [[package]] -name = "proc-macro-error" -version = "1.0.4" +name = "pin-project-internal" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" dependencies = [ - "proc-macro-error-attr", "proc-macro2", "quote", - "syn 1.0.109", - "version_check", + "syn", ] [[package]] -name = "proc-macro-error-attr" -version = "1.0.4" +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "plotters" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" dependencies = [ - "proc-macro2", - "quote", - "version_check", + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", ] [[package]] -name = "proc-macro2" -version = "1.0.86" +name = "plotters-backend" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" dependencies = [ - "unicode-ident", + "plotters-backend", +] + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", ] [[package]] -name = "ptr_meta" -version = "0.1.4" +name = "prettyplease" +version = "0.2.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0738ccf7ea06b608c10564b31debd4f5bc5e197fc8bfe088f68ae5ce81e7a4f1" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ - "ptr_meta_derive", + "proc-macro2", + "syn", ] [[package]] -name = "ptr_meta_derive" -version = "0.1.4" +name = "proc-macro-error-attr3" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16b845dbfca988fa33db069c0e230574d15a3088f147a87b64c7589eb662c9ac" +checksum = "34e4dd828515431dd6c4a030d26f7eaed7dd4778226e9d2bb968d65ca4ec3d4d" dependencies = [ "proc-macro2", "quote", - "syn 1.0.109", ] [[package]] -name = "quantization" -version = "0.0.0" -source = "git+ssh://git@github.com/tensorchord/pgvecto.rs.git?branch=rabbithole-2#29df8b43d45861fc741034bc7f8f304ca18822ca" +name = "proc-macro-error3" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ee475e440453418ff1335189eddf7101ba502cd818ab7ae04209bc83aa925aa" dependencies = [ - "base", - "common", - "detect", - "k_means", - "log", - "nalgebra", - "rand", - "rand_chacha", - "rand_distr", - "serde", - "serde_json", + "proc-macro-error-attr3", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", ] [[package]] name = "quote" -version = "1.0.37" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" dependencies = [ "proc-macro2", ] [[package]] -name = "rabbithole" +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rabitq" version = "0.0.0" dependencies = [ - "base", - "common", - "detect", - "k_means", - "nalgebra", - "paste", - "pgrx", - "quantization", "rand", "rand_chacha", - "rand_distr", - "rkyv", - "seq-macro", - "serde", - "toml", - "validator", + "simd", + "zerocopy", ] [[package]] @@ -996,55 +1359,42 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.8.5" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" dependencies = [ - "libc", - "rand_chacha", - "rand_core", + "chacha20", + "getrandom", + "rand_core 0.10.1", ] [[package]] name = "rand_chacha" -version = "0.3.1" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.10.1", ] [[package]] name = "rand_core" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] - -[[package]] -name = "rand_distr" -version = "0.4.3" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" -dependencies = [ - "num-traits", - "rand", -] +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" [[package]] -name = "rawpointer" -version = "0.2.1" +name = "rand_core" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] name = "rayon" -version = "1.10.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" dependencies = [ "either", "rayon-core", @@ -1052,9 +1402,9 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.12.1" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" dependencies = [ "crossbeam-deque", "crossbeam-utils", @@ -1062,9 +1412,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.10.6" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4219d74c6b67a3654a9fbebc4b419e22126d13d2f3c4a07ee0cb61ff79a79619" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -1074,9 +1424,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.7" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38caf58cc5ef2fed281f89292ef23f6365465ed9a41b7a7754eb4e26496c92df" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", @@ -1085,89 +1435,67 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.4" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" [[package]] -name = "rend" -version = "0.4.2" +name = "rsqlite-vfs" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71fe3824f5629716b1589be05dacd749f6aa084c87e00e016714a8cdfccc997c" +checksum = "a8a1f2315036ef6b1fbacd1972e8ee7688030b0a2121edfc2a6550febd41574d" dependencies = [ - "bytecheck", + "hashbrown 0.16.1", + "thiserror", ] [[package]] -name = "rkyv" -version = "0.7.45" +name = "rusqlite" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9008cd6385b9e161d8229e1f6549dd23c3d022f132a2ea37ac3a10ac4935779b" +checksum = "a0d2b0146dd9661bf67bb107c0bb2a55064d556eeb3fc314151b957f313bcd4e" dependencies = [ - "bitvec", - "bytecheck", - "bytes", - "hashbrown 0.12.3", - "ptr_meta", - "rend", - "rkyv_derive", - "seahash", - "tinyvec", - "uuid", + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", + "sqlite-wasm-rs", ] [[package]] -name = "rkyv_derive" -version = "0.7.45" +name = "rustc-hash" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503d1d27590a2b0a3a4ca4c94755aa2875657196ecbf401a42eff41d7de532c0" -dependencies = [ - "proc-macro2", - "quote", - "syn 1.0.109", -] - -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - -[[package]] -name = "rustc_version" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0dfe2087c51c460008730de8b57e6a320782fbfb312e1f4d520e6c6fae155ee" -dependencies = [ - "semver", -] +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" [[package]] name = "rustix" -version = "0.38.35" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a85d50532239da68e9addb745ba38ff4612a242c1c7ceea689c4bc7c2f43c36f" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys", ] [[package]] -name = "ryu" -version = "1.0.18" +name = "rustversion" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] -name = "safe_arch" -version = "0.7.2" +name = "ruzstd" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3460605018fdc9612bce72735cba0d27efbcd9904780d44c7e3a9948f96148a" +checksum = "e5ff0cc5e135c8870a775d3320910cd9b564ec036b4dc0b8741629020be63f01" dependencies = [ - "bytemuck", + "twox-hash", ] [[package]] @@ -1187,34 +1515,23 @@ checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" [[package]] name = "semver" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f301af10236f6df4160f7c3f04eec6dbc70ace82d23326abad5edee88801c6b6" -dependencies = [ - "semver-parser", -] - -[[package]] -name = "semver-parser" -version = "0.10.2" +version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00b0bef5b7f9e0df16536d3961cfb6e84331c065b4066afb39768d0e319411f7" -dependencies = [ - "pest", -] +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "seq-macro" -version = "0.3.5" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.209" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99fce0ffe7310761ca6bf9faf5115afbc19688edd00171d81b1bb1b116c63e09" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ + "serde_core", "serde_derive", ] @@ -1228,36 +1545,46 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + [[package]] name = "serde_derive" -version = "1.0.209" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5831b979fd7b5439637af1752d535ff49f4860c0f341d1baeb6faf0f4242170" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] [[package]] name = "serde_json" -version = "1.0.127" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8043c06d9f82bd7271361ed64f415fe5e12a77fdb52e573e7f06a516dea329ad" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ "itoa", "memchr", - "ryu", "serde", + "serde_core", + "zmij", ] [[package]] name = "serde_spanned" -version = "0.6.7" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb5b1b31579f3811bf615c144393417496f152e12ac8b7663bf664f4a815306d" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" dependencies = [ - "serde", + "serde_core", ] [[package]] @@ -1267,41 +1594,60 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] -name = "simba" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3a386a501cd104797982c15ae17aafe8b9261315b5d07e3ec803f2ea26be0fa" +name = "simd" +version = "0.0.0" dependencies = [ - "approx", - "num-complex", - "num-traits", - "paste", - "wide", + "cc", + "criterion", + "half 2.7.1", + "rand", + "seq-macro", + "simd_macros", + "zerocopy", ] [[package]] -name = "simdutf8" -version = "0.1.4" +name = "simd-adler32" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f27f6278552951f1f2b8cf9da965d10969b2efdea95a6ec47987ab46edfe263a" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" [[package]] -name = "smawk" -version = "0.3.2" +name = "simd_macros" +version = "0.0.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "small_iter" +version = "0.0.0" + +[[package]] +name = "smallvec" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] -name = "sptr" -version = "0.3.2" +name = "sqlite-wasm-rs" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b9b39299b249ad65f3b7e96443bad61c02ca5cd3589f46cb6d610a0fd6c0d6a" +checksum = "1b2c760607300407ddeaee518acf28c795661b7108c75421303dbefb237d3a36" +dependencies = [ + "cc", + "js-sys", + "rsqlite-vfs", + "wasm-bindgen", +] [[package]] name = "stable_deref_trait" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "strsim" @@ -1311,19 +1657,18 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "supports-color" -version = "2.1.0" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6398cde53adc3c4557306a96ce67b302968513830a77a95b2b17305d9719a89" +checksum = "c64fc7232dd8d2e4ac5ce4ef302b1d81e0b80d055b9d77c7c4f51f6aa4c867d6" dependencies = [ - "is-terminal", "is_ci", ] [[package]] name = "syn" -version = "1.0.109" +version = "2.0.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" dependencies = [ "proc-macro2", "quote", @@ -1331,14 +1676,14 @@ dependencies = [ ] [[package]] -name = "syn" -version = "2.0.77" +name = "synstructure" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f35bcdf61fd8e7be6caf75f429fdca8beb3ed76584befb503b1569faee373ed" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "unicode-ident", + "syn", ] [[package]] @@ -1347,151 +1692,204 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" +[[package]] +name = "target-triple" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "591ef38edfb78ca4771ee32cf494cb8771944bee237a9b91fc9c1424ac4b777b" + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + [[package]] name = "thiserror" -version = "1.0.63" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.63" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] [[package]] -name = "tinyvec" -version = "1.8.0" +name = "tinystr" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" dependencies = [ - "tinyvec_macros", + "displaydoc", + "zerovec", ] [[package]] -name = "tinyvec_macros" -version = "0.1.1" +name = "tinytemplate" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] [[package]] name = "toml" -version = "0.8.19" +version = "0.9.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1ed1f98e3fdc28d6d910e6737ae6ab1a93bf1985935a1193e68f93eeb68d24e" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" dependencies = [ - "serde", + "indexmap", + "serde_core", "serde_spanned", - "toml_datetime", - "toml_edit", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 0.7.15", ] [[package]] -name = "toml_datetime" -version = "0.6.8" +name = "toml" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" dependencies = [ - "serde", + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 1.0.1", ] [[package]] -name = "toml_edit" -version = "0.22.20" +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "583c44c02ad26b0c3f3066fe629275e50627026c51ac2e595cca4c230ce1ce1d" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" dependencies = [ - "indexmap", - "serde", - "serde_spanned", - "toml_datetime", - "winnow", + "serde_core", ] [[package]] -name = "typenum" -version = "1.17.0" +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] [[package]] -name = "ucd-trie" -version = "0.1.6" +name = "toml_parser" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow 1.0.1", +] [[package]] -name = "unescape" -version = "0.1.0" +name = "toml_writer" +version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" [[package]] -name = "unicode-bidi" -version = "0.3.15" +name = "twox-hash" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08f95100a766bf4f8f28f90d77e0a5461bbdb219042e7679bebe79004fed8d75" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" [[package]] -name = "unicode-ident" -version = "1.0.12" +name = "unescape" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] -name = "unicode-normalization" -version = "0.1.23" +name = "unicode-ident" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a56d1686db2308d901306f92a263857ef59ea39678a5458e7cb17f01415101f5" -dependencies = [ - "tinyvec", -] +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" -version = "1.11.0" +version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" [[package]] name = "unicode-width" -version = "0.1.14" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" [[package]] name = "url" -version = "2.5.2" +version = "2.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" dependencies = [ "form_urlencoded", "idna", "percent-encoding", + "serde", ] +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" -version = "1.10.0" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ "getrandom", + "js-sys", + "wasm-bindgen", ] [[package]] name = "validator" -version = "0.18.1" +version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db79c75af171630a3148bd3e6d7c4f42b6a9a014c2945bc5ed0020cbb8d9478e" +checksum = "43fb22e1a008ece370ce08a3e9e4447a910e92621bb49b85d6e48a45397e7cfa" dependencies = [ "idna", "once_cell", @@ -1505,23 +1903,103 @@ dependencies = [ [[package]] name = "validator_derive" -version = "0.18.2" +version = "0.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df0bcf92720c40105ac4b2dda2a4ea3aa717d4d6a862cc217da653a4bd5c6b10" +checksum = "240e4b81c20a1d6d50d1d7265c658dfbd204e8b9ac4d80f3c931f39462196335" dependencies = [ "darling", - "once_cell", - "proc-macro-error", + "proc-macro-error3", "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] [[package]] -name = "version_check" -version = "0.9.5" +name = "vchord" +version = "0.0.0" +dependencies = [ + "always_equal", + "arrayvec", + "bumpalo", + "crossbeam-channel", + "dary_heap", + "distance", + "feistel", + "humansize", + "index", + "index_accessor", + "k_means", + "libc", + "mimalloc", + "paste", + "pgrx", + "pgrx-catalog", + "rabitq", + "rand", + "rayon", + "rusqlite", + "seq-macro", + "serde", + "simd", + "toml 1.1.2+spec-1.1.0", + "validator", + "vchordg", + "vchordrq", + "vector", + "wyhash", + "zerocopy", +] + +[[package]] +name = "vchordg" +version = "0.0.0" +dependencies = [ + "always_equal", + "distance", + "index", + "index_accessor", + "min-max-heap", + "rabitq", + "rand", + "serde", + "simd", + "validator", + "vector", + "zerocopy", +] + +[[package]] +name = "vchordrq" +version = "0.0.0" +dependencies = [ + "always_equal", + "distance", + "index", + "index_accessor", + "pin-project", + "rabitq", + "rand", + "serde", + "simd", + "validator", + "vector", + "zerocopy", +] + +[[package]] +name = "vcpkg" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "vector" +version = "0.0.0" +dependencies = [ + "distance", + "rabitq", + "simd", +] [[package]] name = "walkdir" @@ -1534,19 +2012,119 @@ dependencies = [ ] [[package]] -name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +name = "wasip2" +version = "1.0.3+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen 0.57.1", +] [[package]] -name = "wide" -version = "0.7.28" +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b828f995bf1e9622031f8009f8481a85406ce1f4d4588ff746d872043e855690" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "bytemuck", - "safe_arch", + "wit-bindgen 0.51.0", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser 0.244.0", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser 0.244.0", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + +[[package]] +name = "wasmparser" +version = "0.245.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f08c9adee0428b7bddf3890fc27e015ac4b761cc608c822667102b8bfd6995e" +dependencies = [ + "bitflags", +] + +[[package]] +name = "web-sys" +version = "0.3.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d" +dependencies = [ + "js-sys", + "wasm-bindgen", ] [[package]] @@ -1567,11 +2145,11 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.9" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys", ] [[package]] @@ -1581,94 +2159,139 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] -name = "windows-sys" -version = "0.52.0" +name = "windows-link" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" -dependencies = [ - "windows-targets", -] +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" [[package]] name = "windows-sys" -version = "0.59.0" +version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ - "windows-targets", + "windows-link", ] [[package]] -name = "windows-targets" -version = "0.52.6" +name = "winnow" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" -dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", -] +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" [[package]] -name = "windows_aarch64_gnullvm" -version = "0.52.6" +name = "winnow" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +checksum = "09dac053f1cd375980747450bfc7250c264eaae0583872e845c0c7cd578872b5" [[package]] -name = "windows_aarch64_msvc" -version = "0.52.6" +name = "wit-bindgen" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] [[package]] -name = "windows_i686_gnu" -version = "0.52.6" +name = "wit-bindgen" +version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] -name = "windows_i686_gnullvm" -version = "0.52.6" +name = "wit-bindgen-core" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] [[package]] -name = "windows_i686_msvc" -version = "0.52.6" +name = "wit-bindgen-rust-macro" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] [[package]] -name = "windows_x86_64_gnu" -version = "0.52.6" +name = "wit-component" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser 0.244.0", + "wit-parser", +] [[package]] -name = "windows_x86_64_gnullvm" -version = "0.52.6" +name = "wit-parser" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser 0.244.0", +] [[package]] -name = "windows_x86_64_msvc" -version = "0.52.6" +name = "writeable" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] -name = "winnow" -version = "0.6.18" +name = "wyhash" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68a9bda4691f099d435ad181000724da8e5899daa10713c2d432552b9ccd3a6f" +checksum = "ca4d373340c479fd1e779f7a763acee85da3e423b1a9a9acccf97babcc92edbb" dependencies = [ - "memchr", + "rand_core 0.9.5", ] [[package]] @@ -1681,31 +2304,117 @@ dependencies = [ ] [[package]] -name = "yansi-term" -version = "0.1.2" +name = "xtask" +version = "0.0.0" +dependencies = [ + "clap", + "object", + "serde", + "serde_json", + "shlex", + "target-triple", + "tempfile", +] + +[[package]] +name = "yoke" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe5c30ade05e61656247b2e334a031dfd0cc466fadef865bdcdea8d537951bf1" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" dependencies = [ - "winapi", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", ] [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ - "byteorder", "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.35" +version = "0.8.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn 2.0.77", + "syn", ] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml index bef4e791..32654e1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,56 +1,124 @@ [package] -name = "rabbithole" -version = "0.0.0" -edition = "2021" +name = "vchord" +version.workspace = true +edition.workspace = true +publish = false +autotests = false [lib] -name = "rabbithole" +name = "vchord" crate-type = ["cdylib", "lib"] [[bin]] -name = "pgrx_embed_rabbithole" +name = "pgrx_embed_vchord" path = "./src/bin/pgrx_embed.rs" [features] -default = [] -pg14 = ["pgrx/pg14"] -pg15 = ["pgrx/pg15"] -pg16 = ["pgrx/pg16"] -pg17 = ["pgrx/pg17"] +default = ["simd/init"] +pg14 = ["pgrx/pg14", "pgrx-catalog/pg14"] +pg15 = ["pgrx/pg15", "pgrx-catalog/pg15"] +pg16 = ["pgrx/pg16", "pgrx-catalog/pg16"] +pg17 = ["pgrx/pg17", "pgrx-catalog/pg17"] +pg18 = ["pgrx/pg18", "pgrx-catalog/pg18"] [dependencies] -pgrx = { version = "=0.12.5", default-features = false, features = [] } -base = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -common = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -detect = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -k_means = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -quantization = { git = "ssh://git@github.com/tensorchord/pgvecto.rs.git", branch = "rabbithole-2" } -paste = "1" -serde = "1" -toml = "0.8.19" -rand = "0.8.5" -nalgebra = "0.33.0" -rand_distr = "0.4.3" -rkyv = { version = "0.7.45", features = ["validation"] } -rand_chacha = "0.3.1" -seq-macro = "0.3.5" -validator = "0.18.1" - -[patch.crates-io] -pgrx = { git = "https://github.com/tensorchord/pgrx.git", branch = "v0.12.5-patch" } +always_equal = { path = "./crates/always_equal" } +distance = { path = "./crates/distance" } +feistel = { path = "./crates/feistel" } +index = { path = "./crates/index" } +index_accessor = { path = "./crates/index_accessor" } +k_means = { path = "./crates/k_means" } +rabitq = { path = "./crates/rabitq" } +simd = { path = "./crates/simd" } +vchordg = { path = "./crates/vchordg" } +vchordrq = { path = "./crates/vchordrq" } +vector = { path = "./crates/vector" } + +arrayvec = "0.7.6" +bumpalo.workspace = true +crossbeam-channel = "0.5.15" +dary_heap.workspace = true +humansize = "2.1.3" +paste.workspace = true +pgrx = { version = "=0.17.0", default-features = false, features = ["cshim"] } +pgrx-catalog = "0.3.2" +rand.workspace = true +rayon.workspace = true +rusqlite = { version = "0.39.0", features = ["bundled"] } +seq-macro.workspace = true +serde.workspace = true +toml = "1.1.2" +validator.workspace = true +wyhash.workspace = true +zerocopy.workspace = true + +[target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] +mimalloc = { version = "0.1.49", features = ["local_dynamic_tls"] } + +[target.'cfg(unix)'.dependencies] +libc = "0.2.182" [lints] +workspace = true + +[workspace] +resolver = "3" +members = ["crates/*"] + +[workspace.package] +version = "0.0.0" +edition = "2024" + +[workspace.dependencies] +bumpalo = "3.20.2" +dary_heap = "0.3.9" +paste = "1.0.15" +rand = "0.10.1" +rand_chacha = "0.10.0" +rayon = "1.12.0" +seq-macro = "0.3.6" +serde = { version = "1.0.228", features = ["derive"] } +validator = { version = "0.20.0", features = ["derive"] } +wyhash = "0.6.0" +zerocopy = { version = "0.8.48", features = ["derive"] } + +[workspace.lints] +# complexity +clippy.identity_op = "allow" +clippy.int_plus_one = "allow" +clippy.manual_is_multiple_of = "allow" +clippy.nonminimal_bool = "allow" +clippy.too_many_arguments = "allow" +clippy.type_complexity = "allow" +# style +clippy.collapsible_if = "allow" +clippy.just_underscores_and_digits = "allow" +clippy.needless_range_loop = "allow" +# unsafe +rust.unsafe_code = "deny" rust.unsafe_op_in_unsafe_fn = "deny" +# unused +rust.unused_extern_crates = "warn" +rust.unused_import_braces = "warn" rust.unused_lifetimes = "warn" +rust.unused_macro_rules = "warn" rust.unused_qualifications = "warn" -[profile.opt] -inherits = "dev" -opt-level = 3 -debug-assertions = false -overflow-checks = false +[profile.dev] +codegen-units = 256 +lto = "off" +opt-level = 1 +debug = "full" +strip = "none" [profile.release] -lto = "fat" codegen-units = 1 -debug = true +lto = "fat" +opt-level = 3 +debug = "none" +strip = "symbols" + +[profile.prof] +inherits = "release" +debug = "full" +strip = "none" diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000..e6d43143 --- /dev/null +++ b/LICENSE @@ -0,0 +1,23 @@ +# Dual License Notice + +This software is licensed under a dual license model. You may choose to use this software under one of the following licenses: + +1. **GNU Affero General Public License v3 (AGPLv3)** + + This program is free software: you can redistribute it and/or modify it under the terms of the AGPLv3. + You should have received a copy of the AGPLv3 along with this program. If not, see . + + The AGPLv3 license allows you to use, modify, and distribute the software, but requires you to share your modifications under the same license. + +2. **Elastic License v2 (ELv2)** + + This software is also licensed under the Elastic License v2. + For details, see . + + The Elastic License v2 allows you to use, modify, and distribute the software, but imposes certain restrictions, particularly regarding the provision of the software as a service. + +## Choosing a License + +You may choose to use this software under either the AGPLv3 or the Elastic License v2. If you choose the AGPLv3, you must comply with its terms. If you choose the Elastic License v2, you must comply with its terms. + +By using this software, you agree to adhere to the terms of the license you select. diff --git a/META.json.in b/META.json.in new file mode 100644 index 00000000..19a2ffde --- /dev/null +++ b/META.json.in @@ -0,0 +1,48 @@ +{ + "name": "vchord", + "abstract": "Scalable, fast, and disk-friendly vector search in Postgres", + "description": "VectorChord (vchord) is a PostgreSQL extension designed for scalable, high-performance, and disk-efficient vector similarity search, it supports L2 distance, inner product, cosine distance and maxsim", + "version": "@DISTVERSION@", + "maintainer": [ + "Hu Xinjing HuXinjing@users.noreply.github.com" + ], + "license": "agpl_3", + "prereqs": { + "runtime": { + "requires": { + "PostgreSQL": "14.0.0" + } + } + }, + "provides": { + "vchord": { + "file": "src/lib.rs", + "docfile": "README.md", + "version": "@DISTVERSION@", + "abstract": "Scalable, fast, and disk-friendly vector search in Postgres" + } + }, + "resources": { + "homepage": "https://github.com/HuXinjing/VectorChord", + "bugtracker": { + "web": "https://github.com/HuXinjing/VectorChord/issues" + }, + "repository": { + "url": "https://github.com/HuXinjing/VectorChord.git", + "web": "https://github.com/HuXinjing/VectorChord", + "type": "git" + } + }, + "generated_by": "HuXinjing/VectorChord", + "meta-spec": { + "version": "1.0.0", + "url": "http://pgxn.org/meta/spec.txt" + }, + "tags": [ + "vector search", + "vectors", + "datatype", + "nearest neighbor search", + "approximate nearest neighbors" + ] +} diff --git a/Makefile b/Makefile new file mode 100644 index 00000000..d453aba9 --- /dev/null +++ b/Makefile @@ -0,0 +1,26 @@ +export PG_CONFIG ?= pg_config +PKGLIBDIR := $(shell $(PG_CONFIG) --pkglibdir) +SHAREDIR := $(shell $(PG_CONFIG) --sharedir) +CP_R ?= cp -r +MKDIR_P ?= mkdir -p + +.PHONY: all build clippy install uninstall + +all: build + +build: + cargo run -p xtask -- build + +clippy: + cargo run -p xtask -- clippy + +install: + $(MKDIR_P) $(DESTDIR)$(PKGLIBDIR) && \ + $(CP_R) ./build/pkglibdir/. $(DESTDIR)$(PKGLIBDIR) && \ + $(MKDIR_P) $(DESTDIR)$(SHAREDIR) && \ + $(CP_R) ./build/sharedir/. $(DESTDIR)$(SHAREDIR) + +uninstall: + $(RM) $(wildcard $(DESTDIR)$(PKGLIBDIR)/vchord.*) && \ + $(RM) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord.*) && \ + $(RM) $(wildcard $(DESTDIR)$(SHAREDIR)/extension/vchord--*.sql) diff --git a/README.md b/README.md new file mode 100644 index 00000000..da55600d --- /dev/null +++ b/README.md @@ -0,0 +1,392 @@ +
+ +# VectorChord TileMaxSim + +**PostgreSQL-native vector and tensor retrieval with exact TileMaxSim and a persistent GPU cache.** + +English | [简体中文](README.zh-CN.md) + +
+ +VectorChord TileMaxSim is an open-source retrieval engine built on PostgreSQL +and derived from the upstream +[VectorChord](https://github.com/supervc-stack/VectorChord) codebase. It preserves +the traditional single-vector path while adding exact late-interaction retrieval +for documents represented as arrays of token vectors. + +This repository is maintained as its own project. Its public README describes +this project's implementation, measured results, limitations, and roadmap; it +does not reproduce the upstream project's README or present upstream product +claims as our own. + +## Project scope + +VectorChord TileMaxSim provides infrastructure primitives: + +- vector, tensor-descriptor, and retrieval-metadata storage in PostgreSQL; +- traditional single-vector retrieval and exact multi-vector TileMaxSim; +- PostgreSQL permission, MVCC, row-visibility, cancellation, and timeout + semantics around retrieval; +- a persistent GPU tensor arena, host-memory cache, and immutable disk shards; +- bounded admission, fair/priority scheduling, cache quotas, health probes, and + metrics for the optional GPU service. + +It does not implement application identity, authentication, ACL policy, entity +registries, facts, events, relationship graphs, communities, or query-intent +routing. Those belong to an authenticated application or knowledge-governance +layer. VectorChord accepts only the already-authorized candidate scope and +opaque scheduling hints supplied by that caller. + +## Core capabilities + +- Exact TileMaxSim reranking on CPU and through the native Rust/CUDA + `tilemaxsimd` backend. +- No artificial cap on a caller-authorized tensor scope. Oversized working sets + are processed in bounded chunks. +- External tensor-source registration, so full token tensors can remain outside + the indexed PostgreSQL value while stable public IDs and descriptors remain + queryable. +- A three-level tensor path: + - L0: persistent GPU page cache; + - L1: bounded host-memory cache; + - L2: immutable tensor shards or content-addressed objects on durable storage. +- Batched shard reads, batched host-to-device transfers, content-addressed + deduplication, TinyLFU/GDSF admission, and a coalescing page-run allocator. +- Optional resident prewarming that completes before the daemon reports ready. +- Per-daemon fair, strict-priority, or fair-priority scheduling, request aging, + tenant weights, deadlines, disconnect cancellation, and bounded CUDA work + quanta. +- PostgreSQL planner statistics and cost estimation for multi-vector queries. +- Prometheus metrics and separate liveness/readiness probes. +- Correctness, registry, protocol, planner, real-GPU, and fault-path tests. + +TileMaxSim is opt-in. Ordinary single-vector operation does not start the CUDA +daemon and does not require a GPU-memory setting. When enabled, `tilemaxsimd` +reserves the configured devices and GiB allocations during startup and exits if +any requested device or allocation cannot be obtained. + +## Architecture + +```text +Authenticated application / retrieval planner + │ authorized IDs + scheduling hints + ▼ + PostgreSQL + VectorChord + │ tensor descriptors + ▼ + tilemaxsimd + ┌─────────┼──────────┐ + │ │ │ + L0 GPU L1 host L2 durable + cache cache shards + │ + └── exact CUDA TileMaxSim +``` + +PostgreSQL remains the source of row visibility and hard filtering. The daemon +does not infer authorization from tenant or priority fields. It receives an +explicit tensor set, resolves the corresponding immutable objects through the +cache hierarchy, and returns exact scores. + +## GPU cache and scheduling + +The default `fair-priority` scheduler combines explicit urgency with weighted +fairness. Large requests are split by candidate count, document tokens, and +actual `query rows × document rows × dimension` work. A request re-enters the +scheduler between CUDA launches, allowing another scheduling domain to run. +An executing CUDA kernel cannot be interrupted, so the FMA quantum bounds the +cooperative, non-preemptible interval. + +GPU and host caches have per-scheduling-domain maximums. An optional +`--tenant-cache-reservation TENANT=GB` protects already-warmed GPU pages from +cross-domain eviction below the configured allocation. This is currently an +aggregate byte reservation, not a permanent partition for one particular +knowledge base: it does not prewarm data, does not provide an L1 host-cache +minimum, and does not distinguish multiple sources owned by one tenant. These +gaps are listed explicitly in the roadmap. + +Memory flags use GiB rather than raw byte counts. A representative launch is: + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --listen 0.0.0.0:9191 \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ + --max-inflight-request-gb 1 \ + --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ + --scheduler-policy fair-priority \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 \ + --scheduler-quantum-fmas 4000000000 \ + --scheduler-quantum-io-gb 1 \ + --tenant-weight foreground=2 \ + --tenant-cache-reservation foreground=4 +``` + +PostgreSQL can connect through a local Unix socket or a `tcp://HOST:PORT` +endpoint. `GET /livez`, `GET /healthz`, and `GET /metrics` expose process +liveness, readiness, and bounded operational metrics. See +[`docs/TILEMAXSIM_CUDA_SIDECAR.md`](docs/TILEMAXSIM_CUDA_SIDECAR.md) and +[`docs/TILEMAXSIM_IPC_V2.md`](docs/TILEMAXSIM_IPC_V2.md) for deployment and +protocol details. + +## Measured performance + +The development corpus contains 34,054 descriptors, 34,027 unique tensors, and +16.28 GB of logical FP16 tensor data. Absolute latency depends on storage, CPU, +and GPU hardware; same-machine comparisons are more meaningful than isolated +numbers. + +| Optimization | Baseline | Optimized | Result | +| --- | ---: | ---: | ---: | +| Immutable shards and batched reads | 333.38 ms, sequential files | 56.52 ms | 5.90x faster | +| Shards versus batched legacy files | 87.56 ms | 56.52 ms | 1.55x faster | +| Batched host-to-device transfer | 37.06 ms, 100 transfers | 14.19 ms, one transfer | 2.61x faster | +| TinyLFU/GDSF admission | 69.98% LRU hit rate | 76.18% hit rate | +6.20 percentage points | +| Rust/CUDA cold request | 855.34 ms, Python/Triton | 93.86 ms | 9.11x faster | +| Rust/CUDA warm request p50 | 14.26 ms, Python/Triton | 2.09 ms | 6.82x faster | +| Rust/CUDA warm request p95 | 14.84 ms, Python/Triton | 2.20 ms | 6.75x faster | +| Small-VRAM cold I/O pipeline | 757.96 ms, serial | 685.85 ms, overlap | 1.11x faster; identical scores | + +The I/O-pipeline row used five cold trials over 1,000 real tensors +(478,016,000 logical bytes) with only 0.5 GiB assigned to the daemon: 0.4 GiB +for tensor pages and 0.1 GiB for CUDA workspace. The 0.05 GiB bounded resolver +stage overlaps SSD/L1 resolution of batch N+1, H2D upload of batch N, and exact +TileMaxSim for batch N-1. The maximum score delta between serial and overlap +was zero. Results are hardware- and batch-size-sensitive, so production must +measure the exposed pipeline metrics rather than assume this exact ratio. + +A 20 GiB run assigned 18 GiB to tensors and 2 GiB to workspace. Prewarming all +34,027 unique tensors took 14.86 seconds in the Rust/CUDA daemon versus 23.07 +seconds in the Python/Triton implementation. The default 32 KiB page-run +allocator reduced allocated space from 17.840 GB with the former 256 KiB buddy +allocator to 16.725 GB, reducing internal rounding waste from 8.82% to 2.74%. + +### Full-range and small-cache stress + +With all tensors resident, an intentionally exhaustive exact scan of all 34,054 +descriptors averaged 1.08 seconds. A 3 GiB tensor arena plus 1 GiB host cache +returned the same exact top-K, but sequential full scans caused 68,106 misses, +61,499 evictions, and 32.53 GB of host-to-device transfer; native round trips +averaged 18.47 seconds. + +The slowdown is cyclic cache thrashing and repeated I/O, not TileMaxSim +arithmetic becoming linearly slower as memory shrinks. Candidate-scoped queries +whose hot working set fits the cache do not exhibit this full-scan behavior. + +### Candidate-scope ablation + +Forty real queries compared the original 1,024-dimensional text embedding with +pgvector HNSW against exact TileMaxSim over a lexical candidate scope: + +| Path | Mean latency | p95 | Quality | +| --- | ---: | ---: | ---: | +| Original text embedding + pgvector HNSW | 5.80 ms | 5.86 ms | document hit@1 1.000 | +| Candidate-scoped TileMaxSim, 18 GiB resident cache | 45.56 ms | 134.99 ms | hit@1 0.475; hit@5 0.650 | +| Candidate-scoped TileMaxSim, 3 GiB cache | 573.43 ms | 1,812.70 ms | identical TileMaxSim ranking | + +The lexical scope covered the relevant document for only 65% of queries. +TileMaxSim kept the relevant document in its top five for every covered query, +so candidate generation—not exact scoring—set the recall ceiling. Lexical and +ordinary graph matches must therefore not be the only hard gate for general +semantic retrieval. Authorization filters remain hard; relational graph scopes +may be used for explicitly relational intent when the application can construct +a complete range. + +Reproducible drivers: + +- [`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py) +- [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) +- [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) +- [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) +- [`services/benchmark_tilemaxsim_io_pipeline.py`](services/benchmark_tilemaxsim_io_pipeline.py) + +## Limitations + +This project is suitable for controlled production pilots. Strict low-latency, +high-availability, or large multi-user deployments must address or explicitly +accept the following limitations before general availability. + +### Retrieval scale + +- Exact TileMaxSim is a scorer, not a tensor-native approximate-nearest- + neighbour index. Full residency removes transfer time but not exact scoring + arithmetic. +- Large low-latency corpora still need a measured high-recall candidate + generator or a future tensor-native ANN before exact TileMaxSim. +- A cache smaller than the active working set remains correct but can thrash; + production capacity must cover the hot set or preserve a safe high-recall + range. +- The current Tutti-inspired pipeline still uses CPU shard reads and pinned + host-to-device copies; it is not GPU io_uring or a GPUDirect Storage backend. + SSD, PCIe, and H2D contention make `--io-batch-gb` hardware-sensitive. Use + `--io-pipeline serial` as the ablation/fallback mode and inspect + `tilemaxsim_io_pipeline_*` before enabling overlap for an SLO. + +### Cache and multi-user isolation + +- Fairness, priority, admission, and reservations are enforced independently by + each daemon. Replicas do not share a global queue, entitlement ledger, or + residency-aware router. +- A GPU reservation protects aggregate bytes for a scheduling tenant, not a + named `source_id` or knowledge base. There is no dynamic cache-domain API, + source-level prewarm contract, or host-cache minimum. +- Reservations are configured at daemon startup. On multi-GPU daemons the + current value is applied to each device, rather than being expressed as one + explicit cluster-wide total. +- Globally resident manifests are pinned under a service-owned domain. If + pinned pages consume the entire arena, unrelated cold tensors cannot evict + them and the request fails instead of waiting for pinned data to leave. +- Preemption occurs only between CUDA launches. A running kernel is not + interruptible. +- Tenant and priority values are scheduling hints, not authorization evidence. + +### Operations and security + +- The TileMaxSim TCP protocol has no built-in TLS or client authentication. It + must remain on a private network or behind a mutually authenticated proxy. +- A replacement GPU replica starts cold. Durable storage preserves correctness, + but availability and warm-cache latency are separate SLOs. +- PostgreSQL WAL archiving, PITR, cross-site disaster recovery, and verified + restore procedures remain database-platform responsibilities. +- Metrics exist, but ready-made SLO dashboards and alert rules are not yet + included. +- Very large tensor registries and garbage collection require scale-specific + operational testing. + +### Release status + +The SQL surface, external-tensor protocol, image tags, and deployment packaging +remain under active development. Real-GPU and fault-path tests do not replace +multi-hour, many-user soak and chaos testing on target hardware. + +## Outlook: an AI infrastructure retrieval plane + +VectorChord TileMaxSim is intended to become the retrieval plane in a broader, +composable AI infrastructure rather than absorb application or model-serving +responsibilities: + +```text +Agent / application layer + │ + ▼ +AI gateway: identity, quota, priority, deadline, trace + │ + ├── Knowledge layer ── PostgreSQL / VectorChord ── tilemaxsimd + │ + └── Ray Serve ── vLLM and other model engines + +Kubernetes / KubeRay: placement, GPU nodes, scaling, recovery +Object storage: model weights, tensor shards, ingestion artifacts +Prometheus + OpenTelemetry: end-to-end observability +``` + +In this design: + +- [Ray Serve LLM](https://docs.ray.io/en/latest/serve/llm/) coordinates + distributed model replicas, autoscaling, and model-aware routing. +- [vLLM](https://docs.vllm.ai/) owns generation batching, model parallelism, + KV-cache management, and prefix-cache reuse. +- VectorChord owns PostgreSQL retrieval semantics and vector/tensor metadata; + `tilemaxsimd` owns exact TileMaxSim, persistent tensor residency, and bounded + retrieval scheduling. +- The authenticated application owns identity, ACL, graph/fact/event + governance, query intent, and the final RAG pipeline. + +The components should not silently compete for the same GPU memory. Initial +deployments should place vLLM and TileMaxSim in separate GPU pools. Kubernetes +or another resource manager must account for the memory reserved by the +long-lived TileMaxSim daemon before scheduling model workers; explicit MIG or +other hardware partitioning can be evaluated later. + +The next infrastructure milestones are: + +1. Separate `scheduler_tenant` from an opaque `cache_domain`, typically derived + by the authorized caller from tenant and source identity. +2. Add GPU and host `min_resident_gb`, `max_gb`, and prewarm contracts per cache + domain, with unambiguous per-device and fleet-wide semantics. +3. Build a bounded global admission layer and a residency-aware router across + TileMaxSim replicas, while leaving exact local quantum scheduling in each + daemon. +4. Propagate one request envelope—request ID, tenant, source, priority, + deadline, trace ID, and cache domain—through retrieval and model serving. +5. Validate cold-start recovery, rolling upgrades, PostgreSQL failover, GPU-pod + loss, multi-tenant isolation, and long-running soak/chaos behavior. +6. Add a tensor-native ANN or another recall-measured first stage for corpora + where exact full-range TileMaxSim cannot meet the latency SLO. + +Ray is an orchestration layer, not durable storage. PostgreSQL and immutable +object storage remain authoritative. `tilemaxsimd` should remain a supervised, +stateful GPU service rather than a short-lived generic Ray task, so its reserved +arena and cache lifetime remain predictable. + +## Security boundary + +PostgreSQL sends scheduled protocol metadata only when +`vchordrq.maxsim_tenant` is set. Priority affects latency ordering only; it +cannot bypass row visibility, ACL, source, or other mandatory filters. Request +logs expose a stable hash rather than raw tenant identifiers. + +This public repository contains implementation and public-facing project +information only. Private application and planning documents are intentionally +excluded. + +## Build and test + +Useful entry points include: + +```shell +make build +cargo test --locked --workspace --exclude vchord --no-fail-fast +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked +``` + +The CUDA container and its test stage are defined in +[`services/Dockerfile.tilemaxsimd`](services/Dockerfile.tilemaxsimd). PostgreSQL +tensor-source integration is documented in +[`docs/MAXSIM_TENSOR_SOURCES.md`](docs/MAXSIM_TENSOR_SOURCES.md). + +## Acknowledgements + +We thank the maintainers and contributors of the upstream +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord) +project. Its PostgreSQL extension, vector types, `vchordrq` index, and existing +MaxSim support provide the foundation on which this project builds. + +The project's IO-aware GPU MaxSim direction and the `TileMaxSim` name are +informed by Ashutosh Sharma's paper, +[“TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization”](https://arxiv.org/abs/2606.26439), +arXiv:2606.26439 (2026), and its accompanying open-source +[ashutoshuiuc/tilemaxsim](https://github.com/ashutoshuiuc/tilemaxsim) +Triton implementation. We gratefully acknowledge that work and its published +analysis of fused, tiled MaxSim scoring. + +The bounded asynchronous tensor pipeline is informed by +[“Tutti: Making SSD-Backed KV Cache Practical for Long-Context LLM Serving”](https://arxiv.org/abs/2605.03375), +particularly its GPU-native object, asynchronous I/O, and slack-aware +scheduling principles. This implementation currently adopts the scheduling +and pipelining ideas only; it does not include Tutti's GPU io_uring storage +stack. + +This repository is a separate PostgreSQL/Rust/CUDA systems integration. Its +IPC protocols, persistent three-level tensor cache, allocator, multi-user +scheduler, database bindings, and operational packaging are maintained here; +the acknowledgement does not imply endorsement or shared maintenance by the +upstream projects or paper author. + +## Project provenance and license + +This project is derived from +[supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord). +Original source files retain their upstream copyright and license notices. +TileMaxSim project additions are Copyright (c) 2026 Hu Xinjing. + +The repository is offered under the license terms in [`LICENSE`](LICENSE), +including AGPLv3 and Elastic License v2 options where applicable. Use this +project's [issue tracker](https://github.com/HuXinjing/VectorChord/issues) for +support and project-specific reports. diff --git a/README.zh-CN.md b/README.zh-CN.md new file mode 100644 index 00000000..cbf9bb20 --- /dev/null +++ b/README.zh-CN.md @@ -0,0 +1,332 @@ +
+ +# VectorChord TileMaxSim + +**基于 PostgreSQL 的向量/张量检索、精确 TileMaxSim 与常驻 GPU 缓存。** + +[English](README.md) | 简体中文 + +
+ +VectorChord TileMaxSim 是一个构建在 PostgreSQL 上的开源检索引擎，代码源自 +上游 [VectorChord](https://github.com/supervc-stack/VectorChord)。项目保留传统 +单向量检索通路，并增加面向“每个文档由一组 token 向量表示”的精确 +late-interaction 检索。 + +本仓库作为独立项目维护。README 只介绍本项目的实现、实测结果、局限和路线图， +不再附带上游官方 README，也不会把上游项目的宣传、镜像或产品指标表述成本项目 +自己的能力。 + +## 项目定位 + +VectorChord TileMaxSim 提供基础设施能力： + +- 在 PostgreSQL 中存储向量、张量描述符和检索元数据； +- 兼容传统单向量检索，并提供精确多向量 TileMaxSim； +- 让检索遵循 PostgreSQL 权限、MVCC、行可见性、取消和超时语义； +- 提供常驻 GPU 张量 arena、host 内存缓存和不可变磁盘分片； +- 为可选 GPU 服务提供有界准入、公平/优先级调度、缓存配额、健康探针和指标。 + +本项目不负责应用身份、认证、ACL 策略、实体注册表、Fact、事件账本、关系图谱、 +社区或查询意图路由。这些属于经过认证的应用层或知识治理层。VectorChord 只接收 +上层已经完成授权的候选范围，以及不具有授权含义的调度提示。 + +## 核心能力 + +- CPU 精确 TileMaxSim，以及原生 Rust/CUDA `tilemaxsimd` 后端。 +- 不人为限制调用方授权范围内的张量候选数量，超出显存的工作集按有界批次处理。 +- 外部张量源注册：PostgreSQL 保存稳定公开 ID 和描述符，完整 token 张量可存放在 + 被注册的行外存储中。 +- GPU/内存/磁盘三级张量路径： + - L0：常驻 GPU 页面缓存； + - L1：有界 host 内存缓存； + - L2：持久存储上的不可变张量分片或内容寻址对象。 +- 批量分片读取、批量 Host→Device 传输、内容寻址去重、TinyLFU/GDSF 准入，以及 + 可合并的 page-run 分配器。 +- 可选 resident 预热；预热完成前 daemon 不会报告 ready。 +- daemon 内的公平、严格 priority 或 fair-priority 调度，请求 aging、租户权重、 + deadline、断连取消和有界 CUDA quantum。 +- 多向量查询的 PostgreSQL 规划器统计与成本估算。 +- Prometheus 指标，以及相互独立的存活和就绪探针。 +- 正确性、注册表、协议、规划器、真实 GPU 和故障路径测试。 + +TileMaxSim 是可选功能。传统单向量通路不会启动 CUDA daemon，也不要求配置显存。 +启用时，`tilemaxsimd` 会在启动阶段一次性预占用户指定设备上的 GiB 级显存；任何 +设备或容量无法获得都会直接退出。 + +## 架构 + +```text +认证后的应用 / 检索规划器 + │ 已授权 ID + 调度提示 + ▼ + PostgreSQL + VectorChord + │ 张量描述符 + ▼ + tilemaxsimd + ┌─────────┼──────────┐ + │ │ │ +L0 GPU L1 内存 L2 持久存储 +缓存 缓存 分片 + │ + └── CUDA 精确 TileMaxSim +``` + +PostgreSQL 是行可见性和硬过滤的权威来源。daemon 不会根据 tenant 或 priority 推断 +权限，只会对调用方明确传入的张量集合进行缓存解析和精确打分。 + +## GPU 缓存与调度 + +默认 `fair-priority` 调度器组合显式优先级和加权公平性。大请求按候选数、文档 +token 数和实际 `query rows × document rows × dimension` 计算量切分。每个 CUDA +launch 结束后，请求重新进入调度器，因此其他调度域可以插入执行。已经开始的 +CUDA kernel 不能被中断，FMA quantum 用来限制这段协作式不可抢占时间。 + +GPU 和 host cache 都有调度域最大占用限制。可选的 +`--tenant-cache-reservation TENANT=GB` 会保护已经预热的 GPU 页面，使其他调度域 +不能把它淘汰到低于指定容量。不过它当前只是“聚合字节保底”，不是某个知识库的 +永久显存分区：它不会自动预热，不为 L1 内存提供最低保留量，也不能区分同一租户 +拥有的多个 source。这些缺口已明确列入展望。 + +所有内存参数使用 GiB，不使用裸字节数。示例： + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --listen 0.0.0.0:9191 \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ + --max-inflight-request-gb 1 \ + --contract-root MODEL_CONTRACT_ID=/srv/vectorchord/tensors \ + --scheduler-policy fair-priority \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 \ + --scheduler-quantum-fmas 4000000000 \ + --scheduler-quantum-io-gb 1 \ + --tenant-weight foreground=2 \ + --tenant-cache-reservation foreground=4 +``` + +PostgreSQL 可以通过本地 Unix socket 或 `tcp://HOST:PORT` 连接 daemon。 +`GET /livez`、`GET /healthz` 和 `GET /metrics` 分别提供存活、就绪和有界运行指标。 +部署和协议细节见 +[`docs/TILEMAXSIM_CUDA_SIDECAR.md`](docs/TILEMAXSIM_CUDA_SIDECAR.md) 与 +[`docs/TILEMAXSIM_IPC_V2.md`](docs/TILEMAXSIM_IPC_V2.md)。 + +## 性能实测 + +开发机语料包含 34,054 个张量描述符、34,027 个唯一张量，逻辑 FP16 数据量为 +16.28 GB。绝对延迟会随存储、CPU 和 GPU 改变，因此同机对照比孤立数字更有意义。 + +| 优化项 | 基线 | 优化后 | 结果 | +| --- | ---: | ---: | ---: | +| 不可变分片与批量读取 | 333.38 ms，顺序文件 | 56.52 ms | 快 5.90 倍 | +| 分片相对批量旧文件 | 87.56 ms | 56.52 ms | 快 1.55 倍 | +| 批量 Host→Device 传输 | 37.06 ms，100 次传输 | 14.19 ms，1 次传输 | 快 2.61 倍 | +| TinyLFU/GDSF 准入 | LRU 命中率 69.98% | 命中率 76.18% | +6.20 个百分点 | +| Rust/CUDA 冷请求 | Python/Triton 855.34 ms | 93.86 ms | 快 9.11 倍 | +| Rust/CUDA 热请求 p50 | Python/Triton 14.26 ms | 2.09 ms | 快 6.82 倍 | +| Rust/CUDA 热请求 p95 | Python/Triton 14.84 ms | 2.20 ms | 快 6.75 倍 | +| 小显存冷 I/O 流水线 | 串行 757.96 ms | 重叠 685.85 ms | 快 1.11 倍；分数一致 | + +I/O 流水线一行使用 1,000 个真实张量（逻辑 478,016,000 字节）做了 5 次冷请求， +daemon 只分配 0.5 GiB：0.4 GiB 张量页面加 0.1 GiB CUDA 工作区。0.05 GiB +有界解析批次让 N+1 批 SSD/L1 解析、N 批 H2D 上传与 N-1 批精确 TileMaxSim +重叠执行。串行与重叠路径的最大分数差为 0。实际收益对硬件和批次大小敏感，生产 +环境应以流水线指标实测，不能直接套用该倍数。 + +20 GiB 实验将 18 GiB 分给张量、2 GiB 分给工作区。Rust/CUDA daemon 预热全部 +34,027 个唯一张量耗时 14.86 秒，Python/Triton 实现为 23.07 秒。默认 32 KiB +page-run 分配器将原 256 KiB buddy allocator 的 17.840 GB 分配量降至 +16.725 GB，内部取整浪费从 8.82% 降至 2.74%。 + +### 全范围与小缓存压力 + +全部张量常驻时，刻意对 34,054 个描述符做全量精确扫描平均为 1.08 秒。3 GiB +张量 arena 加 1 GiB host cache 返回相同的精确 top-K，但连续全量扫描产生 +68,106 次 miss、61,499 次 eviction 和 32.53 GB Host→Device 传输；原生 +round trip 平均达到 18.47 秒。 + +该差距来自循环缓存抖动和重复 I/O，不是显存缩小几倍后 TileMaxSim 算术就线性 +变慢。高召回范围内的热工作集能放进缓存时，不会出现这种全库扫描行为。 + +### 候选范围消融 + +40 个真实查询使用原始 1,024 维文本 embedding/pgvector HNSW，与词法范围内的 +精确 TileMaxSim 对照： + +| 路径 | 平均延迟 | p95 | 质量 | +| --- | ---: | ---: | ---: | +| 原始文本 embedding + pgvector HNSW | 5.80 ms | 5.86 ms | 文档 hit@1 1.000 | +| 18 GiB 常驻缓存、候选范围内 TileMaxSim | 45.56 ms | 134.99 ms | hit@1 0.475；hit@5 0.650 | +| 3 GiB 缓存、候选范围内 TileMaxSim | 573.43 ms | 1,812.70 ms | TileMaxSim 排序相同 | + +词法范围只覆盖了 65% 查询的相关文档；对于已经覆盖的查询，TileMaxSim 全部把相关 +文档保留在 top 5。因此召回上限由候选生成而不是精确打分决定。普通词法和图谱命中 +不能作为一般语义查询的唯一硬门；授权过滤仍然必须硬执行，明确的关系类意图可以在 +上层能构造完整关系范围时使用图谱范围。 + +复现实验入口： + +- [`services/benchmark_tilemaxsim_ablation.py`](services/benchmark_tilemaxsim_ablation.py) +- [`services/benchmark_full_corpus_tilemaxsim.py`](services/benchmark_full_corpus_tilemaxsim.py) +- [`services/benchmark_gbrain_scoped_tilemaxsim.py`](services/benchmark_gbrain_scoped_tilemaxsim.py) +- [`services/benchmark_postgres_single_vector.py`](services/benchmark_postgres_single_vector.py) +- [`services/benchmark_tilemaxsim_io_pipeline.py`](services/benchmark_tilemaxsim_io_pipeline.py) + +## 已知局限 + +本项目适合受控生产试点。严格低延迟、高可用或大规模多用户场景在正式 GA 前，必须 +解决或明确接受以下限制。 + +### 检索规模 + +- 精确 TileMaxSim 是打分器，不是张量原生 ANN。全部常驻只能消除传输，不能消除 + 精确计算量。 +- 大规模低延迟知识库仍需要经过实测的高召回候选生成器，或未来的张量原生 ANN， + 再执行精确 TileMaxSim。 +- 缓存小于活跃工作集时结果仍然正确，但可能发生抖动；生产容量必须覆盖热工作集， + 或保留经过验证的高召回范围。 +- 当前受 Tutti 启发的流水线仍通过 CPU 读取分片并经 pinned host memory 上传， + 还不是 GPU io_uring 或 GPUDirect Storage 后端。SSD、PCIe 与 H2D 竞争使 + `--io-batch-gb` 对硬件敏感；应使用 `--io-pipeline serial` 做消融/回退，并在 + SLO 上线前检查 `tilemaxsim_io_pipeline_*` 指标。 + +### 缓存与多用户隔离 + +- 公平、priority、准入和 reservation 由每个 daemon 独立执行。多个副本之间没有 + 全局队列、全局 entitlement 账本或感知缓存驻留的路由器。 +- GPU reservation 保护的是某个调度 tenant 的聚合字节数，不是命名的 + `source_id` 或知识库；当前没有动态 cache-domain API、source 级预热契约或 + host-cache 最低保留量。 +- reservation 在 daemon 启动时静态配置。多 GPU daemon 当前会把该数值应用到 + 每张卡，而不是一个语义明确的集群总量。 +- 全局 resident manifest 使用服务内部域固定。如果 pinned 页面占满 arena， + 无关的冷张量不能淘汰它们，请求会失败，而不是等待 pinned 数据离开。 +- 抢占只能发生在 CUDA launch 之间，正在执行的 kernel 不可中断。 +- tenant 和 priority 只是调度提示，不是授权证据。 + +### 运维与安全 + +- TileMaxSim TCP 协议没有内置 TLS 和客户端认证，必须放在私网或双向认证代理后。 +- GPU 副本替换后是冷缓存。持久存储保证正确性，但可用性和热缓存延迟是两个独立 + SLO。 +- PostgreSQL WAL 归档、PITR、跨站灾备和真实恢复验证仍属于数据库平台职责。 +- 已提供指标，但还没有开箱即用的 SLO dashboard 和告警规则。 +- 超大张量注册表与垃圾回收需要按实际规模做运维压测。 + +### 发布状态 + +SQL 接口、外部张量协议、镜像 tag 和部署封装仍在持续演进。真实 GPU 和故障路径 +测试不能替代目标硬件上的长时间、多用户 soak/chaos 测试。 + +## 展望：AI Infra 的检索平面 + +VectorChord TileMaxSim 的目标是在可组合 AI Infra 中承担检索平面，而不是把应用 +治理和模型服务职责都塞进自身： + +```text +Agent / 应用层 + │ + ▼ +AI Gateway：身份、配额、priority、deadline、Trace + │ + ├── 知识层 ── PostgreSQL / VectorChord ── tilemaxsimd + │ + └── Ray Serve ── vLLM 与其他模型引擎 + +Kubernetes / KubeRay：节点放置、GPU、扩缩容、故障恢复 +对象存储：模型权重、张量分片、注入产物 +Prometheus + OpenTelemetry：端到端可观测性 +``` + +在这套分工中： + +- [Ray Serve LLM](https://docs.ray.io/en/latest/serve/llm/) 负责分布式模型副本、 + 自动扩缩容和模型感知路由； +- [vLLM](https://docs.vllm.ai/) 负责生成模型的连续批处理、模型并行、KV cache 和 + prefix cache； +- VectorChord 负责 PostgreSQL 检索语义以及向量/张量元数据，`tilemaxsimd` 负责 + 精确 TileMaxSim、张量驻留和有界检索调度； +- 经过认证的应用层负责身份、ACL、图谱/Fact/事件治理、查询意图和最终 RAG 链路。 + +各组件不能在未感知对方的情况下竞争同一块显存。初期部署应让 vLLM 与 TileMaxSim +使用独立 GPU 池；Kubernetes 或其他资源管理器必须知道常驻 TileMaxSim daemon +已经预占的显存，之后再评估 MIG 或其他硬件分区方案。 + +下一阶段基础设施里程碑： + +1. 将 `scheduler_tenant` 与不透明的 `cache_domain` 分离；后者通常由已认证上层根据 + tenant 和 source 派生，但 VectorChord 不解释其中的业务含义。 +2. 为每个 cache domain 增加 GPU 和 host 的 `min_resident_gb`、`max_gb` 与预热 + 契约，并明确区分单卡和全局语义。 +3. 建立有界的全局准入层和感知张量驻留的多副本路由，同时保留 daemon 内精确的 + quantum 调度。 +4. 让 request ID、tenant、source、priority、deadline、trace ID 和 cache domain + 组成统一请求信封，贯穿检索与模型服务。 +5. 验证冷启动恢复、滚动升级、PostgreSQL 切换、GPU Pod 丢失、多租户隔离和长时间 + soak/chaos 行为。 +6. 为无法用全范围精确 TileMaxSim 满足延迟 SLO 的大知识库增加张量原生 ANN，或 + 其他经过召回率验证的第一阶段。 + +Ray 是编排层，不是持久存储。PostgreSQL 和不可变对象存储仍是权威数据源。 +`tilemaxsimd` 应继续作为受监督的有状态 GPU 服务，而不是短生命周期的普通 Ray +任务，这样才能保证预占 arena 和缓存生命周期可预测。 + +## 安全边界 + +只有设置 `vchordrq.maxsim_tenant` 时，PostgreSQL 才发送带调度信息的协议。 +priority 只改变延迟顺序，不能绕过行可见性、ACL、source 或其他硬过滤。日志只 +暴露稳定哈希，不输出原始 tenant 标识。 + +本公开仓库只包含实现和对外项目信息，私有应用文档与规划文档不会提交到本仓库。 + +## 构建与测试 + +常用入口： + +```shell +make build +cargo test --locked --workspace --exclude vchord --no-fail-fast +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked +``` + +CUDA 容器及测试阶段定义在 +[`services/Dockerfile.tilemaxsimd`](services/Dockerfile.tilemaxsimd)，PostgreSQL +张量源集成见 [`docs/MAXSIM_TENSOR_SOURCES.md`](docs/MAXSIM_TENSOR_SOURCES.md)。 + +## 致谢 + +感谢上游 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord) +项目的维护者与贡献者。其 PostgreSQL 扩展、向量类型、`vchordrq` 索引和已有 +MaxSim 能力构成了本项目继续开发的基础。 + +本项目的 I/O 感知 GPU MaxSim 方向和 `TileMaxSim` 名称受到 Ashutosh Sharma +论文 +[《TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization》](https://arxiv.org/abs/2606.26439) +（arXiv:2606.26439，2026）及其配套开源 +[ashutoshuiuc/tilemaxsim](https://github.com/ashutoshuiuc/tilemaxsim) Triton 实现 +的启发。感谢该工作公开的融合、分块 MaxSim 设计与性能分析。 + +有界异步张量流水线也受到 +[《Tutti: Making SSD-Backed KV Cache Practical for Long-Context LLM Serving》](https://arxiv.org/abs/2605.03375) +中 GPU 原生对象、异步 I/O 与 slack-aware 调度原则的启发。当前实现只引入其调度 +与流水思想，不包含 Tutti 的 GPU io_uring 存储栈。 + +本仓库是独立的 PostgreSQL/Rust/CUDA 系统集成，IPC 协议、GPU/内存/磁盘三级 +张量缓存、分配器、多用户调度、数据库绑定和部署运维由本项目维护。上述致谢不表示 +上游项目或论文作者对本项目背书，也不表示双方共同维护。 + +## 项目来源与许可 + +本项目源自 [supervc-stack/VectorChord](https://github.com/supervc-stack/VectorChord)， +原始源文件保留其上游版权和许可声明。TileMaxSim 项目新增部分 Copyright (c) 2026 +Hu Xinjing。 + +仓库按 [`LICENSE`](LICENSE) 中的许可条款提供，在适用范围内包括 AGPLv3 和 +Elastic License v2 选项。项目问题与支持请求请提交到本项目的 +[Issue Tracker](https://github.com/HuXinjing/VectorChord/issues)。 diff --git a/build.rs b/build.rs new file mode 100644 index 00000000..642a488a --- /dev/null +++ b/build.rs @@ -0,0 +1,64 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::collections::HashMap; +use std::env::{var, var_os}; +use std::error::Error; +use std::process::{Command, Stdio}; + +fn main() -> Result<(), Box> { + println!("cargo::rerun-if-changed=vchord.control"); + let version = 'version: { + for line in std::fs::read_to_string("./vchord.control")?.lines() { + if let Some(prefix_stripped) = line.strip_prefix("default_version = '") + && let Some(stripped) = prefix_stripped.strip_suffix("'") + { + eprintln!("VectorChord version: {stripped}"); + break 'version stripped.to_string(); + } + } + return Err("VectorChord version is not defined.".into()); + }; + println!("cargo::rustc-env=VCHORD_VERSION={version}"); + if var("CARGO_CFG_TARGET_OS")? == "linux" { + println!("cargo::rustc-link-arg-cdylib=-Wl,-Bsymbolic"); + } + if var("CARGO_CFG_TARGET_OS")? == "macos" { + if let Some(path) = var_os("PGRX_PG_CONFIG_PATH") { + let map = { + let mut command = Command::new(&path); + command.stderr(Stdio::inherit()); + let command_output = command.output()?; + let command_stdout = String::from_utf8(command_output.stdout)?; + let mut map = HashMap::new(); + for line in command_stdout.lines() { + if let Some((key, value)) = line.split_once(" = ") { + map.insert(key.to_string(), value.to_string()); + eprintln!("Config `{key}`: {value}"); + } + } + map + }; + let bindir = &map["BINDIR"]; + println!("cargo::rustc-link-arg-cdylib=-Wl,-bundle,-bundle_loader,{bindir}/postgres",); + } else { + println!("cargo::rustc-link-arg-cdylib=-Wl,-undefined,dynamic_lookup"); + } + } + if var("CARGO_CFG_TARGET_OS")? == "emscripten" { + println!("cargo::rustc-link-arg-cdylib=-sSIDE_MODULE=2"); + println!("cargo::rustc-link-arg-bins=-sEXPORTED_FUNCTIONS=[_main]"); + } + Ok(()) +} diff --git a/crates/always_equal/Cargo.toml b/crates/always_equal/Cargo.toml new file mode 100644 index 00000000..887633ab --- /dev/null +++ b/crates/always_equal/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "always_equal" +version.workspace = true +edition.workspace = true +publish = false + +[lints] +workspace = true diff --git a/crates/always_equal/src/lib.rs b/crates/always_equal/src/lib.rs new file mode 100644 index 00000000..f95d470a --- /dev/null +++ b/crates/always_equal/src/lib.rs @@ -0,0 +1,49 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::cmp::Ordering; +use std::hash::Hash; + +#[derive(Debug, Clone, Copy, Default)] +#[repr(transparent)] +pub struct AlwaysEqual(pub T); + +impl PartialEq for AlwaysEqual { + #[inline(always)] + fn eq(&self, _: &Self) -> bool { + true + } +} + +impl Eq for AlwaysEqual {} + +#[allow(clippy::non_canonical_partial_ord_impl)] +impl PartialOrd for AlwaysEqual { + #[inline(always)] + fn partial_cmp(&self, _: &Self) -> Option { + Some(Ordering::Equal) + } +} + +impl Ord for AlwaysEqual { + #[inline(always)] + fn cmp(&self, _: &Self) -> Ordering { + Ordering::Equal + } +} + +impl Hash for AlwaysEqual { + #[inline(always)] + fn hash(&self, _: &mut H) {} +} diff --git a/crates/distance/Cargo.toml b/crates/distance/Cargo.toml new file mode 100644 index 00000000..348805cf --- /dev/null +++ b/crates/distance/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "distance" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +zerocopy.workspace = true + +[lints] +workspace = true diff --git a/crates/distance/src/lib.rs b/crates/distance/src/lib.rs new file mode 100644 index 00000000..cdaca6d8 --- /dev/null +++ b/crates/distance/src/lib.rs @@ -0,0 +1,115 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[derive( + Clone, + Copy, + Default, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +#[repr(transparent)] +pub struct Distance(i32); + +impl std::fmt::Debug for Distance { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.to_f32()) + } +} + +impl Distance { + pub const ZERO: Self = Distance::from_f32(0.0f32); + pub const INFINITY: Self = Distance::from_f32(f32::INFINITY); + pub const NEG_INFINITY: Self = Distance::from_f32(f32::NEG_INFINITY); + pub const NAN: Self = Distance::from_f32(f32::NAN); + + #[inline(always)] + pub const fn from_f32(value: f32) -> Self { + let bits = value.to_bits() as i32; + let mask = ((bits >> 31) as u32) >> 1; + let res = bits ^ (mask as i32); + Self(res) + } + + #[inline(always)] + pub const fn to_f32(self) -> f32 { + let bits = self.0; + let mask = ((bits >> 31) as u32) >> 1; + let res = bits ^ (mask as i32); + f32::from_bits(res as u32) + } + + #[inline(always)] + pub const fn to_i32(self) -> i32 { + self.0 + } +} + +impl From for Distance { + #[inline(always)] + fn from(value: f32) -> Self { + Distance::from_f32(value) + } +} + +impl From for f32 { + #[inline(always)] + fn from(value: Distance) -> Self { + Distance::to_f32(value) + } +} + +#[test] +fn distance_conversions() { + assert_eq!(Distance::from(0.0f32), Distance::ZERO); + assert_eq!(Distance::from(f32::INFINITY), Distance::INFINITY); + assert_eq!(Distance::from(f32::NEG_INFINITY), Distance::NEG_INFINITY); + for i in -100..100 { + let val = (i as f32) * 0.1; + assert_eq!(f32::from(Distance::from(val)).to_bits(), val.to_bits()); + } + assert_eq!( + f32::from(Distance::from(0.0f32)).to_bits(), + 0.0f32.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-0.0f32)).to_bits(), + (-0.0f32).to_bits() + ); + assert_eq!( + f32::from(Distance::from(f32::NAN)).to_bits(), + f32::NAN.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-f32::NAN)).to_bits(), + (-f32::NAN).to_bits() + ); + assert_eq!( + f32::from(Distance::from(f32::INFINITY)).to_bits(), + f32::INFINITY.to_bits() + ); + assert_eq!( + f32::from(Distance::from(-f32::INFINITY)).to_bits(), + (-f32::INFINITY).to_bits() + ); +} diff --git a/crates/feistel/Cargo.toml b/crates/feistel/Cargo.toml new file mode 100644 index 00000000..826e327f --- /dev/null +++ b/crates/feistel/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "feistel" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] + +[dev-dependencies] +rand.workspace = true +wyhash.workspace = true + +[lints] +workspace = true diff --git a/crates/feistel/src/lib.rs b/crates/feistel/src/lib.rs new file mode 100644 index 00000000..b7987f54 --- /dev/null +++ b/crates/feistel/src/lib.rs @@ -0,0 +1,118 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub fn feistel(width: u32, x: I, round: u32, secret: impl Fn(u32, I) -> I) -> I +where + I: Copy, + I: Ord, + I: std::ops::Add, + I: std::ops::Sub, + I: std::ops::BitAnd, + I: std::ops::BitOr, + I: std::ops::BitXor, + I: std::ops::Shl, + I: std::ops::Shr, + I: Zero, + I: One, + I: std::fmt::Debug, +{ + assert_eq!(width % 2, 0); + assert_eq!(x >> width, I::zero()); + + let half_width = width >> 1; + let half_mask = (I::one() << half_width) - I::one(); + + let mut left = (x >> half_width) & half_mask; + let mut right = x & half_mask; + + for i in 0..round { + (left, right) = (right, left ^ (secret(i, right) & half_mask)); + } + + (left << half_width) | right +} + +pub trait Zero { + fn zero() -> Self; +} + +pub trait One { + fn one() -> Self; +} + +macro_rules! impl_traits { + ($t:ty) => { + impl Zero for $t { + fn zero() -> Self { + 0 + } + } + + impl One for $t { + fn one() -> Self { + 1 + } + } + }; +} + +impl_traits!(u8); +impl_traits!(u16); +impl_traits!(u32); +impl_traits!(u64); +impl_traits!(u128); + +// This is a standalone crate, simply because we want to run tests for it. + +#[test] +fn is_a_permutation() { + let key_0 = [7u8; _]; + let key_1 = [8u8; _]; + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + for width in (0..if cfg!(not(miri)) { 20 } else { 10 }).step_by(2) { + let mut y = Vec::new(); + for x in 0..1 << width { + y.push(feistel::(width, x, 8, secret)); + } + if width <= 8 { + eprintln!("feistel({width}, _, 8, *) = {y:?}"); + } + y.sort_unstable(); + for x in 0..1 << width { + assert_eq!(y[x as usize], x); + } + } +} + +#[test] +fn sample() { + let n = if cfg!(not(miri)) { 6370_u32 } else { 637_u32 }; + let width = (n.ilog2() + 1).next_multiple_of(2); + let key_0 = rand::RngExt::random(&mut rand::rng()); + let key_1 = rand::RngExt::random(&mut rand::rng()); + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + let mut permutation = (0..1 << width) + .map(move |i| feistel(width, i, 8, secret)) + .filter(move |&x| x < n); + for _ in 0..n { + assert!(permutation.next().is_some()); + } + assert_eq!(permutation.next(), None); +} diff --git a/crates/index/Cargo.toml b/crates/index/Cargo.toml new file mode 100644 index 00000000..6d21f44c --- /dev/null +++ b/crates/index/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "index" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +always_equal = { path = "../always_equal" } +small_iter = { path = "../small_iter" } + +bumpalo.workspace = true +dary_heap.workspace = true +zerocopy.workspace = true + +[lints] +workspace = true diff --git a/crates/index/src/accessor.rs b/crates/index/src/accessor.rs new file mode 100644 index 00000000..e10a41d2 --- /dev/null +++ b/crates/index/src/accessor.rs @@ -0,0 +1,513 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use distance::Distance; +use rabitq::byte::CodeMetadata; +use simd::{Floating, f16}; +use std::marker::PhantomData; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; + +#[derive(Debug, Clone, Copy)] +pub struct L2S; + +#[derive(Debug, Clone, Copy)] +pub struct Dot; + +pub trait Accessor2 { + type Output; + fn push(&mut self, input: &[E0], target: &[E1]); + fn finish(self, input: M0, target: M1) -> Self::Output; +} + +impl Accessor2 for () { + type Output = (); + + #[inline(always)] + fn push(&mut self, _: &[E0], _: &[E1]) {} + + #[inline(always)] + fn finish(self, _: M0, _: M1) -> Self::Output {} +} + +impl> Accessor2 for (A,) { + type Output = (A::Output,); + + #[inline(always)] + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + } + + #[inline(always)] + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target),) + } +} + +impl, B: Accessor2> + Accessor2 for (A, B) +{ + type Output = (A::Output, B::Output); + + #[inline(always)] + fn push(&mut self, input: &[E0], target: &[E1]) { + self.0.push(input, target); + self.1.push(input, target); + } + + #[inline(always)] + fn finish(self, input: M0, target: M1) -> Self::Output { + (self.0.finish(input, target), self.1.finish(input, target)) + } +} + +pub trait Accessor1 { + type Output; + fn push(&mut self, input: &[E]); + fn finish(self, input: M) -> Self::Output; +} + +impl Accessor1 for () { + type Output = (); + + #[inline(always)] + fn push(&mut self, _: &[E]) {} + + #[inline(always)] + fn finish(self, _: M) -> Self::Output {} +} + +impl Accessor1 for (A,) +where + A: Accessor1, +{ + type Output = (A::Output,); + + #[inline(always)] + fn push(&mut self, input: &[E]) { + self.0.push(input); + } + + #[inline(always)] + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input),) + } +} + +impl Accessor1 for (A, B) +where + A: Accessor1, + B: Accessor1, +{ + type Output = (A::Output, B::Output); + + #[inline(always)] + fn push(&mut self, input: &[E]) { + self.0.push(input); + self.1.push(input); + } + + #[inline(always)] + fn finish(self, input: M) -> Self::Output { + (self.0.finish(input), self.1.finish(input)) + } +} + +pub struct FunctionalAccessor { + data: T, + p: P, + f: F, +} + +impl FunctionalAccessor { + #[inline(always)] + pub fn new(data: T, p: P, f: F) -> Self { + Self { data, p, f } + } +} + +impl Accessor1 for FunctionalAccessor +where + P: for<'a> FnMut(&'a mut T, &'a [E]), + F: FnOnce(T, M) -> R, +{ + type Output = R; + + #[inline(always)] + fn push(&mut self, input: &[E]) { + (self.p)(&mut self.data, input); + } + + #[inline(always)] + fn finish(self, input: M) -> Self::Output { + (self.f)(self.data, input) + } +} + +pub struct LAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LAccess<'a, E, M, A> { + #[inline(always)] + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for LAccess<'_, E0, M0, A> { + type Output = A::Output; + + #[inline(always)] + fn push(&mut self, rhs: &[E1]) { + let (lhs, elements) = self.elements.split_at(rhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + #[inline(always)] + fn finish(self, rhs: M1) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); + self.accessor.finish(self.metadata, rhs) + } +} + +pub struct RAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> RAccess<'a, E, M, A> { + #[inline(always)] + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> Accessor1 for RAccess<'_, E1, M1, A> { + type Output = A::Output; + + #[inline(always)] + fn push(&mut self, lhs: &[E0]) { + let (rhs, elements) = self.elements.split_at(lhs.len()); + self.accessor.push(lhs, rhs); + self.elements = elements; + } + + #[inline(always)] + fn finish(self, lhs: M0) -> Self::Output { + assert!(self.elements.is_empty(), "goal is shorter than expected"); + self.accessor.finish(lhs, self.metadata) + } +} + +pub trait TryAccessor1: Sized { + type Output; + #[must_use] + fn push(&mut self, input: &[E]) -> Option<()>; + #[must_use] + fn finish(self, input: M) -> Option; +} + +impl TryAccessor1 for () { + type Output = (); + + #[inline(always)] + fn push(&mut self, _: &[E]) -> Option<()> { + Some(()) + } + + #[inline(always)] + fn finish(self, _: M) -> Option { + Some(()) + } +} + +impl TryAccessor1 for (A,) +where + A: TryAccessor1, +{ + type Output = (A::Output,); + + #[inline(always)] + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + Some(()) + } + + #[inline(always)] + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?,)) + } +} + +impl TryAccessor1 for (A, B) +where + A: TryAccessor1, + B: TryAccessor1, +{ + type Output = (A::Output, B::Output); + + #[inline(always)] + fn push(&mut self, input: &[E]) -> Option<()> { + self.0.push(input)?; + self.1.push(input)?; + Some(()) + } + + #[inline(always)] + fn finish(self, input: M) -> Option { + Some((self.0.finish(input)?, self.1.finish(input)?)) + } +} + +pub struct LTryAccess<'a, E, M, A> { + elements: &'a [E], + metadata: M, + accessor: A, +} + +impl<'a, E, M, A> LTryAccess<'a, E, M, A> { + #[inline(always)] + pub fn new((elements, metadata): (&'a [E], M), accessor: A) -> Self { + Self { + elements, + metadata, + accessor, + } + } +} + +impl> TryAccessor1 + for LTryAccess<'_, E0, M0, A> +{ + type Output = A::Output; + + #[inline(always)] + fn push(&mut self, rhs: &[E1]) -> Option<()> { + let (lhs, elements) = self.elements.split_at_checked(rhs.len())?; + self.accessor.push(lhs, rhs); + self.elements = elements; + Some(()) + } + + #[inline(always)] + fn finish(self, rhs: M1) -> Option { + if !self.elements.is_empty() { + return None; + } + Some(self.accessor.finish(self.metadata, rhs)) + } +} + +#[derive(Debug)] +pub struct DistanceAccessor(f32, PhantomData V>, PhantomData D>); + +impl Default for DistanceAccessor { + #[inline(always)] + fn default() -> Self { + Self(0.0, PhantomData, PhantomData) + } +} + +impl Accessor2 for DistanceAccessor, L2S> { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_d2(target, input) + } + + #[inline(always)] + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[f32], input: &[f32]) { + self.0 += f32::reduce_sum_of_xy(target, input) + } + + #[inline(always)] + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} + +impl Accessor2 for DistanceAccessor, L2S> { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_d2(target, input) + } + + #[inline(always)] + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(self.0) + } +} + +impl Accessor2 for DistanceAccessor, Dot> { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[f16], input: &[f16]) { + self.0 += f16::reduce_sum_of_xy(target, input) + } + + #[inline(always)] + fn finish(self, (): (), (): ()) -> Self::Output { + Distance::from_f32(-self.0) + } +} + +pub trait DefaultWithDimension { + fn default_with_dimension(dim: u32) -> Self; +} + +impl DefaultWithDimension for T { + fn default_with_dimension(_dim: u32) -> Self { + Self::default() + } +} + +#[derive(Debug)] +pub struct ByteDistanceAccessor(u32, u32, PhantomData V>, PhantomData D>); + +impl DefaultWithDimension for ByteDistanceAccessor { + #[inline(always)] + fn default_with_dimension(dim: u32) -> Self { + Self(dim, 0, PhantomData, PhantomData) + } +} + +impl Accessor2 for ByteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::byte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_l2s( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + +impl Accessor2 for ByteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::byte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::byte::binary::half_process_dot( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + +#[derive(Debug)] +pub struct HalfbyteDistanceAccessor( + u32, + u32, + PhantomData V>, + PhantomData D>, +); + +impl DefaultWithDimension for HalfbyteDistanceAccessor { + #[inline(always)] + fn default_with_dimension(dim: u32) -> Self { + Self(dim, 0, PhantomData, PhantomData) + } +} + +impl Accessor2 for HalfbyteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::halfbyte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::halfbyte::binary::half_process_l2s( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} + +impl Accessor2 for HalfbyteDistanceAccessor { + type Output = Distance; + + #[inline(always)] + fn push(&mut self, target: &[u8], input: &[u8]) { + self.1 += rabitq::halfbyte::binary::accumulate(target, input); + } + + #[inline(always)] + fn finish(self, target: [f32; 4], input: [f32; 4]) -> Self::Output { + Distance::from_f32( + rabitq::halfbyte::binary::half_process_dot( + self.0, + self.1, + CodeMetadata::from_array(std::array::from_fn(|i| target[i])), + CodeMetadata::from_array(std::array::from_fn(|i| input[i])), + ) + .0, + ) + } +} diff --git a/crates/index/src/bump.rs b/crates/index/src/bump.rs new file mode 100644 index 00000000..8416a686 --- /dev/null +++ b/crates/index/src/bump.rs @@ -0,0 +1,32 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub trait Bump: 'static { + #[allow(clippy::mut_from_ref)] + fn alloc(&self, value: T) -> &mut T; + #[allow(clippy::mut_from_ref)] + fn alloc_slice(&self, slice: &[T]) -> &mut [T]; +} + +impl Bump for bumpalo::Bump { + #[inline] + fn alloc(&self, value: T) -> &mut T { + self.alloc(value) + } + + #[inline] + fn alloc_slice(&self, slice: &[T]) -> &mut [T] { + self.alloc_slice_copy(slice) + } +} diff --git a/crates/index/src/fetch.rs b/crates/index/src/fetch.rs new file mode 100644 index 00000000..a8ec1137 --- /dev/null +++ b/crates/index/src/fetch.rs @@ -0,0 +1,128 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::packed::PackedRefMut; +use always_equal::AlwaysEqual; + +pub type BorrowedIter<'b> = small_iter::borrowed::Iter<'b, u32, 1>; + +pub trait Fetch<'b> { + type Iter: ExactSizeIterator + 'b; + #[must_use] + fn fetch(&self) -> Self::Iter; +} + +impl Fetch<'_> for u32 { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + std::iter::once(*self) + } +} + +impl<'b, T, A, B, W: 'b + PackedRefMut)>> Fetch<'b> + for (T, AlwaysEqual) +{ + type Iter = BorrowedIter<'b>; + #[inline(always)] + fn fetch(&self) -> BorrowedIter<'b> { + let (.., list) = self.1.0.get(); + *list + } +} + +impl<'b, T, A, B, C> Fetch<'b> for (T, AlwaysEqual<&mut (A, B, C, BorrowedIter<'b>)>) { + type Iter = BorrowedIter<'b>; + #[inline(always)] + fn fetch(&self) -> BorrowedIter<'b> { + let (_, AlwaysEqual((.., list))) = self; + *list + } +} + +impl Fetch<'_> for (T, AlwaysEqual<(u32, u16)>) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + let (_, AlwaysEqual((x, _))) = self; + std::iter::once(*x) + } +} + +impl Fetch<'_> for (u32, u16) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + std::iter::once(self.0) + } +} + +impl Fetch<'_> for (T, AlwaysEqual<((u32, u16), (u32, u16))>) { + type Iter = std::iter::Once; + #[inline(always)] + fn fetch(&self) -> std::iter::Once { + let (_, AlwaysEqual(((x, _), _))) = self; + std::iter::once(*x) + } +} + +pub trait Fetch1 { + fn fetch_1(&self) -> u32; +} + +#[repr(transparent)] +pub struct Fetch1Iter { + iter: I, +} + +impl Iterator for Fetch1Iter +where + I::Item: Fetch1, +{ + type Item = u32; + + #[inline(always)] + fn next(&mut self) -> Option { + self.iter.next().map(|x| x.fetch_1()) + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } +} + +impl ExactSizeIterator for Fetch1Iter where I::Item: Fetch1 {} + +impl<'b, T, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<&'b [F]>) { + type Iter = Fetch1Iter>>; + + #[inline(always)] + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.iter().copied(), + } + } +} + +impl<'b, T, U, F: Fetch1 + Copy> Fetch<'b> for (T, AlwaysEqual<(&'b [F], U)>) { + type Iter = Fetch1Iter>>; + + #[inline(always)] + fn fetch(&self) -> Self::Iter { + Fetch1Iter { + iter: self.1.0.0.iter().copied(), + } + } +} diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs new file mode 100644 index 00000000..7513f10c --- /dev/null +++ b/crates/index/src/lib.rs @@ -0,0 +1,21 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +// pub mod accessor; +pub mod bump; +pub mod fetch; +pub mod packed; +pub mod prefetcher; +pub mod relation; +pub mod tuples; diff --git a/crates/index/src/packed.rs b/crates/index/src/packed.rs new file mode 100644 index 00000000..f8ee67c9 --- /dev/null +++ b/crates/index/src/packed.rs @@ -0,0 +1,61 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub trait PackedRefMut { + type T; + fn get(&self) -> &Self::T; + fn get_mut(&mut self) -> &mut Self::T; +} + +impl PackedRefMut for &mut T { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self + } +} + +#[repr(Rust, packed(4))] +pub struct PackedRefMut4<'b, T>(pub &'b mut T); + +impl<'b, T> PackedRefMut for PackedRefMut4<'b, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} + +#[repr(Rust, packed(8))] +pub struct PackedRefMut8<'b, T>(pub &'b mut T); + +impl<'a, T> PackedRefMut for PackedRefMut8<'a, T> { + type T = T; + #[inline(always)] + fn get(&self) -> &T { + self.0 + } + #[inline(always)] + fn get_mut(&mut self) -> &mut T { + self.0 + } +} diff --git a/crates/index/src/prefetcher.rs b/crates/index/src/prefetcher.rs new file mode 100644 index 00000000..626d4a64 --- /dev/null +++ b/crates/index/src/prefetcher.rs @@ -0,0 +1,419 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::fetch::Fetch; +use crate::relation::{ + Hints, ReadStream, RelationPrefetch, RelationRead, RelationReadStream, RelationReadTypes, +}; +use dary_heap::DaryHeap; +use std::collections::{BinaryHeap, VecDeque}; + +pub const WINDOW_SIZE: usize = 32; +const _: () = assert!(WINDOW_SIZE > 0); + +pub trait Sequence { + type Item; + type Inner: Iterator; + #[must_use] + fn next(&mut self) -> Option; + #[must_use] + fn peek(&mut self) -> Option<&Self::Item>; + #[must_use] + fn into_inner(self) -> Self::Inner; +} + +impl Sequence for BinaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + +impl Sequence for DaryHeap { + type Item = T; + type Inner = std::vec::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + (self as &Self).peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_vec().into_iter() + } +} + +impl Sequence for std::iter::Peekable { + type Item = I::Item; + type Inner = std::iter::Peekable; + #[inline] + fn next(&mut self) -> Option { + Iterator::next(self) + } + #[inline] + fn peek(&mut self) -> Option<&I::Item> { + self.peek() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self + } +} + +impl Sequence for VecDeque { + type Item = T; + type Inner = std::collections::vec_deque::IntoIter; + #[inline] + fn next(&mut self) -> Option { + self.pop_front() + } + #[inline] + fn peek(&mut self) -> Option<&T> { + self.front() + } + #[inline] + fn into_inner(self) -> Self::Inner { + self.into_iter() + } +} + +pub trait Prefetcher<'b>: IntoIterator +where + Self::Item: Fetch<'b>, +{ + type R: RelationRead + 'b; + type Guards: ExactSizeIterator::ReadGuard<'b>>; + + #[must_use] + fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; + #[must_use] + fn next_if( + &mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(Self::Item, Self::Guards)>; +} + +pub struct PlainPrefetcher<'b, R, S: Sequence> { + relation: &'b R, + sequence: S, +} + +impl<'b, R, S: Sequence> PlainPrefetcher<'b, R, S> { + #[inline] + pub fn new(relation: &'b R, sequence: S) -> Self { + Self { relation, sequence } + } +} + +impl<'b, R, S: Sequence> IntoIterator for PlainPrefetcher<'b, R, S> { + type Item = S::Item; + + type IntoIter = S::Inner; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.sequence.into_inner() + } +} + +impl<'b, R: RelationRead, S: Sequence> Prefetcher<'b> for PlainPrefetcher<'b, R, S> +where + S::Item: Fetch<'b>, +{ + type R = R; + type Guards = PlainPrefetcherGuards<'b, R, >::Iter>; + + #[inline] + fn next( + &mut self, + ) -> Option<( + Self::Item, + PlainPrefetcherGuards<'b, R, >::Iter>, + )> { + let e = self.sequence.next()?; + let list = e.fetch(); + Some(( + e, + PlainPrefetcherGuards { + relation: self.relation, + list, + }, + )) + } + + #[inline] + fn next_if( + &mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<( + S::Item, + PlainPrefetcherGuards<'b, R, >::Iter>, + )> { + if !predicate(self.sequence.peek()?) { + return None; + } + let e = self.sequence.next()?; + let list = e.fetch(); + Some(( + e, + PlainPrefetcherGuards { + relation: self.relation, + list, + }, + )) + } +} + +pub struct PlainPrefetcherGuards<'b, R, L> { + relation: &'b R, + list: L, +} + +impl<'b, R: RelationRead, L: Iterator> Iterator for PlainPrefetcherGuards<'b, R, L> { + type Item = R::ReadGuard<'b>; + + #[inline] + fn next(&mut self) -> Option { + let id = self.list.next()?; + Some(self.relation.read(id)) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } +} + +impl<'b, R: RelationRead, L: Iterator> ExactSizeIterator + for PlainPrefetcherGuards<'b, R, L> +{ +} + +pub struct SimplePrefetcher<'b, R, S: Sequence> { + relation: &'b R, + window: VecDeque, + sequence: S, +} + +impl<'b, R, S: Sequence> SimplePrefetcher<'b, R, S> { + #[inline] + pub fn new(relation: &'b R, sequence: S) -> Self { + Self { + relation, + window: VecDeque::new(), + sequence, + } + } +} + +impl<'b, R, S: Sequence> IntoIterator for SimplePrefetcher<'b, R, S> { + type Item = S::Item; + + type IntoIter = std::iter::Chain, S::Inner>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.window.into_iter().chain(self.sequence.into_inner()) + } +} + +impl<'b, R: RelationRead + RelationPrefetch, S: Sequence> Prefetcher<'b> + for SimplePrefetcher<'b, R, S> +where + S::Item: Fetch<'b>, +{ + type R = R; + type Guards = SimplePrefetcherGuards<'b, R, >::Iter>; + + #[inline] + fn next( + &mut self, + ) -> Option<( + Self::Item, + SimplePrefetcherGuards<'b, R, >::Iter>, + )> { + while self.window.len() < WINDOW_SIZE + && let Some(e) = self.sequence.next() + { + for id in e.fetch() { + self.relation.prefetch(id); + } + self.window.push_back(e); + } + let e = self.window.pop_front()?; + let list = e.fetch(); + Some(( + e, + SimplePrefetcherGuards { + relation: self.relation, + list, + }, + )) + } + + #[inline] + fn next_if( + &mut self, + predicate: impl FnOnce(&S::Item) -> bool, + ) -> Option<( + S::Item, + SimplePrefetcherGuards<'b, R, >::Iter>, + )> { + while self.window.len() < WINDOW_SIZE + && let Some(e) = self.sequence.next() + { + for id in e.fetch() { + self.relation.prefetch(id); + } + self.window.push_back(e); + } + let e = self.window.pop_front_if(move |x| predicate(x))?; + let list = e.fetch(); + Some(( + e, + SimplePrefetcherGuards { + relation: self.relation, + list, + }, + )) + } +} + +pub struct SimplePrefetcherGuards<'b, R, L> { + relation: &'b R, + list: L, +} + +impl<'b, R: RelationRead, L: Iterator> Iterator for SimplePrefetcherGuards<'b, R, L> { + type Item = R::ReadGuard<'b>; + + #[inline] + fn next(&mut self) -> Option { + let id = self.list.next()?; + Some(self.relation.read(id)) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } +} + +impl<'b, R: RelationRead, L: ExactSizeIterator> ExactSizeIterator + for SimplePrefetcherGuards<'b, R, L> +{ +} + +pub struct StreamPrefetcherSequence(S); + +impl Iterator for StreamPrefetcherSequence { + type Item = S::Item; + + #[inline] + fn next(&mut self) -> Option { + self.0.next() + } +} + +pub struct StreamPrefetcher<'b, R: RelationReadStream + 'b, S: Sequence> +where + S::Item: Fetch<'b>, +{ + stream: R::ReadStream<'b, StreamPrefetcherSequence>, +} + +impl<'b, R: RelationReadStream, S: Sequence> StreamPrefetcher<'b, R, S> +where + S::Item: Fetch<'b>, +{ + #[inline] + pub fn new(relation: &'b R, sequence: S, hints: Hints) -> Self { + let stream = relation.read_stream(StreamPrefetcherSequence(sequence), hints); + Self { stream } + } +} + +impl<'b, R: RelationReadStream, S: Sequence> IntoIterator for StreamPrefetcher<'b, R, S> +where + S::Item: Fetch<'b>, +{ + type Item = S::Item; + + type IntoIter = > as ReadStream<'b>>::Inner; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.stream.into_inner() + } +} + +impl<'b, R: RelationRead + RelationReadStream, S: Sequence> Prefetcher<'b> + for StreamPrefetcher<'b, R, S> +where + S::Item: Fetch<'b>, +{ + type R = R; + + type Guards = > as ReadStream<'b>>::Guards; + + #[inline] + fn next(&mut self) -> Option<(S::Item, Self::Guards)> { + self.stream.next() + } + + #[inline] + fn next_if( + &mut self, + predicate: impl FnOnce(&Self::Item) -> bool, + ) -> Option<(S::Item, Self::Guards)> { + self.stream.next_if(predicate) + } +} + +pub trait PrefetcherSequenceFamily<'b, R> { + type P: Prefetcher<'b, R = R, Item = S::Item> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>; + + fn is_not_plain(&self) -> bool; +} + +pub trait PrefetcherHeapFamily<'b, R> { + type P: Prefetcher<'b, R = R, Item = T> + where + T: Ord + Fetch<'b>; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch<'b>; + + fn is_not_plain(&self) -> bool; +} diff --git a/crates/index/src/relation.rs b/crates/index/src/relation.rs new file mode 100644 index 00000000..d6c4e832 --- /dev/null +++ b/crates/index/src/relation.rs @@ -0,0 +1,136 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::fetch::Fetch; +use std::ops::{Deref, DerefMut}; +use zerocopy::{FromBytes, IntoBytes}; + +/// # Safety +/// +/// * `Opaque` must aligned to 8 bytes. +#[allow(unsafe_code)] +pub unsafe trait Opaque: Copy + Send + Sync + FromBytes + IntoBytes + 'static {} + +pub trait Page: Sized + 'static { + type Opaque: Opaque; + + #[must_use] + fn get_opaque(&self) -> &Self::Opaque; + #[must_use] + fn get_opaque_mut(&mut self) -> &mut Self::Opaque; + #[must_use] + fn len(&self) -> u16; + #[must_use] + #[inline] + fn is_empty(&self) -> bool { + self.len() == 0 + } + #[must_use] + fn get(&self, i: u16) -> Option<&[u8]>; + #[must_use] + fn get_mut(&mut self, i: u16) -> Option<&mut [u8]>; + #[must_use] + fn alloc(&mut self, data: &[u8]) -> Option; + fn free(&mut self, i: u16); + #[must_use] + fn freespace(&self) -> u16; + fn clear(&mut self, opaque: Self::Opaque); +} + +pub trait PageGuard { + fn id(&self) -> u32; +} + +pub trait ReadStream<'b> { + type Relation: RelationReadTypes; + type Guards: ExactSizeIterator::ReadGuard<'b>>; + type Item; + type Inner: Iterator; + fn next(&mut self) -> Option<(Self::Item, Self::Guards)>; + fn next_if bool>( + &mut self, + predicate: P, + ) -> Option<(Self::Item, Self::Guards)>; + fn into_inner(self) -> Self::Inner; +} + +pub trait Relation { + type Page: Page; +} + +pub trait RelationReadTypes: Relation { + type ReadGuard<'b>: PageGuard + Deref; +} + +pub trait RelationRead: RelationReadTypes { + fn read(&self, id: u32) -> Self::ReadGuard<'_>; +} + +pub trait RelationWriteTypes: Relation { + type WriteGuard<'b>: PageGuard + DerefMut; +} + +pub trait RelationWrite: RelationWriteTypes { + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_>; + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_>; + fn search(&self, freespace: usize) -> Option>; +} + +pub trait RelationPrefetch: Relation { + fn prefetch(&self, id: u32); +} + +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct Hints { + pub full: bool, + pub batch: bool, +} + +impl Default for Hints { + #[inline] + fn default() -> Self { + Self { + full: false, + batch: false, + } + } +} + +impl Hints { + #[inline] + pub fn full(self, full: bool) -> Self { + Self { full, ..self } + } + #[inline] + pub fn batch(self, batch: bool) -> Self { + Self { batch, ..self } + } +} + +pub trait RelationReadStreamTypes: RelationReadTypes { + type ReadStream<'b, I: Iterator>: ReadStream<'b, Item = I::Item, Relation = Self> + where + I::Item: Fetch<'b>; +} + +pub trait RelationReadStream: RelationReadStreamTypes { + fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> + where + I::Item: Fetch<'b>; +} diff --git a/crates/index/src/tuples.rs b/crates/index/src/tuples.rs new file mode 100644 index 00000000..0934ce7e --- /dev/null +++ b/crates/index/src/tuples.rs @@ -0,0 +1,237 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::marker::PhantomData; +use std::num::NonZero; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; + +pub struct RefChecker<'a> { + bytes: &'a [u8], +} + +impl<'a> RefChecker<'a> { + #[inline(always)] + pub fn new(bytes: &'a [u8]) -> Self { + Self { bytes } + } + #[inline] + pub fn prefix( + &self, + s: impl Into + Copy, + ) -> &'a T { + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") + } + #[inline] + pub fn bytes( + &self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a T { + let start = Into::::into(s); + let end = Into::::into(e); + let bytes = &self.bytes[start..end]; + FromBytes::ref_from_bytes(bytes).expect("deserialization: bad bytes") + } + /// # Safety + /// + /// * `FromBytes` could be implemented for `T`. + #[inline] + #[allow(unsafe_code)] + pub unsafe fn bytes_slice_unchecked( + &self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a [T] { + let start = Into::::into(s); + let end = Into::::into(e); + let bytes = &self.bytes[start..end]; + if size_of::() == 0 || bytes.len() % size_of::() == 0 { + let ptr = bytes as *const [u8] as *const T; + if ptr.is_aligned() { + unsafe { std::slice::from_raw_parts(ptr, bytes.len() / size_of::()) } + } else { + panic!("deserialization: bad bytes") + } + } else { + panic!("deserialization: bad bytes") + } + } +} + +pub struct MutChecker<'a> { + flag: usize, + bytes: *mut [u8], + phantom: PhantomData<&'a mut [u8]>, +} + +impl<'a> MutChecker<'a> { + #[inline(always)] + pub fn new(bytes: &'a mut [u8]) -> Self { + Self { + flag: 0, + bytes, + phantom: PhantomData, + } + } + #[inline] + pub fn prefix( + &mut self, + s: impl Into + Copy, + ) -> &'a mut T { + let start = Into::::into(s); + let end = Into::::into(s) + size_of::(); + if !(start <= end && end <= self.bytes.len()) { + panic!("deserialization: bad bytes"); + } + if !(self.flag <= start) { + panic!("deserialization: bad bytes"); + } else { + self.flag = end; + } + #[allow(unsafe_code)] + let bytes = unsafe { + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") + } + #[inline] + pub fn bytes( + &mut self, + s: impl Into + Copy, + e: impl Into + Copy, + ) -> &'a mut T { + let start = Into::::into(s); + let end = Into::::into(e); + if !(start <= end && end <= self.bytes.len()) { + panic!("deserialization: bad bytes"); + } + if !(self.flag <= start) { + panic!("deserialization: bad bytes"); + } else { + self.flag = end; + } + #[allow(unsafe_code)] + let bytes = unsafe { + std::slice::from_raw_parts_mut((self.bytes as *mut u8).add(start), end - start) + }; + FromBytes::mut_from_bytes(bytes).expect("deserialization: bad bytes") + } +} + +#[test] +fn aliasing_test() { + #[repr(C, align(8))] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct ExampleHeader { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], + } + let serialized = { + let elements = (0u32..1111).collect::>(); + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + buffer[..size_of::()].copy_from_slice( + ExampleHeader { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + }; + let mut source = vec![0u64; serialized.len().next_multiple_of(8)]; + source.as_mut_bytes()[..serialized.len()].copy_from_slice(&serialized); + let deserialized = { + let mut checker = MutChecker::new(source.as_mut_bytes()); + let header: &mut ExampleHeader = checker.prefix(0_u16); + let elements: &mut [u32] = checker.bytes(header.elements_s, header.elements_e); + (header, elements) + }; + assert_eq!( + deserialized.1, + (0u32..1111).collect::>().as_slice() + ); +} + +#[repr(transparent)] +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Padding(Option>); + +impl Padding { + pub const ZERO: Self = Self(None); +} + +#[repr(transparent)] +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Bool(u8); + +impl Bool { + pub const FALSE: Self = Self(0); + pub const TRUE: Self = Self(1); +} + +impl From for bool { + #[inline] + fn from(value: Bool) -> Self { + value != Bool::FALSE + } +} + +impl From for Bool { + #[inline] + fn from(value: bool) -> Self { + std::hint::select_unpredictable(value, Self::TRUE, Self::FALSE) + } +} diff --git a/crates/index_accessor/Cargo.toml b/crates/index_accessor/Cargo.toml new file mode 100644 index 00000000..e23da0f3 --- /dev/null +++ b/crates/index_accessor/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "index_accessor" +version.workspace = true +edition.workspace = true +publish = false + +[lib] +path = "../index/src/accessor.rs" + +[dependencies] +distance = { path = "../distance" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +[lints] +workspace = true diff --git a/crates/k_means/Cargo.toml b/crates/k_means/Cargo.toml new file mode 100644 index 00000000..0ddf105f --- /dev/null +++ b/crates/k_means/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "k_means" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } + +rand.workspace = true +rayon.workspace = true + +[lints] +workspace = true diff --git a/crates/k_means/src/flat.rs b/crates/k_means/src/flat.rs new file mode 100644 index 00000000..3ebf88ec --- /dev/null +++ b/crates/k_means/src/flat.rs @@ -0,0 +1,112 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::{flat_index as prefect_index, flat_index as index}; +use crate::square::{Square, SquareMut}; +use crate::{KMeans, This}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use rayon::prelude::*; + +struct Flat<'a> { + this: This<'a>, + samples: SquareMut<'a>, + centroids: Square, + targets: Vec, +} + +impl<'a> KMeans for Flat<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.centroids)) + } + + fn assign(&mut self) { + let this = &mut self.this; + let samples = &mut self.samples; + let centroids = &self.centroids; + let targets = &mut self.targets; + let index = index(centroids); + this.pool.install(|| { + targets + .par_iter_mut() + .zip(samples.par_iter_mut()) + .for_each(|(target, sample)| { + *target = index(sample).1; + }); + }); + } + + fn update(&mut self) { + crate::index::update( + &mut self.this, + &self.samples, + &self.targets, + &mut self.centroids, + ); + } + + fn finish(self: Box) -> Square { + self.centroids + } +} + +pub fn new<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + let mut rng = StdRng::from_seed(seed); + + let mut centroids = Square::with_capacity(d, c); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); + + let targets = vec![0; samples.len()]; + + Box::new(Flat { + this: This { + pool, + rng, + d, + c, + is_spherical, + }, + samples, + centroids, + targets, + }) +} diff --git a/crates/k_means/src/hierarchical.rs b/crates/k_means/src/hierarchical.rs new file mode 100644 index 00000000..7531db64 --- /dev/null +++ b/crates/k_means/src/hierarchical.rs @@ -0,0 +1,321 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::square::{Square, SquareMut}; +use crate::{KMeans, This}; +use always_equal::AlwaysEqual; +use distance::Distance; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator}; +use std::collections::BinaryHeap; + +struct Hierarchical<'a> { + this: This<'a>, + partitions: Vec>, + coarse_centroids: Square, + offsets: Vec, +} + +impl<'a> KMeans for Hierarchical<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.partitions, &self.offsets)) + } + + fn index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(index( + &self.partitions, + &self.coarse_centroids, + &self.offsets, + )) + } + + fn assign(&mut self) { + for partial in self.partitions.iter_mut() { + partial.assign(); + } + } + + fn update(&mut self) { + for partial in self.partitions.iter_mut() { + partial.update(); + } + } + + fn finish(self: Box) -> Square { + let mut centroids = Square::new(self.this.d); + for k_means in self.partitions { + let partial_centroids = k_means.finish(); + for centroid in partial_centroids.into_iter() { + centroids.push_slice(centroid); + } + } + centroids + } +} + +fn prefect_index( + partitions: &[Box], + offsets: &[usize], +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + let indexes = partitions + .iter() + .map(|p| p.prefect_index()) + .collect::>(); + move |sample| { + let mut result = (f32::INFINITY, 0); + for (id, index) in indexes.iter().enumerate() { + let partial_result = index(sample); + if partial_result.0 <= result.0 { + result = (partial_result.0, offsets[id] + partial_result.1); + } + } + result + } +} + +fn index( + partitions: &[Box], + coarse_centroids: &Square, + offsets: &[usize], +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + let indexes = partitions.iter().map(|p| p.index()).collect::>(); + move |sample| { + use simd::Floating; + let mut result = (f32::INFINITY, 0); + for i in 0..coarse_centroids.len() { + let dis = f32::reduce_sum_of_d2(sample, &coarse_centroids[i]); + if dis <= result.0 { + result = (dis, i); + } + } + let id = result.1; + let result = indexes[id](sample); + (result.0, offsets[id] + result.1) + } +} + +const COARSE_SAMPLING_FACTOR: usize = 256; +const COARSE_ITERATIONS: usize = 10; + +pub fn new<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + mut samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + let mut rng = StdRng::from_seed(seed); + + let mut coarse_samples = { + let mut coarse_samples = Square::new(d); + let s = c.isqrt().saturating_mul(COARSE_SAMPLING_FACTOR); + for index in rand::seq::index::sample(&mut rng, samples.len(), s.min(samples.len())) { + coarse_samples.push_slice(&samples[index]); + } + coarse_samples + }; + let coarse_k_means = { + let mut coarse_k_means = crate::lloyd_k_means( + pool, + d, + coarse_samples.as_mut_view(), + c.isqrt(), + seed, + is_spherical, + ); + for _ in 0..COARSE_ITERATIONS { + coarse_k_means.assign(); + coarse_k_means.update(); + } + coarse_k_means + }; + let coarse_assign = { + let coarse_index = coarse_k_means.prefect_index(); + let mut coarse_assign = vec![0; samples.len()]; + pool.install(|| { + coarse_assign + .par_iter_mut() + .zip(samples.par_iter_mut()) + .for_each(|(target, sample)| { + *target = coarse_index(sample).1; + }); + }); + coarse_assign + }; + let coarse_centroids = coarse_k_means.finish(); + + let (weight, groups) = { + let mut weight = vec![0_usize; coarse_centroids.len()]; + let mut groups = vec![vec![]; coarse_centroids.len()]; + for (i, &target) in coarse_assign.iter().enumerate() { + weight[target] += 1; + groups[target].push(i); + } + (weight, groups) + }; + let seats = modified_webster_method(c, &weight); + + let mut partitions = vec![]; + let mut offsets = vec![]; + let mut offset = 0; + for (samples, c) in std::iter::zip(partition(samples, &groups), seats) { + partitions.push(crate::lloyd_k_means( + pool, + d, + samples, + c, + seed, + is_spherical, + )); + offsets.push(offset); + offset += c; + } + + Box::new(Hierarchical { + this: This { + pool, + d, + c, + rng, + is_spherical, + }, + coarse_centroids, + partitions, + offsets, + }) +} + +// https://en.wikipedia.org/wiki/Sainte-Lagu%C3%AB_method +fn modified_webster_method(n: usize, weight: &[usize]) -> Vec { + assert!(n >= weight.len()); + let mut seats = vec![0_usize; weight.len()]; + let mut quotients = Vec::new(); + for index in 0..weight.len() { + seats[index] += 1; + let quotient = weight[index] as f64 / (seats[index] as f64 + 0.5); + quotients.push((Distance::from_f32(quotient as _), AlwaysEqual(index))); + } + let mut quotients = BinaryHeap::<_>::from(quotients); + for _ in weight.len()..n { + let Some((_, AlwaysEqual(index))) = quotients.pop() else { + break; + }; + seats[index] += 1; + let quotient = weight[index] as f64 / (seats[index] as f64 + 0.5); + quotients.push((Distance::from_f32(quotient as _), AlwaysEqual(index))); + } + seats +} + +fn partition<'a>(mut a: SquareMut<'a>, groups: &[impl AsRef<[usize]>]) -> Vec> { + let n = a.len(); + let permutation = groups + .iter() + .flat_map(AsRef::as_ref) + .copied() + .collect::>(); + let mut marked = vec![false; a.len()]; + let mut buffer = vec![0.0; a.d()]; + for i in 0..n { + if marked[i] { + continue; + } + buffer.copy_from_slice(&a[i]); + let (mut src, mut dst) = (permutation[i], i); + while src != i { + a.copy_within(src..src + 1, dst); + marked[dst] = true; + (src, dst) = (permutation[src], src); + } + a[dst].copy_from_slice(&buffer); + marked[dst] = true; + } + let mut result = Vec::with_capacity(groups.len()); + let (d, mut p) = a.into_inner(); + for group in groups { + let group = group.as_ref(); + let head; + (head, p) = std::mem::take(&mut p).split_at_mut(group.len() * d); + result.push(SquareMut::new(d, head)); + } + result +} + +#[test] +fn test_modified_webster_method() { + let seats = modified_webster_method(51, &[10, 10, 10, 11, 9]); + assert_eq!(seats[0], 10); + assert_eq!(seats[1], 10); + assert_eq!(seats[2], 10); + assert_eq!(seats[3], 12); + assert_eq!(seats[4], 9); +} + +#[test] +fn test_partition() { + fn gen_random_alloc(rows: usize, groups: usize, rng: &mut impl RngExt) -> Vec> { + let mut idx: Vec = (0..rows).collect(); + idx.shuffle(rng); + let mut alloc = Vec::with_capacity(groups); + let mut start = 0; + for g in 0..groups { + let rem = groups - g; + let take = if rem == 1 { + rows - start + } else { + rng.random_range(1..=(rows - start - (rem - 1))) + }; + alloc.push(idx[start..start + take].to_vec()); + start += take; + } + alloc + } + + use rand::prelude::*; + let mut rng = StdRng::seed_from_u64(7); + for trial in 0..if cfg!(not(miri)) { 1000 } else { 1 } { + let d = rng.random_range(1..10); + let rows = rng.random_range(1000..2000); + let groups = rng.random_range(10..=rows.min(20)); + let mut s = { + let mut result = Square::with_capacity(d, rows); + for _ in 0..rows { + result.push_iter((0..d).map(|_| rng.random_range(-1000.0..1000.0))); + } + result + }; + let golden: Vec> = (0..s.len()).map(|i| s[i].to_vec()).collect(); + let alloc = gen_random_alloc(rows, groups, &mut rng); + let views = partition(s.as_mut_view(), &alloc); + assert_eq!(views.len(), alloc.len(), "trial {}", trial); + + for (g, group) in alloc.iter().enumerate() { + let v = &views[g]; + assert_eq!(v.len(), group.len(), "trial {}, group {}", trial, g); + assert_eq!(v.d(), d); + for (r, &row_idx) in group.iter().enumerate() { + assert_eq!( + v.row(r), + &golden[row_idx][..], + "trial {}, group {}, row {}", + trial, + g, + r + ); + } + } + } +} diff --git a/crates/k_means/src/index.rs b/crates/k_means/src/index.rs new file mode 100644 index 00000000..6da6524d --- /dev/null +++ b/crates/k_means/src/index.rs @@ -0,0 +1,157 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::This; +use crate::square::{Square, SquareMut}; +use rand::RngExt; +use rayon::prelude::*; +use simd::Floating; + +pub fn flat_index(centroids: &Square) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + move |sample| { + let mut result = (f32::INFINITY, 0); + for (i, centroid) in centroids.into_iter().enumerate() { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + result + } +} + +pub fn rabitq_index( + pool: &rayon::ThreadPool, + centroids: &Square, +) -> impl Fn(&[f32]) -> (f32, usize) + Sync { + use rabitq::packing::{pack_to_u4, padding_pack}; + let (metadata, blocks) = pool.install(|| { + let metadata = centroids + .par_iter() + .map(rabitq::bit::code_metadata) + .collect::>(); + + let blocks = centroids + .par_iter() + .chunks(32) + .map(|chunk| { + let f = |x: &&_| pack_to_u4(&rabitq::bit::code_elements(x)); + padding_pack(chunk.iter().map(f)) + }) + .collect::>(); + + (metadata, blocks) + }); + move |sample| { + let lut = rabitq::bit::block::preprocess(sample); + let mut result = (f32::INFINITY, 0); + let mut sum = [0u32; 32]; + for (i, centroid) in centroids.into_iter().enumerate() { + if i % 32 == 0 { + sum = rabitq::bit::block::accumulate(&blocks[i / 32], &lut.1); + } + let (rough, err) = + rabitq::bit::block::half_process_l2s(sum[i % 32], metadata[i], lut.0); + let lowerbound = rough - err * 1.9; + if lowerbound < result.0 { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + } + result + } +} + +pub fn update( + this: &mut This<'_>, + samples: &SquareMut<'_>, + targets: &[usize], + centroids: &mut Square, +) { + this.pool.install(|| { + const DELTA: f32 = 9.7656e-4_f32; + + let d = this.d; + let n = samples.len(); + let c = this.c; + + let list = rayon::broadcast({ + |ctx| { + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for i in (ctx.index()..samples.len()).step_by(ctx.num_threads()) { + let target = targets[i]; + let sample = &samples[i]; + f32::vector_add_inplace(&mut sum[target], sample); + count[target] += 1.0; + } + (sum, count) + } + }); + let mut sum = Square::from_zeros(d, c); + let mut count = vec![0.0f32; c]; + for (sum_1, count_1) in list { + for i in 0..c { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + } + + sum.par_iter_mut() + .zip(count.par_iter()) + .for_each(|(sum, count)| f32::vector_mul_scalar_inplace(sum, 1.0 / count)); + + *centroids = sum; + + for i in 0..c { + if count[i] != 0.0f32 { + continue; + } + let mut o = 0; + loop { + let alpha = this.rng.random_range(0.0..1.0f32); + let beta = (count[o] - 1.0) / (n - c) as f32; + if alpha < beta { + break; + } + o = (o + 1) % c; + } + centroids.copy_within(o..o + 1, i); + vector_mul_scalars_inplace(&mut centroids[i], [1.0 + DELTA, 1.0 - DELTA]); + vector_mul_scalars_inplace(&mut centroids[o], [1.0 - DELTA, 1.0 + DELTA]); + count[i] = count[o] / 2.0; + count[o] -= count[i]; + } + + if this.is_spherical { + centroids.into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); +} + +fn vector_mul_scalars_inplace(this: &mut [f32], scalars: [f32; 2]) { + let n: usize = this.len(); + for i in 0..n { + if i % 2 == 0 { + this[i] *= scalars[0]; + } else { + this[i] *= scalars[1]; + } + } +} diff --git a/crates/k_means/src/lib.rs b/crates/k_means/src/lib.rs new file mode 100644 index 00000000..510a62e7 --- /dev/null +++ b/crates/k_means/src/lib.rs @@ -0,0 +1,85 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod flat; +mod hierarchical; +mod index; +mod quick; +mod rabitq; + +pub mod square; + +use crate::square::{Square, SquareMut}; +use rand::rngs::StdRng; + +pub struct This<'a> { + pool: &'a rayon::ThreadPool, + rng: StdRng, + d: usize, + c: usize, + is_spherical: bool, +} + +pub trait KMeans { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_>; + fn index(&self) -> Box (f32, usize) + Sync + '_> { + self.prefect_index() + } + fn assign(&mut self); + fn update(&mut self); + fn finish(self: Box) -> Square; +} + +pub fn k_means_lookup(sample: &[f32], centroids: &Square) -> usize { + use simd::Floating; + let mut result = (f32::INFINITY, 0); + for (i, centroid) in centroids.into_iter().enumerate() { + let dis = f32::reduce_sum_of_d2(sample, centroid); + if dis <= result.0 { + result = (dis, i); + } + } + result.1 +} + +pub fn lloyd_k_means<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + assert!(d > 0 && c > 0); + let n = samples.len(); + if n <= c { + quick::new(pool, d, samples, c, seed, is_spherical) + } else if n <= c * 2 { + flat::new(pool, d, samples, c, seed, is_spherical) + } else { + rabitq::new(pool, d, samples, c, seed, is_spherical) + } +} + +pub fn hierarchical_k_means<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + assert!(d > 0 && c > 0); + hierarchical::new(pool, d, samples, c, seed, is_spherical) +} diff --git a/crates/k_means/src/quick.rs b/crates/k_means/src/quick.rs new file mode 100644 index 00000000..ac69da6a --- /dev/null +++ b/crates/k_means/src/quick.rs @@ -0,0 +1,75 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::KMeans; +use crate::index::flat_index as prefect_index; +use crate::square::{Square, SquareMut}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +struct Quick { + centroids: Square, +} + +impl KMeans for Quick { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + Box::new(prefect_index(&self.centroids)) + } + + fn assign(&mut self) {} + + fn update(&mut self) {} + + fn finish(self: Box) -> Square { + self.centroids + } +} + +pub fn new<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + let mut rng = StdRng::from_seed(seed); + + let mut centroids = Square::with_capacity(d, c); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); + + Box::new(Quick { centroids }) +} diff --git a/crates/k_means/src/rabitq.rs b/crates/k_means/src/rabitq.rs new file mode 100644 index 00000000..61999979 --- /dev/null +++ b/crates/k_means/src/rabitq.rs @@ -0,0 +1,137 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::{flat_index as prefect_index, rabitq_index as index}; +use crate::square::{Square, SquareMut}; +use crate::{KMeans, This}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use rayon::prelude::*; + +struct RaBitQ<'a> { + this: This<'a>, + samples: SquareMut<'a>, + centroids: Square, + targets: Vec, +} + +impl<'a> KMeans for RaBitQ<'a> { + fn prefect_index(&self) -> Box (f32, usize) + Sync + '_> { + let index = prefect_index(&self.centroids); + Box::new(move |sample| { + let rotated = rabitq::rotate::rotate(sample); + let sample = rotated.as_slice(); + index(sample) + }) + } + + fn index(&self) -> Box (f32, usize) + Sync + '_> { + let index = index(self.this.pool, &self.centroids); + Box::new(move |sample| { + let rotated = rabitq::rotate::rotate(sample); + let sample = rotated.as_slice(); + index(sample) + }) + } + + fn assign(&mut self) { + let this = &mut self.this; + let samples = &mut self.samples; + let centroids = &self.centroids; + let targets = &mut self.targets; + let index = index(this.pool, centroids); + this.pool.install(|| { + targets + .par_iter_mut() + .zip(samples.par_iter_mut()) + .for_each(|(target, sample)| { + *target = index(sample).1; + }); + }); + } + + fn update(&mut self) { + crate::index::update( + &mut self.this, + &self.samples, + &self.targets, + &mut self.centroids, + ); + } + + fn finish(mut self: Box) -> Square { + self.this.pool.install(|| { + self.centroids.par_iter_mut().for_each(|centroid| { + rabitq::rotate::rotate_reversed_inplace(centroid); + }); + }); + self.centroids + } +} + +pub fn new<'a>( + pool: &'a rayon::ThreadPool, + d: usize, + mut samples: SquareMut<'a>, + c: usize, + seed: [u8; 32], + is_spherical: bool, +) -> Box { + let mut rng = StdRng::from_seed(seed); + + pool.install(|| { + samples.par_iter_mut().for_each(|sample| { + rabitq::rotate::rotate_inplace(sample); + }); + }); + + let mut centroids = Square::with_capacity(d, c); + + for index in rand::seq::index::sample(&mut rng, samples.len(), c.min(samples.len())) { + centroids.push_slice(&samples[index]); + } + + if centroids.is_empty() && c == 1 { + centroids.push_iter(std::iter::repeat_n(0.0, d as _)); + } + + while centroids.len() < c { + centroids.push_iter((0..d).map(|_| rng.random_range(-1.0f32..1.0f32))); + } + + pool.install(|| { + if is_spherical { + use simd::Floating; + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + }); + + let targets = vec![0; samples.len()]; + + Box::new(RaBitQ { + this: This { + pool, + d, + c, + rng, + is_spherical, + }, + samples, + centroids, + targets, + }) +} diff --git a/crates/k_means/src/square.rs b/crates/k_means/src/square.rs new file mode 100644 index 00000000..769f1fde --- /dev/null +++ b/crates/k_means/src/square.rs @@ -0,0 +1,214 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[derive(Debug, Clone)] +pub struct Square { + d: usize, + p: Vec, +} + +impl Square { + pub fn d(&self) -> usize { + self.d + } + pub fn len(&self) -> usize { + self.p.len() / self.d + } + pub fn is_empty(&self) -> bool { + self.p.is_empty() + } + pub fn with_capacity(d: usize, p: usize) -> Self { + Self { + d, + p: Vec::with_capacity(usize::saturating_mul(d, p)), + } + } + pub fn new(d: usize) -> Self { + Self { d, p: Vec::new() } + } + pub fn push_slice(&mut self, slice: &[f32]) { + assert_eq!(slice.len(), self.d); + self.p.extend_from_slice(slice); + } + pub fn push_iter(&mut self, iter: impl ExactSizeIterator) { + assert_eq!(iter.len(), self.d); + self.p.extend(iter); + } + pub fn copy_within>(&mut self, src: R, dest: usize) { + let src_start = src.start_bound().map(|x| self.d * x); + let src_end = src.end_bound().map(|x| self.d * x); + self.p.copy_within((src_start, src_end), self.d * dest); + } + pub fn from_zeros(d: usize, p: usize) -> Self { + Self { + d, + p: vec![0.0; d * p], + } + } + pub fn truncate(&mut self, len: usize) { + self.p.truncate(self.d * len); + } + #[inline] + pub fn as_mut_view(&mut self) -> SquareMut<'_> { + SquareMut { + d: self.d, + p: self.p.as_mut_slice(), + } + } +} + +impl std::ops::Index for Square { + type Output = [f32]; + + fn index(&self, index: usize) -> &Self::Output { + &self.p[self.d * index..][..self.d] + } +} + +impl std::ops::IndexMut for Square { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + &mut self.p[self.d * index..][..self.d] + } +} + +impl<'a> IntoIterator for &'a Square { + type Item = &'a [f32]; + + type IntoIter = std::slice::ChunksExact<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact(self.d) + } +} + +impl<'a> IntoIterator for &'a mut Square { + type Item = &'a mut [f32]; + + type IntoIter = std::slice::ChunksExactMut<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact_mut(self.d) + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a Square { + type Item = &'a [f32]; + + type Iter = rayon::slice::ChunksExact<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSlice::par_chunks_exact(self.p.as_slice(), self.d) + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a mut Square { + type Item = &'a mut [f32]; + + type Iter = rayon::slice::ChunksExactMut<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p.as_mut_slice(), self.d) + } +} + +#[derive(Debug)] +pub struct SquareMut<'a> { + d: usize, + p: &'a mut [f32], +} + +impl<'a> SquareMut<'a> { + #[inline] + pub fn d(&self) -> usize { + self.d + } + #[inline] + pub fn len(&self) -> usize { + self.p.len() / self.d + } + #[inline] + pub fn is_empty(&self) -> bool { + self.p.is_empty() + } + pub fn new(d: usize, p: &'a mut [f32]) -> Self { + Self { d, p } + } + #[inline] + pub fn iter_mut<'b>(&'b mut self) -> std::slice::ChunksExactMut<'b, f32> + where + 'a: 'b, + { + self.p.chunks_exact_mut(self.d) + } + #[inline] + pub fn par_iter_mut<'b>(&'b mut self) -> rayon::slice::ChunksExactMut<'b, f32> + where + 'a: 'b, + { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p, self.d) + } + #[inline] + pub fn row(&self, i: usize) -> &[f32] { + let d = self.d; + &self.p[i * d..(i + 1) * d] + } + #[inline] + pub fn row_mut(&mut self, i: usize) -> &mut [f32] { + let d = self.d; + &mut self.p[i * d..(i + 1) * d] + } + pub fn copy_within>(&mut self, src: R, dest: usize) { + let src_start = src.start_bound().map(|x| self.d * x); + let src_end = src.end_bound().map(|x| self.d * x); + self.p.copy_within((src_start, src_end), self.d * dest); + } + #[inline] + pub fn into_inner(self) -> (usize, &'a mut [f32]) { + (self.d, self.p) + } +} + +impl std::ops::Index for SquareMut<'_> { + type Output = [f32]; + + fn index(&self, index: usize) -> &Self::Output { + &self.p[self.d * index..][..self.d] + } +} + +impl std::ops::IndexMut for SquareMut<'_> { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + &mut self.p[self.d * index..][..self.d] + } +} + +impl<'a> IntoIterator for &'a mut SquareMut<'a> { + type Item = &'a mut [f32]; + + type IntoIter = std::slice::ChunksExactMut<'a, f32>; + + fn into_iter(self) -> Self::IntoIter { + self.p.chunks_exact_mut(self.d) + } +} + +impl<'a> rayon::prelude::IntoParallelIterator for &'a mut SquareMut<'a> { + type Item = &'a mut [f32]; + + type Iter = rayon::slice::ChunksExactMut<'a, f32>; + + fn into_par_iter(self) -> Self::Iter { + rayon::slice::ParallelSliceMut::par_chunks_exact_mut(self.p, self.d) + } +} diff --git a/crates/rabitq/Cargo.toml b/crates/rabitq/Cargo.toml new file mode 100644 index 00000000..72e65fc6 --- /dev/null +++ b/crates/rabitq/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "rabitq" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +simd = { path = "../simd" } + +zerocopy.workspace = true + +[build-dependencies] +rand.workspace = true +rand_chacha.workspace = true + +[lints] +workspace = true diff --git a/crates/rabitq/build.rs b/crates/rabitq/build.rs new file mode 100644 index 00000000..007f32ef --- /dev/null +++ b/crates/rabitq/build.rs @@ -0,0 +1,32 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::env::var; +use std::error::Error; +use std::path::PathBuf; + +fn main() -> Result<(), Box> { + use rand::{RngExt, SeedableRng}; + use rand_chacha::ChaCha12Rng; + let mut rng = ChaCha12Rng::from_seed([7; 32]); + let mut bits = (0..262144).map(|_| rng.random::()).collect::>(); + if var("CARGO_CFG_TARGET_ENDIAN")?.as_str() == "big" { + for chunk in bits.as_chunks_mut::<8>().0 { + chunk.reverse(); + } + } + let out_dir = var("OUT_DIR")?; + std::fs::write(PathBuf::from(out_dir).join("bits"), bits)?; + Ok(()) +} diff --git a/crates/rabitq/src/bit.rs b/crates/rabitq/src/bit.rs new file mode 100644 index 00000000..900d4528 --- /dev/null +++ b/crates/rabitq/src/bit.rs @@ -0,0 +1,504 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use binary::BinaryLut; +use block::BlockLut; +use simd::Floating; + +#[derive(Debug, Clone, Copy, Default)] +pub struct CodeMetadata { + pub dis_u_2: f32, + pub factor_cnt: f32, + pub factor_ip: f32, + pub factor_err: f32, +} + +impl CodeMetadata { + #[inline(always)] + pub fn into_tuple(self) -> (f32, f32, f32, f32) { + ( + self.dis_u_2, + self.factor_cnt, + self.factor_ip, + self.factor_err, + ) + } + #[inline(always)] + pub fn into_array(self) -> [f32; 4] { + [ + self.dis_u_2, + self.factor_cnt, + self.factor_ip, + self.factor_err, + ] + } + #[inline(always)] + pub fn from_tuple((dis_u_2, factor_cnt, factor_ip, factor_err): (f32, f32, f32, f32)) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + } + } + #[inline(always)] + pub fn from_array([dis_u_2, factor_cnt, factor_ip, factor_err]: [f32; 4]) -> Self { + Self { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + } + } +} + +pub type Code = (CodeMetadata, Vec); + +pub fn code_metadata(vector: &[f32]) -> CodeMetadata { + let n = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.iter().filter(|x| x.is_sign_positive()).count(); + let cnt_neg = vector.iter().filter(|x| x.is_sign_negative()).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + } +} + +pub fn code_elements(vector: &[f32]) -> Vec { + let n = vector.len(); + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i].is_sign_positive()); + } + signs +} + +pub fn code(vector: &[f32]) -> Code { + let n = vector.len(); + let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); + let sum_of_x_2 = f32::reduce_sum_of_x2(vector); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.iter().filter(|x| x.is_sign_positive()).count(); + let cnt_neg = vector.iter().filter(|x| x.is_sign_negative()).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let mut signs = Vec::new(); + for i in 0..n { + signs.push(vector[i].is_sign_positive()); + } + signs + }, + ) +} + +pub fn preprocess(vector: &[f32]) -> (BlockLut, BinaryLut) { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + ( + block::preprocess_with_distance(vector, dis_v_2), + binary::preprocess_with_distance(vector, dis_v_2), + ) +} + +pub mod binary { + pub fn pack_code(input: &[bool]) -> Vec { + let f = |t: &[bool; 64]| { + let mut result = 0_u64; + for i in 0..64 { + result |= (t[i] as u64) << i; + } + result + }; + let (arrays, remainder) = input.as_chunks::<64>(); + let mut buffer = [false; 64]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } + + use super::CodeMetadata; + use simd::Floating; + + const BITS: usize = 6; + + #[derive(Debug, Clone, Copy)] + pub struct BinaryLutMetadata { + dis_v_2: f32, + b: f32, + k: f32, + qvector_sum: f32, + } + + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u64]); + + #[inline(always)] + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + preprocess_with_distance(vector, dis_v_2) + } + + pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BinaryLut { + let (k, b, qvector) = simd::quantize::quantize(vector, ((1 << BITS) - 1) as f32); + let qvector_sum = simd::byte::reduce_sum_of_x(&qvector) as f32; + ( + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }, + binarize(&qvector), + ) + } + + #[inline(always)] + pub fn accumulate(x: &[u64], y: &[impl AsRef<[u64]>; BITS]) -> u32 { + let mut result = 0_u32; + for i in 0..BITS { + result += simd::bit::reduce_sum_of_and(x, y[i].as_ref()) << i; + } + result + } + + #[inline(always)] + pub fn half_process_l2s( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_l2s_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2: _, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt, + factor_ip, + factor_err, + }: CodeMetadata, + BinaryLutMetadata { + dis_v_2, + b, + k, + qvector_sum, + }: BinaryLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + let e = k * ((2.0 * sum as f32) - qvector_sum) + b * factor_cnt; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) + } + + #[inline(always)] + pub(crate) fn binarize(vector: &[u8]) -> [Vec; BITS] { + let n = vector.len(); + let mut t: [_; BITS] = std::array::from_fn(|_| vec![0_u64; n.div_ceil(64)]); + for i in 0..BITS { + for j in 0..n { + let bit = (vector[j] >> i) & 1; + t[i][j / 64] |= (bit as u64) << (j % 64); + } + } + t + } +} + +pub mod block { + use super::CodeMetadata; + use simd::Floating; + + const BITS: usize = 8; + pub const STEP: usize = 65535_usize / ((1_usize << BITS) - 1); + + #[derive(Debug, Clone, Copy)] + pub struct BlockLutMetadata { + dis_v_2: f32, + k: f32, + c: f32, + } + + pub type BlockLut = (BlockLutMetadata, Vec<[u8; 16]>); + pub type BlockCode<'a> = ( + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [f32; 32], + &'a [[u8; 16]], + ); + + #[inline(always)] + pub fn preprocess(vector: &[f32]) -> BlockLut { + let dis_v_2 = f32::reduce_sum_of_x2(vector); + preprocess_with_distance(vector, dis_v_2) + } + + pub(crate) fn preprocess_with_distance(vector: &[f32], dis_v_2: f32) -> BlockLut { + let cvector = compress(vector); + let (k, b, cqvector) = + simd::quantize::quantize(cvector.as_flattened(), ((1 << BITS) - 1) as f32); + ( + BlockLutMetadata { + dis_v_2, + k, + c: b * vector.len().div_ceil(4) as f32, + }, + cqvector.as_chunks::<16>().0.to_vec(), + ) + } + + #[inline(always)] + pub fn accumulate(t: &[[u8; 16]], s: &[[u8; 16]]) -> [u32; 32] { + use std::iter::zip; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + sum + } + + #[inline(always)] + pub fn full_process_l2s( + (dis_u_2, _, factor_ip, factor_err, t): BlockCode<'_>, + lut: &BlockLut, + ) -> [(f32, f32); 32] { + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + std::array::from_fn(|i| { + let e = k * (sum[i] as f32) + c; + let rough = dis_u_2[i] + dis_v_2 - 2.0 * e * factor_ip[i]; + let err = 2.0 * factor_err[i] * dis_v_2.sqrt(); + (rough, err) + }) + } + + #[inline(always)] + pub fn full_process_dot( + (_, _, factor_ip, factor_err, t): BlockCode<'_>, + lut: &BlockLut, + ) -> [(f32, f32); 32] { + use std::iter::zip; + let &(BlockLutMetadata { dis_v_2, k, c }, ref s) = lut; + let mut sum = [0_u32; 32]; + for (t, s) in zip(t.chunks(STEP), s.chunks(STEP)) { + let delta = simd::fast_scan::scan(t, s); + simd::fast_scan::accu(&mut sum, &delta); + } + std::array::from_fn(|i| { + let e = k * (sum[i] as f32) + c; + let rough = -e * factor_ip[i]; + let err = factor_err[i] * dis_v_2.sqrt(); + (rough, err) + }) + } + + #[inline(always)] + pub fn half_process_l2s( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_v_2 - 2.0 * e * factor_ip; + let err = 2.0 * factor_err * dis_v_2.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_l2s_residual( + sum: u32, + CodeMetadata { + dis_u_2, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2: _, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = dis_u_2 + dis_f - 2.0 * e * factor_ip + delta; + let err = 2.0 * factor_err * dis_f.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_dot( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip; + let err = factor_err * dis_v_2.sqrt(); + (rough, err) + } + + #[inline(always)] + pub fn half_process_dot_residual( + sum: u32, + CodeMetadata { + dis_u_2: _, + factor_cnt: _, + factor_ip, + factor_err, + }: CodeMetadata, + BlockLutMetadata { dis_v_2, k, c }: BlockLutMetadata, + dis_f: f32, + delta: f32, + norm: f32, + ) -> (f32, f32) { + let e = k * (sum as f32) + c; + let rough = -e * factor_ip + dis_f + delta; + let err = factor_err * (dis_v_2 + norm * norm + 2.0 * dis_f).sqrt(); + (rough, err) + } + + fn compress(vector: &[f32]) -> Vec<[f32; 16]> { + let f = |&[t_0, t_1, t_2, t_3]: &[f32; 4]| { + [ + 0.0 - t_3 - t_2 - t_1 - t_0, + 0.0 - t_3 - t_2 - t_1 + t_0, + 0.0 - t_3 - t_2 + t_1 - t_0, + 0.0 - t_3 - t_2 + t_1 + t_0, + 0.0 - t_3 + t_2 - t_1 - t_0, + 0.0 - t_3 + t_2 - t_1 + t_0, + 0.0 - t_3 + t_2 + t_1 - t_0, + 0.0 - t_3 + t_2 + t_1 + t_0, + 0.0 + t_3 - t_2 - t_1 - t_0, + 0.0 + t_3 - t_2 - t_1 + t_0, + 0.0 + t_3 - t_2 + t_1 - t_0, + 0.0 + t_3 - t_2 + t_1 + t_0, + 0.0 + t_3 + t_2 - t_1 - t_0, + 0.0 + t_3 + t_2 - t_1 + t_0, + 0.0 + t_3 + t_2 + t_1 - t_0, + 0.0 + t_3 + t_2 + t_1 + t_0, + ] + }; + let (arrays, remainder) = vector.as_chunks::<4>(); + let mut buffer = [0.0f32; 4]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() + } +} diff --git a/crates/rabitq/src/bits.rs b/crates/rabitq/src/bits.rs new file mode 100644 index 00000000..a1853916 --- /dev/null +++ b/crates/rabitq/src/bits.rs @@ -0,0 +1,117 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +#[derive(Debug, Clone, Copy)] +#[repr(u8)] +pub enum Bits { + _1 = 1, + _2 = 2, +} + +impl TryFrom for Bits { + type Error = (); + + fn try_from(value: u8) -> Result { + match value { + 1 => Ok(Self::_1), + 2 => Ok(Self::_2), + _ => Err(()), + } + } +} + +pub fn code(bits: Bits, vector: &[f32]) -> Code { + match bits { + Bits::_1 => crate::extended::code::<1>(vector), + Bits::_2 => crate::extended::code::<2>(vector), + } +} + +pub fn ugly_code(bits: Bits, vector: &[f32]) -> Code { + match bits { + Bits::_1 => crate::extended::ugly_code::<1>(vector), + Bits::_2 => crate::extended::ugly_code::<2>(vector), + } +} + +pub fn pack_code(bits: Bits, input: &[u8]) -> Vec { + match bits { + Bits::_1 => crate::extended::pack_code::<1>(input) + .into_iter() + .flatten() + .collect(), + Bits::_2 => crate::extended::pack_code::<2>(input) + .into_iter() + .flatten() + .collect(), + } +} + +pub mod binary { + use crate::bits::Bits; + use crate::extended::CodeMetadata; + + const BITS: usize = 4; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, [Vec; BITS]); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, crate::extended::pack_code::(&elements)) + } + + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, crate::extended::pack_code::(&elements)) + } + + pub fn accumulate(bits: Bits, lhs: &[u64], rhs: &[Vec; BITS]) -> u32 { + match bits { + Bits::_1 => crate::extended::accumulate::<1, BITS>(lhs, rhs), + Bits::_2 => crate::extended::accumulate::<2, BITS>(lhs, rhs), + } + } + + pub fn half_process_dot( + bits: Bits, + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = match bits { + Bits::_1 => crate::extended::half_process_dot::<1, BITS>(dim, sum, code, lut), + Bits::_2 => crate::extended::half_process_dot::<2, BITS>(dim, sum, code, lut), + }; + (rough,) + } + + pub fn half_process_l2s( + bits: Bits, + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = match bits { + Bits::_1 => crate::extended::half_process_l2s::<1, BITS>(dim, sum, code, lut), + Bits::_2 => crate::extended::half_process_l2s::<2, BITS>(dim, sum, code, lut), + }; + (rough,) + } +} diff --git a/crates/rabitq/src/byte.rs b/crates/rabitq/src/byte.rs new file mode 100644 index 00000000..71681c71 --- /dev/null +++ b/crates/rabitq/src/byte.rs @@ -0,0 +1,83 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +#[deprecated] +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<8>(vector) +} + +pub fn ugly_code(vector: &[f32]) -> Code { + crate::extended::ugly_code::<8>(vector) +} + +pub fn pack_code(input: &[u8]) -> Vec { + input.to_vec() +} + +pub mod binary { + use crate::extended::CodeMetadata; + + const BITS: usize = 8; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + #[deprecated] + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::byte::reduce_sum_of_xy(x, y) + } + + pub fn half_process_dot( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<8, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_l2s( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2s::<8, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_cos( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_cos::<8, BITS>(dim, sum, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/extended.rs b/crates/rabitq/src/extended.rs new file mode 100644 index 00000000..789211f0 --- /dev/null +++ b/crates/rabitq/src/extended.rs @@ -0,0 +1,289 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use simd::Floating; + +#[derive(Debug, Clone, Copy)] +pub struct CodeMetadata { + pub dis_u_2: f32, + pub norm_of_lattice: f32, + pub sum_of_code: f32, +} + +impl CodeMetadata { + pub fn into_tuple(self) -> (f32, f32, f32) { + (self.dis_u_2, self.norm_of_lattice, self.sum_of_code) + } + pub fn into_array(self) -> [f32; 3] { + [self.dis_u_2, self.norm_of_lattice, self.sum_of_code] + } + pub fn from_tuple((dis_u_2, norm_of_lattice, sum_of_code): (f32, f32, f32)) -> Self { + Self { + dis_u_2, + norm_of_lattice, + sum_of_code, + } + } + pub fn from_array([dis_u_2, norm_of_lattice, sum_of_code]: [f32; 3]) -> Self { + Self { + dis_u_2, + norm_of_lattice, + sum_of_code, + } + } +} + +pub type Code = (CodeMetadata, Vec); + +pub fn code(vector: &[f32]) -> Code { + assert!((1..=8).contains(&BITS)); + + let n = vector.len(); + let dis_u_2 = f32::reduce_sum_of_x2(vector); + let code = { + let min = -(1 << (BITS - 1)); + let max = (1 << (BITS - 1)) - 1; + let normalized_vector = { + let mut vector = vector.to_vec(); + f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt()); + vector + }; + let scale = { + let mut o = normalized_vector.clone(); + f32::vector_abs_inplace(&mut o); + find_scale::(&o) + }; + let mut code = Vec::with_capacity(n as _); + for i in 0..n { + let v = scale * normalized_vector[i]; + let c = v.floor().clamp(min as f32, max as f32) as i32; + code.push((c + (1 << (BITS - 1))) as _); + } + code + }; + let norm_of_lattice = { + let base = -0.5 * ((1 << BITS) - 1) as f32; + let mut y = 0.0; + for i in 0..n { + let x = base + code[i] as f32; + y += x * x; + } + y.sqrt() + }; + let sum_of_code = { + let mut y = 0; + for i in 0..n { + let x = code[i] as u32; + y += x; + } + y as f32 + }; + ( + CodeMetadata { + dis_u_2, + norm_of_lattice, + sum_of_code, + }, + code, + ) +} + +pub fn ugly_code(vector: &[f32]) -> Code { + assert!((1..=8).contains(&BITS)); + + let n = vector.len(); + let dis_u_2 = f32::reduce_sum_of_x2(vector); + let code = { + let min = -(1 << (BITS - 1)); + let max = (1 << (BITS - 1)) - 1; + let normalized_vector = { + let mut vector = vector.to_vec(); + f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt()); + vector + }; + let scale = { + let mut o = normalized_vector.clone(); + f32::vector_abs_inplace(&mut o); + ugly_find_scale::(&o) + }; + let mut code = Vec::with_capacity(n as _); + for i in 0..n { + let v = scale * normalized_vector[i]; + let c = v.floor().clamp(min as f32, max as f32) as i32; + code.push((c + (1 << (BITS - 1))) as _); + } + code + }; + let norm_of_lattice = { + let base = -0.5 * ((1 << BITS) - 1) as f32; + let mut y = 0.0; + for i in 0..n { + let x = base + code[i] as f32; + y += x * x; + } + y.sqrt() + }; + let sum_of_code = { + let mut y = 0; + for i in 0..n { + let x = code[i] as u32; + y += x; + } + y as f32 + }; + ( + CodeMetadata { + dis_u_2, + norm_of_lattice, + sum_of_code, + }, + code, + ) +} + +pub fn half_process_l2s( + dim: u32, + sum: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; + lhs.dis_u_2 + rhs.dis_u_2 + - 2.0 + * ip + * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) + * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) +} + +pub fn half_process_dot( + dim: u32, + sum: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; + -ip * (lhs.dis_u_2.sqrt() / lhs.norm_of_lattice) * (rhs.dis_u_2.sqrt() / rhs.norm_of_lattice) +} + +pub fn half_process_cos( + dim: u32, + sum: u32, + lhs: CodeMetadata, + rhs: CodeMetadata, +) -> f32 { + assert!((1..=8).contains(&X)); + assert!((1..=8).contains(&Y)); + + let c_x = ((1 << X) - 1) as f32 * 0.5; + let c_y = ((1 << Y) - 1) as f32 * 0.5; + + let ip = sum as f32 - (c_y * lhs.sum_of_code + c_x * rhs.sum_of_code) + dim as f32 * c_x * c_y; + -ip / lhs.norm_of_lattice / rhs.norm_of_lattice +} + +fn find_scale(o: &[f32]) -> f32 { + assert!((1..=8).contains(&B)); + + let mask = (1_u32 << (B - 1)) - 1; + let dim = o.len(); + + let mut code = Vec::::with_capacity(dim); + let mut numerator = 0.0f64; + let mut sqr_denominator = 0.0f64; + + let (mut y_m, mut x_m); + + for i in 0..dim { + code.push(0); + let value = 0.5; + numerator += value * o[i] as f64; + sqr_denominator += value * value; + } + { + let x = 0.0; + let y = numerator / sqr_denominator.sqrt(); + (y_m, x_m) = (y, x); + } + + let mut events = Vec::<(f64, usize)>::new(); + for i in 0..dim { + for c in 1..=mask { + let x = (c as f64) / o[i] as f64; + events.push((x, i)); + } + } + events.sort_unstable_by(|(x, _), (y, _)| f64::total_cmp(x, y)); + for (x, i) in events.into_iter() { + code[i] += 1; + numerator += o[i] as f64; + sqr_denominator += 2.0 * code[i] as f64; + + let y = numerator / sqr_denominator.sqrt(); + if y > y_m { + (y_m, x_m) = (y, x); + } + } + + x_m as f32 + f32::EPSILON +} + +fn ugly_find_scale(o: &[f32]) -> f32 { + assert!((1..=8).contains(&B)); + + (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1 +} + +pub fn pack_code(input: &[u8]) -> [Vec; BITS] { + #[inline(always)] + fn f(array: &[u8; 64], bit: usize) -> u64 { + let mut result = 0_u64; + for i in 0..64 { + result |= ((array[i] as u64 >> bit) & 1) << i; + } + result + } + let (arrays, remainder) = input.as_chunks::<64>(); + let mut buffer = [0_u8; 64]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + std::array::from_fn(|bit| arrays.iter().chain(tailing).map(|t| f(t, bit)).collect()) +} + +pub fn accumulate(lhs: &[u64], rhs: &[Vec; Y]) -> u32 { + assert!(lhs.len().is_multiple_of(X)); + let d = lhs.len() / X; + let mut result = 0_u32; + for i in 0..X { + for j in 0..Y { + result += simd::bit::reduce_sum_of_and(&lhs[i * d..][..d], &rhs[j]) << (i + j); + } + } + result +} diff --git a/crates/rabitq/src/halfbyte.rs b/crates/rabitq/src/halfbyte.rs new file mode 100644 index 00000000..725b7d4f --- /dev/null +++ b/crates/rabitq/src/halfbyte.rs @@ -0,0 +1,92 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub use crate::extended::{Code, CodeMetadata}; + +#[deprecated] +pub fn code(vector: &[f32]) -> Code { + crate::extended::code::<4>(vector) +} + +pub fn ugly_code(vector: &[f32]) -> Code { + crate::extended::ugly_code::<4>(vector) +} + +pub fn pack_code(input: &[u8]) -> Vec { + let f = |t: &[u8; 2]| t[0] | t[1] << 4; + let (arrays, remainder) = input.as_chunks::<2>(); + let mut buffer = [0u8; 2]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() +} + +pub mod binary { + use crate::extended::CodeMetadata; + + const BITS: usize = 4; + + pub type BinaryLutMetadata = CodeMetadata; + pub type BinaryLut = (BinaryLutMetadata, Vec); + pub type BinaryCode<'a> = ((f32, f32, f32, f32), &'a [u8]); + + #[deprecated] + pub fn preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn ugly_preprocess(vector: &[f32]) -> BinaryLut { + let (metadata, elements) = crate::extended::ugly_code::(vector); + (metadata, super::pack_code(&elements)) + } + + pub fn accumulate(x: &[u8], y: &[u8]) -> u32 { + simd::halfbyte::reduce_sum_of_xy(x, y) + } + + pub fn half_process_dot( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_dot::<4, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_l2s( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_l2s::<4, BITS>(dim, sum, code, lut); + (rough,) + } + + pub fn half_process_cos( + dim: u32, + sum: u32, + code: CodeMetadata, + lut: BinaryLutMetadata, + ) -> (f32,) { + let rough = crate::extended::half_process_cos::<4, BITS>(dim, sum, code, lut); + (rough,) + } +} diff --git a/crates/rabitq/src/lib.rs b/crates/rabitq/src/lib.rs new file mode 100644 index 00000000..cc412b7f --- /dev/null +++ b/crates/rabitq/src/lib.rs @@ -0,0 +1,21 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod bit; +pub mod bits; +pub mod byte; +mod extended; +pub mod halfbyte; +pub mod packing; +pub mod rotate; diff --git a/crates/rabitq/src/packing.rs b/crates/rabitq/src/packing.rs new file mode 100644 index 00000000..ce21115d --- /dev/null +++ b/crates/rabitq/src/packing.rs @@ -0,0 +1,115 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub fn pack(x: [&[u8]; 32]) -> Vec<[u8; 16]> { + let n = { + let l = x.each_ref().map(|i| i.len()); + for i in 1..32 { + assert!(l[0] == l[i]); + } + l[0] + }; + let mut result = Vec::with_capacity(n); + for i in 0..n { + result.push([ + x[0][i] | (x[16][i] << 4), + x[8][i] | (x[24][i] << 4), + x[1][i] | (x[17][i] << 4), + x[9][i] | (x[25][i] << 4), + x[2][i] | (x[18][i] << 4), + x[10][i] | (x[26][i] << 4), + x[3][i] | (x[19][i] << 4), + x[11][i] | (x[27][i] << 4), + x[4][i] | (x[20][i] << 4), + x[12][i] | (x[28][i] << 4), + x[5][i] | (x[21][i] << 4), + x[13][i] | (x[29][i] << 4), + x[6][i] | (x[22][i] << 4), + x[14][i] | (x[30][i] << 4), + x[7][i] | (x[23][i] << 4), + x[15][i] | (x[31][i] << 4), + ]); + } + result +} + +pub fn unpack(x: &[[u8; 16]]) -> [Vec; 32] { + let n = x.len(); + let mut result = std::array::from_fn(|_| Vec::with_capacity(n)); + for i in 0..n { + result[0].push(x[i][0] & 0xf); + result[1].push(x[i][2] & 0xf); + result[2].push(x[i][4] & 0xf); + result[3].push(x[i][6] & 0xf); + result[4].push(x[i][8] & 0xf); + result[5].push(x[i][10] & 0xf); + result[6].push(x[i][12] & 0xf); + result[7].push(x[i][14] & 0xf); + result[8].push(x[i][1] & 0xf); + result[9].push(x[i][3] & 0xf); + result[10].push(x[i][5] & 0xf); + result[11].push(x[i][7] & 0xf); + result[12].push(x[i][9] & 0xf); + result[13].push(x[i][11] & 0xf); + result[14].push(x[i][13] & 0xf); + result[15].push(x[i][15] & 0xf); + result[16].push(x[i][0] >> 4); + result[17].push(x[i][2] >> 4); + result[18].push(x[i][4] >> 4); + result[19].push(x[i][6] >> 4); + result[20].push(x[i][8] >> 4); + result[21].push(x[i][10] >> 4); + result[22].push(x[i][12] >> 4); + result[23].push(x[i][14] >> 4); + result[24].push(x[i][1] >> 4); + result[25].push(x[i][3] >> 4); + result[26].push(x[i][5] >> 4); + result[27].push(x[i][7] >> 4); + result[28].push(x[i][9] >> 4); + result[29].push(x[i][11] >> 4); + result[30].push(x[i][13] >> 4); + result[31].push(x[i][15] >> 4); + } + result +} + +pub fn padding_pack(x: impl IntoIterator>) -> Vec<[u8; 16]> { + let x = x.into_iter().collect::>(); + let x = x.iter().map(|x| x.as_ref()).collect::>(); + if x.is_empty() || x.len() > 32 { + panic!("too few or too many slices"); + } + let n = x[0].len(); + let t = vec![0; n]; + pack(std::array::from_fn(|i| { + if i < x.len() { x[i] } else { t.as_slice() } + })) +} + +pub fn any_pack(mut x: impl Iterator) -> [T; 32] { + std::array::from_fn(|_| x.next()).map(|x| x.unwrap_or_default()) +} + +pub fn pack_to_u4(input: &[bool]) -> Vec { + let f = |t: &[bool; 4]| t[0] as u8 | (t[1] as u8) << 1 | (t[2] as u8) << 2 | (t[3] as u8) << 3; + let (arrays, remainder) = input.as_chunks::<4>(); + let mut buffer = [false; 4]; + let tailing = if !remainder.is_empty() { + buffer[..remainder.len()].copy_from_slice(remainder); + Some(&buffer) + } else { + None + }; + arrays.iter().chain(tailing).map(f).collect() +} diff --git a/crates/rabitq/src/rotate.rs b/crates/rabitq/src/rotate.rs new file mode 100644 index 00000000..93806163 --- /dev/null +++ b/crates/rabitq/src/rotate.rs @@ -0,0 +1,151 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +const BITS: &[u8; 262144] = include_bytes!(concat!(env!("OUT_DIR"), "/bits")); + +#[cfg(target_endian = "little")] +const _: () = { + assert!(BITS[0] == 246); + assert!(BITS[1] == 133); + assert!(BITS[2] == 163); + assert!(BITS[3] == 106); + assert!(BITS[4] == 54); + assert!(BITS[5] == 126); + assert!(BITS[6] == 9); + assert!(BITS[7] == 115); +}; + +#[cfg(target_endian = "big")] +const _: () = { + assert!(BITS[7] == 246); + assert!(BITS[6] == 133); + assert!(BITS[5] == 163); + assert!(BITS[4] == 106); + assert!(BITS[3] == 54); + assert!(BITS[2] == 126); + assert!(BITS[1] == 9); + assert!(BITS[0] == 115); +}; + +static BITS_0: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[0]); +static BITS_1: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[1]); +static BITS_2: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[2]); +static BITS_3: [u64; 1024] = zerocopy::transmute!(BITS.as_chunks::<8192>().0[3]); + +fn kacs_walk(result: &mut [f32]) { + let n = result.len(); + let m = n / 2; + let (l, t) = result.split_at_mut(m); + let (_, r) = t.split_at_mut(n - 2 * m); + simd::rotate::givens(l, r); +} + +pub fn rotate(vector: &[f32]) -> Vec { + let mut vector = vector.to_vec(); + rotate_inplace(&mut vector); + vector +} + +pub fn rotate_inplace(result: &mut [f32]) { + use simd::Floating; + use std::ops::Bound::{Excluded, Included, Unbounded}; + + let n = result.len(); + let base = n.ilog2(); + let scale = 1.0 / ((1_usize << base) as f32).sqrt(); + + let l = (Unbounded, Excluded(1_usize << base)); + let r = (Included(n - (1_usize << base)), Unbounded); + + simd::rotate::flip(&BITS_0, result); + simd::fht::fht(&mut result[l]); + f32::vector_mul_scalar_inplace(&mut result[l], scale); + if n != (1_usize << base) { + kacs_walk(result); + } + + simd::rotate::flip(&BITS_1, result); + simd::fht::fht(&mut result[r]); + f32::vector_mul_scalar_inplace(&mut result[r], scale); + if n != (1_usize << base) { + kacs_walk(result); + } + + simd::rotate::flip(&BITS_2, result); + simd::fht::fht(&mut result[l]); + f32::vector_mul_scalar_inplace(&mut result[l], scale); + if n != (1_usize << base) { + kacs_walk(result); + } + + simd::rotate::flip(&BITS_3, result); + simd::fht::fht(&mut result[r]); + f32::vector_mul_scalar_inplace(&mut result[r], scale); + if n != (1_usize << base) { + kacs_walk(result); + } +} + +pub fn rotate_reversed_inplace(result: &mut [f32]) { + use simd::Floating; + use std::ops::Bound::{Excluded, Included, Unbounded}; + + let n = result.len(); + let base = n.ilog2(); + let scale = 1.0 / ((1_usize << base) as f32).sqrt(); + + let l = (Unbounded, Excluded(1_usize << base)); + let r = (Included(n - (1_usize << base)), Unbounded); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[r], scale); + simd::fht::fht(&mut result[r]); + simd::rotate::flip(&BITS_3, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[l], scale); + simd::fht::fht(&mut result[l]); + simd::rotate::flip(&BITS_2, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[r], scale); + simd::fht::fht(&mut result[r]); + simd::rotate::flip(&BITS_1, result); + + if n != (1_usize << base) { + kacs_walk(result); + } + f32::vector_mul_scalar_inplace(&mut result[l], scale); + simd::fht::fht(&mut result[l]); + simd::rotate::flip(&BITS_0, result); +} + +#[test] +fn reverse() { + let mut x = vec![2.0, 3.0, 4.0]; + rotate_inplace(&mut x); + assert!((x[0] - 3.981917).abs() < 1e-6); + assert!((x[1] - 1.8043789).abs() < 1e-6); + assert!((x[2] - 3.1446066).abs() < 1e-6); + rotate_reversed_inplace(&mut x); + assert!((x[0] - 2.0).abs() < 1e-6); + assert!((x[1] - 3.0).abs() < 1e-6); + assert!((x[2] - 4.0).abs() < 1e-6); +} diff --git a/crates/simd/Cargo.toml b/crates/simd/Cargo.toml new file mode 100644 index 00000000..1b6faf59 --- /dev/null +++ b/crates/simd/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "simd" +version.workspace = true +edition.workspace = true +publish = false + +[[bench]] +name = "bench" +harness = false +required-features = ["internal"] + +[features] +init = [] +internal = [] +nightly_f16 = ["zerocopy/float-nightly"] + +[dependencies] +simd_macros = { path = "../simd_macros" } + +half = { version = "2.7.1", features = ["zerocopy"] } +seq-macro.workspace = true +zerocopy.workspace = true + +[dev-dependencies] +criterion = "0.8.2" +rand.workspace = true + +[build-dependencies] +cc = "1.2.60" + +[lints] +workspace = true diff --git a/crates/simd/benches/bench.rs b/crates/simd/benches/bench.rs new file mode 100644 index 00000000..db9ba0df --- /dev/null +++ b/crates/simd/benches/bench.rs @@ -0,0 +1,337 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#![allow(unsafe_code)] + +use criterion::{Criterion, criterion_group, criterion_main}; + +fn floating_f32_reduce_sum_of_xy(c: &mut Criterion) { + use rand::RngExt; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + let y = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f32::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f32::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") && simd::is_feature_detected!("fma") { + c.bench_function("floating_f32::reduce_sum_of_xy::v2_fma", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_v2_fma(&x, &y) + }) + }); + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if simd::is_cpu_detected!("a3.256") { + c.bench_function("floating_f32::reduce_sum_of_xy::a3_256", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_a3_256(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("floating_f32::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } +} + +fn floating_f32_reduce_sum_of_d2(c: &mut Criterion) { + use rand::RngExt; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + let y = (0..4095) + .map(|_| rng.random_range(-1.0..=1.0f32)) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f32::reduce_sum_of_d2::v4", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f32::reduce_sum_of_d2::v3", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") && simd::is_feature_detected!("fma") { + c.bench_function("floating_f32::reduce_sum_of_d2::v2_fma", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_v2_fma(&x, &y) + }) + }); + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if simd::is_cpu_detected!("a3.256") { + c.bench_function("floating_f32::reduce_sum_of_d2::a3_256", |b| { + b.iter(|| unsafe { + simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_a3_256(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("floating_f32::reduce_sum_of_d2::a2", |b| { + b.iter(|| unsafe { simd::floating_f32::reduce_sum_of_d2::reduce_sum_of_d2_a2(&x, &y) }) + }); + } +} + +fn floating_f16_reduce_sum_of_xy(c: &mut Criterion) { + use rand::RngExt; + use simd::{F16, f16}; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let y = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512fp16") { + c.bench_function("floating_f16::reduce_sum_of_xy::v4_avx512fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512fp16(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f16::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f16::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a3.512") { + c.bench_function("floating_f16::reduce_sum_of_xy::a3_512", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_a3_512(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("fp16") { + c.bench_function("floating_f16::reduce_sum_of_xy::a2_fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_xy::reduce_sum_of_xy_a2_fp16(&x, &y) + }) + }); + } +} + +fn floating_f16_reduce_sum_of_d2(c: &mut Criterion) { + use rand::RngExt; + use simd::{F16, f16}; + let mut rng = rand::rng(); + let x = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let y = (0..4095) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512fp16") { + c.bench_function("floating_f16::reduce_sum_of_d2::v4_avx512fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v4_avx512fp16(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("floating_f16::reduce_sum_of_d2::v4", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("floating_f16::reduce_sum_of_d2::v3", |b| { + b.iter(|| unsafe { simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a3.512") { + c.bench_function("floating_f16::reduce_sum_of_d2::a3_512", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_a3_512(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("fp16") { + c.bench_function("floating_f16::reduce_sum_of_d2::a2_fp16", |b| { + b.iter(|| unsafe { + simd::floating_f16::reduce_sum_of_d2::reduce_sum_of_d2_a2_fp16(&x, &y) + }) + }); + } +} + +fn byte_reduce_sum_of_xy(c: &mut Criterion) { + use rand::RngExt; + let mut rng = rand::rng(); + let x = (0..4095).map(|_| rng.random::()).collect::>(); + let y = (0..4095).map(|_| rng.random::()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { + c.bench_function("byte::reduce_sum_of_xy::v4_avx512vnni", |b| { + b.iter(|| unsafe { + simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512vnni(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("byte::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("byte::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("byte::reduce_sum_of_xy::v2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("dotprod") { + c.bench_function("byte::reduce_sum_of_xy::a2_prod", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_a2_dotprod(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("byte::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } +} + +fn byte_reduce_sum_of_x(c: &mut Criterion) { + use rand::RngExt; + let mut rng = rand::rng(); + let this = (0..4095).map(|_| rng.random::()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("byte::reduce_sum_of_x::v4", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v4(&this) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("byte::reduce_sum_of_x::v3", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v3(&this) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("byte::reduce_sum_of_x::v2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_v2(&this) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("byte::reduce_sum_of_x::a2", |b| { + b.iter(|| unsafe { simd::byte::reduce_sum_of_x::reduce_sum_of_x_a2(&this) }) + }); + } +} + +fn halfbyte_reduce_sum_of_xy(c: &mut Criterion) { + use rand::RngExt; + let mut rng = rand::rng(); + let x = (0..2047).map(|_| rng.random::()).collect::>(); + let y = (0..2047).map(|_| rng.random::()).collect::>(); + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") && simd::is_feature_detected!("avx512vnni") { + c.bench_function("halfbyte::reduce_sum_of_xy::v4_avx512vnni", |b| { + b.iter(|| unsafe { + simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v4_avx512vnni(&x, &y) + }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v4") { + c.bench_function("halfbyte::reduce_sum_of_xy::v4", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v4(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v3") { + c.bench_function("halfbyte::reduce_sum_of_xy::v3", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v3(&x, &y) }) + }); + } + #[cfg(target_arch = "x86_64")] + if simd::is_cpu_detected!("v2") { + c.bench_function("halfbyte::reduce_sum_of_xy::v2", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_v2(&x, &y) }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") && simd::is_feature_detected!("dotprod") { + c.bench_function("halfbyte::reduce_sum_of_xy::a2_prod", |b| { + b.iter(|| unsafe { + simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_a2_dotprod(&x, &y) + }) + }); + } + #[cfg(target_arch = "aarch64")] + if simd::is_cpu_detected!("a2") { + c.bench_function("halfbyte::reduce_sum_of_xy::a2", |b| { + b.iter(|| unsafe { simd::halfbyte::reduce_sum_of_xy::reduce_sum_of_xy_a2(&x, &y) }) + }); + } +} + +criterion_group!( + benches, + floating_f32_reduce_sum_of_xy, + floating_f32_reduce_sum_of_d2, + floating_f16_reduce_sum_of_xy, + floating_f16_reduce_sum_of_d2, + byte_reduce_sum_of_xy, + byte_reduce_sum_of_x, + halfbyte_reduce_sum_of_xy, +); +criterion_main!(benches); diff --git a/crates/simd/build.rs b/crates/simd/build.rs new file mode 100644 index 00000000..57a73c01 --- /dev/null +++ b/crates/simd/build.rs @@ -0,0 +1,38 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::env::var; +use std::error::Error; + +fn main() -> Result<(), Box> { + println!("cargo::rerun-if-changed=cshim"); + let target_arch = var("CARGO_CFG_TARGET_ARCH")?; + let target_endian = var("CARGO_CFG_TARGET_ENDIAN")?; + #[allow(clippy::single_match)] + match target_arch.as_str() { + "aarch64" => { + let mut build = cc::Build::new(); + build.file("./cshim/aarch64_byte.c"); + build.file("./cshim/aarch64_halfbyte.c"); + if target_endian == "little" { + build.file("./cshim/aarch64_sve_fp16.c"); + build.file("./cshim/aarch64_sve_fp32.c"); + } + build.opt_level(3); + build.compile("simd_cshim"); + } + _ => {} + } + Ok(()) +} diff --git a/crates/simd/cshim/aarch64_byte.c b/crates/simd/cshim/aarch64_byte.c new file mode 100644 index 00000000..a9519510 --- /dev/null +++ b/crates/simd/cshim/aarch64_byte.c @@ -0,0 +1,52 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#include +#include + +__attribute__((target("+dotprod"))) uint32_t +byte_reduce_sum_of_xy_a2_dotprod(size_t n, uint8_t *restrict a, + uint8_t *restrict b) { + uint32x4_t sum = vdupq_n_u32(0); + while (n >= 16) { + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + sum = vdotq_u32(sum, x, y); + n -= 16, a += 16, b += 16; + } + if (n > 0) { + uint8_t _a[16] = {}, _b[16] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + sum = vdotq_u32(sum, x, y); + } + return vaddvq_u32(sum); +} diff --git a/crates/simd/cshim/aarch64_halfbyte.c b/crates/simd/cshim/aarch64_halfbyte.c new file mode 100644 index 00000000..3b0faa07 --- /dev/null +++ b/crates/simd/cshim/aarch64_halfbyte.c @@ -0,0 +1,64 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#include +#include + +__attribute__((target("+dotprod"))) uint32_t +halfbyte_reduce_sum_of_xy_a2_dotprod(size_t n, uint8_t *restrict a, + uint8_t *restrict b) { + uint8x16_t lo = vdupq_n_u8(0x0f); + uint32x4_t _0 = vdupq_n_u32(0); + uint32x4_t _1 = vdupq_n_u32(0); + while (n >= 16) { + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + uint8x16_t x_0 = vandq_u8(x, lo); + uint8x16_t x_1 = vshrq_n_u8(x, 4); + uint8x16_t y_0 = vandq_u8(y, lo); + uint8x16_t y_1 = vshrq_n_u8(y, 4); + _0 = vdotq_u32(_0, x_0, y_0); + _1 = vdotq_u32(_1, x_1, y_1); + n -= 16, a += 16, b += 16; + } + if (n > 0) { + uint8_t _a[16] = {}, _b[16] = {}; + for (size_t i = 0; i < n; i += 1) { + _a[i] = a[i], _b[i] = b[i]; + } + a = _a, b = _b; + uint8x16_t x = vld1q_u8(a); + uint8x16_t y = vld1q_u8(b); + uint8x16_t x_0 = vandq_u8(x, lo); + uint8x16_t x_1 = vshrq_n_u8(x, 4); + uint8x16_t y_0 = vandq_u8(y, lo); + uint8x16_t y_1 = vshrq_n_u8(y, 4); + _0 = vdotq_u32(_0, x_0, y_0); + _1 = vdotq_u32(_1, x_1, y_1); + } + return vaddvq_u32(vaddq_u32(_0, _1)); +} diff --git a/crates/simd/cshim/aarch64_sve_fp16.c b/crates/simd/cshim/aarch64_sve_fp16.c new file mode 100644 index 00000000..064ba4be --- /dev/null +++ b/crates/simd/cshim/aarch64_sve_fp16.c @@ -0,0 +1,113 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#if defined(__AARCH64EL__) +#include +#endif +#include +#include + +typedef __fp16 f16; +typedef float f32; + +__attribute__((target("+sve"))) float +fp16_reduce_sum_of_xy_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { + const size_t vl = svcnth(); + svfloat16_t _0 = svdup_f16(0.0); + svfloat16_t _1 = svdup_f16(0.0); + while (n >= vl * 2) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t x_1 = svld1_f16(mask, a + 1 * vl); + svfloat16_t y_1 = svld1_f16(mask, b + 1 * vl); + _0 = svmla_f16_x(mask, _0, x_0, y_0); + _1 = svmla_f16_x(mask, _1, x_1, y_1); + n -= vl * 2, a += vl * 2, b += vl * 2; + } + if (n >= vl) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + _0 = svmla_f16_x(mask, _0, x_0, y_0); + n -= vl * 1, a += vl * 1, b += vl * 1; + } + if (n > 0) { + svbool_t mask = svwhilelt_b16((int64_t)0, (int64_t)n); + svfloat16_t x_0 = svld1_f16(mask, a); + svfloat16_t y_0 = svld1_f16(mask, b); + _0 = svmla_f16_m(mask, _0, x_0, y_0); + } + svbool_t mask = svptrue_b32(); + svfloat32_t s_0 = svcvt_f32_f16_x(mask, _0); + svfloat32_t s_1 = svcvt_f32_f16_x(mask, svext_f16(_0, _0, 1)); + svfloat32_t s_2 = svcvt_f32_f16_x(mask, _1); + svfloat32_t s_3 = svcvt_f32_f16_x(mask, svext_f16(_1, _1, 1)); + return svaddv_f32(mask, svadd_f32_x(mask, svadd_f32_x(mask, s_0, s_2), + svadd_f32_x(mask, s_1, s_3))); +} + +__attribute__((target("+sve"))) float +fp16_reduce_sum_of_d2_a3_512(size_t n, f16 *restrict a, f16 *restrict b) { + const size_t vl = svcnth(); + svfloat16_t _0 = svdup_f16(0.0); + svfloat16_t _1 = svdup_f16(0.0); + while (n >= vl * 2) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t x_1 = svld1_f16(mask, a + 1 * vl); + svfloat16_t y_1 = svld1_f16(mask, b + 1 * vl); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + svfloat16_t d_1 = svsub_f16_z(mask, x_1, y_1); + _0 = svmla_f16_x(mask, _0, d_0, d_0); + _1 = svmla_f16_x(mask, _1, d_1, d_1); + n -= vl * 2, a += vl * 2, b += vl * 2; + } + if (n >= vl) { + svbool_t mask = svptrue_b16(); + svfloat16_t x_0 = svld1_f16(mask, a + 0 * vl); + svfloat16_t y_0 = svld1_f16(mask, b + 0 * vl); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + _0 = svmla_f16_x(mask, _0, d_0, d_0); + n -= vl * 1, a += vl * 1, b += vl * 1; + } + if (n > 0) { + svbool_t mask = svwhilelt_b16((int64_t)0, (int64_t)n); + svfloat16_t x_0 = svld1_f16(mask, a); + svfloat16_t y_0 = svld1_f16(mask, b); + svfloat16_t d_0 = svsub_f16_z(mask, x_0, y_0); + _0 = svmla_f16_m(mask, _0, d_0, d_0); + } + svbool_t mask = svptrue_b32(); + svfloat32_t s_0 = svcvt_f32_f16_x(mask, _0); + svfloat32_t s_1 = svcvt_f32_f16_x(mask, svext_f16(_0, _0, 1)); + svfloat32_t s_2 = svcvt_f32_f16_x(mask, _1); + svfloat32_t s_3 = svcvt_f32_f16_x(mask, svext_f16(_1, _1, 1)); + return svaddv_f32(mask, svadd_f32_x(mask, svadd_f32_x(mask, s_0, s_2), + svadd_f32_x(mask, s_1, s_3))); +} diff --git a/crates/simd/cshim/aarch64_sve_fp32.c b/crates/simd/cshim/aarch64_sve_fp32.c new file mode 100644 index 00000000..4ceadeaf --- /dev/null +++ b/crates/simd/cshim/aarch64_sve_fp32.c @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#if defined(__clang__) +#if !(__clang_major__ >= 16) +#error "Clang version must be at least 16." +#endif +#elif defined(__GNUC__) +#if !(__GNUC__ >= 14) +#error "GCC version must be at least 14." +#endif +#else +#error "This file requires Clang or GCC." +#endif + +#include +#if defined(__AARCH64EL__) +#include +#endif +#include +#include + +typedef __fp16 f16; +typedef float f32; + +__attribute__((target("+sve"))) float +fp32_reduce_sum_of_xy_a3_256(size_t n, float *restrict lhs, + float *restrict rhs) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + sum = svmla_f32_m(mask, sum, x, y); + } + return svaddv_f32(svptrue_b32(), sum); +} + +__attribute__((target("+sve"))) float +fp32_reduce_sum_of_d2_a3_256(size_t n, float *restrict lhs, + float *restrict rhs) { + svfloat32_t sum = svdup_f32(0.0); + for (size_t i = 0; i < n; i += svcntw()) { + svbool_t mask = svwhilelt_b32((int64_t)i, (int64_t)n); + svfloat32_t x = svld1_f32(mask, lhs + i); + svfloat32_t y = svld1_f32(mask, rhs + i); + svfloat32_t d = svsub_f32_z(mask, x, y); + sum = svmla_f32_m(mask, sum, d, d); + } + return svaddv_f32(svptrue_b32(), sum); +} diff --git a/crates/simd/src/aligned.rs b/crates/simd/src/aligned.rs new file mode 100644 index 00000000..b06b1b4c --- /dev/null +++ b/crates/simd/src/aligned.rs @@ -0,0 +1,18 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[allow(dead_code)] +#[derive(Debug, Clone, Copy)] +#[repr(C, align(16))] +pub struct Aligned16(pub T); diff --git a/crates/simd/src/bit.rs b/crates/simd/src/bit.rs new file mode 100644 index 00000000..d0160b25 --- /dev/null +++ b/crates/simd/src/bit.rs @@ -0,0 +1,1006 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[inline(always)] +pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_and::reduce_sum_of_and(lhs, rhs) +} + +mod reduce_sum_of_and { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn reduce_sum_of_and_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + } + _mm512_reduce_add_epi64(and) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_and_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mut sum_and = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); + } + _mm512_reduce_add_epi64(sum_and) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_and_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mut sum_and = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); + } + emulate_mm256_reduce_add_epi64(sum_and) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_and(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut and = 0; + for i in 0..n { + and += (lhs[i] & rhs[i]).count_ones(); + } + and + } +} + +#[inline(always)] +pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_or::reduce_sum_of_or(lhs, rhs) +} + +mod reduce_sum_of_or { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn reduce_sum_of_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + _mm512_reduce_add_epi64(or) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_or_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_or_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mut sum_or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); + } + _mm512_reduce_add_epi64(sum_or) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_or_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_or_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_or_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mut sum_or = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); + } + emulate_mm256_reduce_add_epi64(sum_or) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_or_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_or_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_or(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut or = 0; + for i in 0..n { + or += (lhs[i] | rhs[i]).count_ones(); + } + or + } +} + +#[inline(always)] +pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + reduce_sum_of_xor::reduce_sum_of_xor(lhs, rhs) +} + +mod reduce_sum_of_xor { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn reduce_sum_of_xor_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut xor = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + xor = _mm512_add_epi64(xor, _mm512_popcnt_epi64(_mm512_xor_si512(x, y))); + } + _mm512_reduce_add_epi64(xor) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xor_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_xor_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xor_v4(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mut sum_xor = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + let xor = _mm512_xor_si512(x, y); + let xor_lo = _mm512_and_si512(xor, mask_0); + let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); + let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); + let xor_res = _mm512_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm512_sad_epu8(xor_res, _mm512_setzero_si512()); + sum_xor = _mm512_add_epi64(sum_xor, xor_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + let xor = _mm512_xor_si512(x, y); + let xor_lo = _mm512_and_si512(xor, mask_0); + let xor_hi = _mm512_and_si512(_mm512_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm512_shuffle_epi8(lut, xor_lo); + let xor_res_hi = _mm512_shuffle_epi8(lut, xor_hi); + let xor_res = _mm512_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm512_sad_epu8(xor_res, _mm512_setzero_si512()); + sum_xor = _mm512_add_epi64(sum_xor, xor_sad); + } + _mm512_reduce_add_epi64(sum_xor) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xor_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_xor_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xor_v3(lhs: &[u64], rhs: &[u64]) -> u32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mut sum_xor = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let xor = _mm256_xor_si256(x, y); + let xor_lo = _mm256_and_si256(xor, mask_0); + let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm256_shuffle_epi8(lut, xor_lo); + let xor_res_hi = _mm256_shuffle_epi8(lut, xor_hi); + let xor_res = _mm256_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm256_sad_epu8(xor_res, _mm256_setzero_si256()); + sum_xor = _mm256_add_epi64(sum_xor, xor_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let xor = _mm256_xor_si256(x, y); + let xor_lo = _mm256_and_si256(xor, mask_0); + let xor_hi = _mm256_and_si256(_mm256_srli_epi16(xor, 4), mask_0); + let xor_res_lo = _mm256_shuffle_epi8(lut, xor_lo); + let xor_res_hi = _mm256_shuffle_epi8(lut, xor_hi); + let xor_res = _mm256_add_epi8(xor_res_lo, xor_res_hi); + let xor_sad = _mm256_sad_epu8(xor_res, _mm256_setzero_si256()); + sum_xor = _mm256_add_epi64(sum_xor, xor_sad); + } + emulate_mm256_reduce_add_epi64(sum_xor) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xor_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_xor_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xor(lhs: &[u64], rhs: &[u64]) -> u32 { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut xor = 0; + for i in 0..n { + xor += (lhs[i] ^ rhs[i]).count_ones(); + } + xor + } +} + +#[inline(always)] +pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + reduce_sum_of_and_or::reduce_sum_of_and_or(lhs, rhs) +} + +mod reduce_sum_of_and_or { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn reduce_sum_of_and_or_v4_avx512vpopcntdq(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut and = _mm512_setzero_si512(); + let mut or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + and = _mm512_add_epi64(and, _mm512_popcnt_epi64(_mm512_and_si512(x, y))); + or = _mm512_add_epi64(or, _mm512_popcnt_epi64(_mm512_or_si512(x, y))); + } + ( + _mm512_reduce_add_epi64(and) as u32, + _mm512_reduce_add_epi64(or) as u32, + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_or_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_or_v4_avx512vpopcntdq(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_and_or_v4(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mut sum_and = _mm512_setzero_si512(); + let mut sum_or = _mm512_setzero_si512(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let y = unsafe { _mm512_loadu_si512(b.cast()) }; + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi64(mask, b.cast()) }; + let and = _mm512_and_si512(x, y); + let and_lo = _mm512_and_si512(and, mask_0); + let and_hi = _mm512_and_si512(_mm512_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm512_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm512_shuffle_epi8(lut, and_hi); + let and_res = _mm512_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm512_sad_epu8(and_res, _mm512_setzero_si512()); + sum_and = _mm512_add_epi64(sum_and, and_sad); + let or = _mm512_or_si512(x, y); + let or_lo = _mm512_and_si512(or, mask_0); + let or_hi = _mm512_and_si512(_mm512_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm512_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm512_shuffle_epi8(lut, or_hi); + let or_res = _mm512_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm512_sad_epu8(or_res, _mm512_setzero_si512()); + sum_or = _mm512_add_epi64(sum_or, or_sad); + } + ( + _mm512_reduce_add_epi64(sum_and) as u32, + _mm512_reduce_add_epi64(sum_or) as u32, + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_or_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_or_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_and_or_v3(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert!(lhs.len() == rhs.len()); + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mut sum_and = _mm256_setzero_si256(); + let mut sum_or = _mm256_setzero_si256(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut n = lhs.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let and = _mm256_and_si256(x, y); + let and_lo = _mm256_and_si256(and, mask_0); + let and_hi = _mm256_and_si256(_mm256_srli_epi16(and, 4), mask_0); + let and_res_lo = _mm256_shuffle_epi8(lut, and_lo); + let and_res_hi = _mm256_shuffle_epi8(lut, and_hi); + let and_res = _mm256_add_epi8(and_res_lo, and_res_hi); + let and_sad = _mm256_sad_epu8(and_res, _mm256_setzero_si256()); + sum_and = _mm256_add_epi64(sum_and, and_sad); + let or = _mm256_or_si256(x, y); + let or_lo = _mm256_and_si256(or, mask_0); + let or_hi = _mm256_and_si256(_mm256_srli_epi16(or, 4), mask_0); + let or_res_lo = _mm256_shuffle_epi8(lut, or_lo); + let or_res_hi = _mm256_shuffle_epi8(lut, or_hi); + let or_res = _mm256_add_epi8(or_res_lo, or_res_hi); + let or_sad = _mm256_sad_epu8(or_res, _mm256_setzero_si256()); + sum_or = _mm256_add_epi64(sum_or, or_sad); + } + ( + emulate_mm256_reduce_add_epi64(sum_and) as u32, + emulate_mm256_reduce_add_epi64(sum_or) as u32, + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_and_or_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lhs = (0..126).map(|_| rand::random::()).collect::>(); + let rhs = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_and_or_v3(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_and_or(lhs: &[u64], rhs: &[u64]) -> (u32, u32) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut and = 0; + let mut or = 0; + for i in 0..n { + and += (lhs[i] & rhs[i]).count_ones(); + or += (lhs[i] | rhs[i]).count_ones(); + } + (and, or) + } +} + +#[inline(always)] +pub fn reduce_sum_of_x(this: &[u64]) -> u32 { + reduce_sum_of_x::reduce_sum_of_x(this) +} + +mod reduce_sum_of_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vpopcntdq")] + fn reduce_sum_of_x_v4_avx512vpopcntdq(this: &[u64]) -> u32 { + use core::arch::x86_64::*; + let mut sum = _mm512_setzero_si512(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + sum = _mm512_add_epi64(sum, _mm512_popcnt_epi64(x)); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + sum = _mm512_add_epi64(sum, _mm512_popcnt_epi64(x)); + } + _mm512_reduce_add_epi64(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v4_avx512vpopcntdq_test() { + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vpopcntdq") { + println!("test {} ... skipped (v4:avx512vpopcntdq)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_x_v4_avx512vpopcntdq(&this) }; + let fallback = fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[u64]) -> u32 { + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 4] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 4]; + let lut = unsafe { _mm512_loadu_si512((&raw const LUT).cast()) }; + let mask_0 = _mm512_set1_epi8(0x0f); + let mut sum = _mm512_setzero_si512(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 8 { + let x = unsafe { _mm512_loadu_si512(a.cast()) }; + let lo = _mm512_and_si512(x, mask_0); + let hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, lo); + let res_hi = _mm512_shuffle_epi8(lut, hi); + let res = _mm512_add_epi8(res_lo, res_hi); + let sad = _mm512_sad_epu8(res, _mm512_setzero_si512()); + sum = _mm512_add_epi64(sum, sad); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xff, n as u32) as u8; + let x = unsafe { _mm512_maskz_loadu_epi64(mask, a.cast()) }; + let lo = _mm512_and_si512(x, mask_0); + let hi = _mm512_and_si512(_mm512_srli_epi16(x, 4), mask_0); + let res_lo = _mm512_shuffle_epi8(lut, lo); + let res_hi = _mm512_shuffle_epi8(lut, hi); + let res = _mm512_add_epi8(res_lo, res_hi); + let sad = _mm512_sad_epu8(res, _mm512_setzero_si512()); + sum = _mm512_add_epi64(sum, sad); + } + _mm512_reduce_add_epi64(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_x_v4(&this) }; + let fallback = fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[u64]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi64, partial_load}; + use core::arch::x86_64::*; + static LUT: [[i8; 16]; 2] = [[0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; 2]; + let lut = unsafe { _mm256_loadu_si256((&raw const LUT).cast()) }; + let mask_0 = _mm256_set1_epi8(0x0f); + let mut sum = _mm256_setzero_si256(); + let mut a = this.as_ptr(); + let mut n = this.len(); + while n >= 4 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let lo = _mm256_and_si256(x, mask_0); + let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); + let res_lo = _mm256_shuffle_epi8(lut, lo); + let res_hi = _mm256_shuffle_epi8(lut, hi); + let res = _mm256_add_epi8(res_lo, res_hi); + let sad = _mm256_sad_epu8(res, _mm256_setzero_si256()); + sum = _mm256_add_epi64(sum, sad); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let lo = _mm256_and_si256(x, mask_0); + let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), mask_0); + let res_lo = _mm256_shuffle_epi8(lut, lo); + let res_hi = _mm256_shuffle_epi8(lut, hi); + let res = _mm256_add_epi8(res_lo, res_hi); + let sad = _mm256_sad_epu8(res, _mm256_setzero_si256()); + sum = _mm256_add_epi64(sum, sad); + } + emulate_mm256_reduce_add_epi64(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let this = (0..126).map(|_| rand::random::()).collect::>(); + let specialized = unsafe { reduce_sum_of_x_v3(&this) }; + let fallback = fallback(&this); + assert_eq!(specialized, fallback); + } + } + + #[crate::multiversion(@"v4:avx512vpopcntdq", @"v4", @"v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_x(this: &[u64]) -> u32 { + let n = this.len(); + let mut sum = 0; + for i in 0..n { + sum += this[i].count_ones(); + } + sum + } +} + +#[inline(always)] +pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_and::vector_and(lhs, rhs) +} + +mod vector_and { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_and(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] & rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_or::vector_or(lhs, rhs) +} + +mod vector_or { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_or(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] | rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { + vector_xor::vector_xor(lhs, rhs) +} + +mod vector_xor { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_xor(lhs: &[u64], rhs: &[u64]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] ^ rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} diff --git a/crates/simd/src/byte.rs b/crates/simd/src/byte.rs new file mode 100644 index 00000000..0c7c4bed --- /dev/null +++ b/crates/simd/src/byte.rs @@ -0,0 +1,611 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vnni")] + fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let sign = _mm512_set1_epi8(-128); + let i8_1 = _mm512_set1_epi8(1); + let i16_1 = _mm512_set1_epi16(1); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + _0 = _mm512_dpbusd_epi32(_0, x, _mm512_xor_si512(sign, y)); + _1 = _mm512_add_epi32(_1, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + _0 = _mm512_dpbusd_epi32(_0, x, _mm512_xor_si512(sign, y)); + _1 = _mm512_add_epi32(_1, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _mm512_slli_epi32(_1, 7))) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_avx512vnni_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { + println!("test {} ... skipped (v4:avx512vnni)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512vnni(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm512_set1_epi16(0x00ff_i16); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_srli_epi16(x, 8); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_srli_epi16(y, 8); + let z_0 = _mm512_madd_epi16(x_0, y_0); + let z_1 = _mm512_madd_epi16(x_1, y_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_srli_epi16(x, 8); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_srli_epi16(y, 8); + let z_0 = _mm512_madd_epi16(x_0, y_0); + let z_1 = _mm512_madd_epi16(x_1, y_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm256_set1_epi16(0x00ff_i16); + let mut _0 = _mm256_setzero_si256(); + let mut _1 = _mm256_setzero_si256(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_srli_epi16(x, 8); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_srli_epi16(y, 8); + let z_0 = _mm256_madd_epi16(x_0, y_0); + let z_1 = _mm256_madd_epi16(x_1, y_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(32, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_srli_epi16(x, 8); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_srli_epi16(y, 8); + let z_0 = _mm256_madd_epi16(x_0, y_0); + let z_1 = _mm256_madd_epi16(x_1, y_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + } + emulate_mm256_reduce_add_epi32(_mm256_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm_set1_epi16(0x00ff_i16); + let mut _0 = _mm_setzero_si128(); + let mut _1 = _mm_setzero_si128(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_srli_epi16(x, 8); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_srli_epi16(y, 8); + let z_0 = _mm_madd_epi16(x_0, y_0); + let z_1 = _mm_madd_epi16(x_1, y_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_srli_epi16(x, 8); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_srli_epi16(y, 8); + let z_0 = _mm_madd_epi16(x_0, y_0); + let z_1 = _mm_madd_epi16(x_1, y_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + } + emulate_mm_reduce_add_epi32(_mm_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "dotprod")] + fn reduce_sum_of_xy_a2_dotprod(lhs: &[u8], rhs: &[u8]) -> u32 { + unsafe extern "C" { + #[link_name = "byte_reduce_sum_of_xy_a2_dotprod"] + unsafe fn f(n: usize, a: *const u8, b: *const u8) -> u32; + } + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_dotprod_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { + println!("test {} ... skipped (a2:dotprod)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2_dotprod(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = vdupq_n_u32(0); + let mut _1 = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let lo = vmull_u8(vget_low_u8(x), vget_low_u8(y)); + let hi = vmull_u8(vget_high_u8(x), vget_high_u8(y)); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let lo = vmull_u8(vget_low_u8(x), vget_low_u8(y)); + let hi = vmull_u8(vget_high_u8(x), vget_high_u8(y)); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + } + vaddvq_u32(vaddq_u32(_0, _1)) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512vnni", @"v4", @"v3", @"v2", @"a2:dotprod", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; + for i in 0..n { + result += (s[i] as u32) * (t[i] as u32); + } + result + } +} + +#[inline(always)] +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) +} + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_x { + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[u8]) -> u32 { + use core::arch::x86_64::*; + let i8_1 = _mm512_set1_epi8(1); + let i16_1 = _mm512_set1_epi16(1); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + sum = _mm512_add_epi32(sum, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + sum = _mm512_add_epi32(sum, _mm512_madd_epi16(_mm512_maddubs_epi16(x, i8_1), i16_1)); + } + _mm512_reduce_add_epi32(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + let i8_1 = _mm256_set1_epi8(1); + let i16_1 = _mm256_set1_epi16(1); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_si256(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + sum = _mm256_add_epi32(sum, _mm256_madd_epi16(_mm256_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(32, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + sum = _mm256_add_epi32(sum, _mm256_madd_epi16(_mm256_maddubs_epi16(x, i8_1), i16_1)); + } + emulate_mm256_reduce_add_epi32(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_v2(this: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + let i8_1 = _mm_set1_epi8(1); + let i16_1 = _mm_set1_epi16(1); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_si128(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + sum = _mm_add_epi32(sum, _mm_madd_epi16(_mm_maddubs_epi16(x, i8_1), i16_1)); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(16, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + sum = _mm_add_epi32(sum, _mm_madd_epi16(_mm_maddubs_epi16(x, i8_1), i16_1)); + } + emulate_mm_reduce_add_epi32(sum) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v2(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[u8]) -> u32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u8(a.cast()) }; + sum = vaddq_u32(sum, vpaddlq_u16(vpaddlq_u8(x))); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(16, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_u8(a.cast()) }; + sum = vaddq_u32(sum, vpaddlq_u16(vpaddlq_u8(x))); + } + vaddvq_u32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_x(this: &[u8]) -> u32 { + let n = this.len(); + let mut sum = 0; + for i in 0..n { + sum += this[i] as u32; + } + sum + } +} + +#[inline(always)] +pub fn reduce_sum_of_x(vector: &[u8]) -> u32 { + reduce_sum_of_x::reduce_sum_of_x(vector) +} diff --git a/crates/simd/src/emulate.rs b/crates/simd/src/emulate.rs new file mode 100644 index 00000000..1647ee2d --- /dev/null +++ b/crates/simd/src/emulate.rs @@ -0,0 +1,229 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[allow(unused_macros)] +macro_rules! partial_load { + (@internal_0) => { + ::core::mem::zeroed() + }; + (@internal_0 $x:expr) => { + $x + }; + ($N:literal, $n:ident, $($p:ident $(= $x:expr)?),+) => { + ( + $( + { + ::core::hint::assert_unchecked($n < $N); + let mut result = [$crate::emulate::partial_load!(@internal_0 $($x)?); $N]; + // LLVM loves `memcpy`, which is much slower. + // So we use an explicit loop here. + for i in 0..$n { + result[i] = $p.add(i).read(); + } + result + }, + )+ + ) + }; +} + +#[allow(unused_imports)] +pub(crate) use partial_load; + +// VP2INTERSECT emulation. +// Díez-Cañas, G. (2021). Faster-Than-Native Alternatives for x86 VP2INTERSECT +// Instructions. arXiv preprint arXiv:2112.06342. +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f")] +pub fn emulate_mm512_2intersect_epi32( + a: core::arch::x86_64::__m512i, + b: core::arch::x86_64::__m512i, +) -> (core::arch::x86_64::__mmask16, core::arch::x86_64::__mmask16) { + use core::arch::x86_64::*; + + let a1 = _mm512_alignr_epi32(a, a, 4); + let a2 = _mm512_alignr_epi32(a, a, 8); + let a3 = _mm512_alignr_epi32(a, a, 12); + let b1 = _mm512_shuffle_epi32(b, _MM_PERM_ADCB); + let b2 = _mm512_shuffle_epi32(b, _MM_PERM_BADC); + let b3 = _mm512_shuffle_epi32(b, _MM_PERM_CBAD); + let m00 = _mm512_cmpeq_epi32_mask(a, b); + let m01 = _mm512_cmpeq_epi32_mask(a, b1); + let m02 = _mm512_cmpeq_epi32_mask(a, b2); + let m03 = _mm512_cmpeq_epi32_mask(a, b3); + let m10 = _mm512_cmpeq_epi32_mask(a1, b); + let m11 = _mm512_cmpeq_epi32_mask(a1, b1); + let m12 = _mm512_cmpeq_epi32_mask(a1, b2); + let m13 = _mm512_cmpeq_epi32_mask(a1, b3); + let m20 = _mm512_cmpeq_epi32_mask(a2, b); + let m21 = _mm512_cmpeq_epi32_mask(a2, b1); + let m22 = _mm512_cmpeq_epi32_mask(a2, b2); + let m23 = _mm512_cmpeq_epi32_mask(a2, b3); + let m30 = _mm512_cmpeq_epi32_mask(a3, b); + let m31 = _mm512_cmpeq_epi32_mask(a3, b1); + let m32 = _mm512_cmpeq_epi32_mask(a3, b2); + let m33 = _mm512_cmpeq_epi32_mask(a3, b3); + + let m0 = m00 | m10 | m20 | m30; + let m1 = m01 | m11 | m21 | m31; + let m2 = m02 | m12 | m22 | m32; + let m3 = m03 | m13 | m23 | m33; + + let res_a = m00 + | m01 + | m02 + | m03 + | (m10 | m11 | m12 | m13).rotate_left(4) + | (m20 | m21 | m22 | m23).rotate_left(8) + | (m30 | m31 | m32 | m33).rotate_right(4); + + let res_b = m0 + | ((0x7777 & m1) << 1) + | ((m1 >> 3) & 0x1111) + | ((0x3333 & m2) << 2) + | ((m2 >> 2) & 0x3333) + | ((0x1111 & m3) << 3) + | ((m3 >> 1) & 0x7777); + (res_a, res_b) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +pub fn emulate_mm256_reduce_add_ps(mut x: core::arch::x86_64::__m256) -> f32 { + use core::arch::x86_64::*; + x = _mm256_add_ps(x, _mm256_permute2f128_ps(x, x, 1)); + x = _mm256_hadd_ps(x, x); + x = _mm256_hadd_ps(x, x); + _mm256_cvtss_f32(x) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "ssse3")] +pub fn emulate_mm_reduce_add_ps(mut x: core::arch::x86_64::__m128) -> f32 { + use core::arch::x86_64::*; + x = _mm_hadd_ps(x, x); + x = _mm_hadd_ps(x, x); + _mm_cvtss_f32(x) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +pub fn emulate_mm256_reduce_add_epi32(mut x: core::arch::x86_64::__m256i) -> i32 { + use core::arch::x86_64::*; + x = _mm256_add_epi32(x, _mm256_permute2f128_si256(x, x, 1)); + x = _mm256_hadd_epi32(x, x); + x = _mm256_hadd_epi32(x, x); + _mm256_cvtsi256_si32(x) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "ssse3")] +pub fn emulate_mm_reduce_add_epi32(mut x: core::arch::x86_64::__m128i) -> i32 { + use core::arch::x86_64::*; + x = _mm_hadd_epi32(x, x); + x = _mm_hadd_epi32(x, x); + _mm_cvtsi128_si32(x) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +pub fn emulate_mm256_reduce_min_ps(x: core::arch::x86_64::__m256) -> f32 { + use crate::aligned::Aligned16; + use core::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let min = _mm_min_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); + unsafe { + _mm_store_ps(x.0.as_mut_ptr(), min); + } + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "sse")] +pub fn emulate_mm_reduce_min_ps(x: core::arch::x86_64::__m128) -> f32 { + use crate::aligned::Aligned16; + use core::arch::x86_64::*; + let min = x; + let mut x = Aligned16([0.0f32; 4]); + unsafe { + _mm_store_ps(x.0.as_mut_ptr(), min); + } + f32::min(f32::min(x.0[0], x.0[1]), f32::min(x.0[2], x.0[3])) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +pub fn emulate_mm256_reduce_max_ps(x: core::arch::x86_64::__m256) -> f32 { + use crate::aligned::Aligned16; + use core::arch::x86_64::*; + let lo = _mm256_castps256_ps128(x); + let hi = _mm256_extractf128_ps(x, 1); + let max = _mm_max_ps(lo, hi); + let mut x = Aligned16([0.0f32; 4]); + unsafe { + _mm_store_ps(x.0.as_mut_ptr(), max); + } + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "sse")] +pub fn emulate_mm_reduce_max_ps(x: core::arch::x86_64::__m128) -> f32 { + use crate::aligned::Aligned16; + use core::arch::x86_64::*; + let max = x; + let mut x = Aligned16([0.0f32; 4]); + unsafe { + _mm_store_ps(x.0.as_mut_ptr(), max); + } + f32::max(f32::max(x.0[0], x.0[1]), f32::max(x.0[2], x.0[3])) +} + +#[inline] +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +pub fn emulate_mm256_reduce_add_epi64(mut x: core::arch::x86_64::__m256i) -> i64 { + use core::arch::x86_64::*; + x = _mm256_add_epi64(x, _mm256_permute2f128_si256(x, x, 1)); + _mm256_extract_epi64(x, 0) + _mm256_extract_epi64(x, 1) +} + +#[inline] +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub fn emulate_vreinterpret_f16_u16( + x: core::arch::aarch64::uint16x4_t, +) -> core::arch::aarch64::float16x4_t { + unsafe { core::mem::transmute(x) } +} + +#[inline] +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub fn emulate_vreinterpretq_f16_u16( + x: core::arch::aarch64::uint16x8_t, +) -> core::arch::aarch64::float16x8_t { + unsafe { core::mem::transmute(x) } +} diff --git a/crates/simd/src/fast_scan.rs b/crates/simd/src/fast_scan.rs new file mode 100644 index 00000000..a02dfc33 --- /dev/null +++ b/crates/simd/src/fast_scan.rs @@ -0,0 +1,775 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +/* + +## code layout for 4-bit quantizer + +group i = | vector i | (total bytes = n/2) + +byte: | 0 | 1 | 2 | ... | n/2 - 1 | +bits 0..3: | code 0 | code 2 | code 4 | ... | code n-2 | +bits 4..7: | code 1 | code 3 | code 5 | ... | code n-1 | + +## packed_code layout for 4-bit quantizer + +group i = | vector 32i | vector 32i+1 | vector 32i+2 | ... | vector 32i+31 | (total bytes = n * 16) + +byte | 0 | 1 | 2 | ... | 14 | 15 | +bits 0..3 | code 0,vector 0 | code 0,vector 8 | code 0,vector 1 | ... | code 0,vector 14 | code 0,vector 15 | +bits 4..7 | code 0,vector 16 | code 0,vector 24 | code 0,vector 17 | ... | code 0,vector 30 | code 0,vector 31 | + +byte | 16 | 17 | 18 | ... | 30 | 31 | +bits 0..3 | code 1,vector 0 | code 1,vector 8 | code 1,vector 1 | ... | code 1,vector 14 | code 1,vector 15 | +bits 4..7 | code 1,vector 16 | code 1,vector 24 | code 1,vector 17 | ... | code 1,vector 30 | code 1,vector 31 | + +byte | 32 | 33 | 34 | ... | 46 | 47 | +bits 0..3 | code 2,vector 0 | code 2,vector 8 | code 2,vector 1 | ... | code 2,vector 14 | code 2,vector 15 | +bits 4..7 | code 2,vector 16 | code 2,vector 24 | code 2,vector 17 | ... | code 2,vector 30 | code 2,vector 31 | + +... + +byte | n*32-32 | n*32-31 | ... | n*32-1 | +bits 0..3 | code (n-1),vector 0 | code (n-1),vector 8 | ... | code (n-1),vector 15 | +bits 4..7 | code (n-1),vector 16 | code (n-1),vector 24 | ... | code (n-1),vector 31 | + +*/ + +mod scan { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn scan_v4(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::x86_64::*; + + #[inline] + #[crate::target_cpu(enable = "v4")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } + + #[inline] + #[crate::target_cpu(enable = "v4")] + fn combine4x2(x0x1x2x3: __m512i, y0y1y2y3: __m512i) -> __m256i { + let x0x1 = _mm512_castsi512_si256(x0x1x2x3); + let x2x3 = _mm512_extracti64x4_epi64(x0x1x2x3, 1); + let y0y1 = _mm512_castsi512_si256(y0y1y2y3); + let y2y3 = _mm512_extracti64x4_epi64(y0y1y2y3, 1); + let x01y01 = combine2x2(x0x1, y0y1); + let x23y23 = combine2x2(x2x3, y2y3); + _mm256_add_epi16(x01y01, x23y23) + } + + let mut accu_0 = _mm512_setzero_si512(); + let mut accu_1 = _mm512_setzero_si512(); + let mut accu_2 = _mm512_setzero_si512(); + let mut accu_3 = _mm512_setzero_si512(); + + let mut i = 0_usize; + while i + 4 <= n { + let code = unsafe { _mm512_loadu_si512(code.as_ptr().add(i).cast()) }; + + let mask = _mm512_set1_epi8(0xf); + let clo = _mm512_and_si512(code, mask); + let chi = _mm512_and_si512(_mm512_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm512_loadu_si512(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_shuffle_epi8(lut, clo); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_shuffle_epi8(lut, chi); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 4; + } + if i + 2 <= n { + let code = unsafe { _mm256_loadu_si256(code.as_ptr().add(i).cast()) }; + + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm256_loadu_si256(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi256_si512(_mm256_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 2; + } + if i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm512_add_epi16(accu_0, res_lo); + accu_1 = _mm512_add_epi16(accu_1, _mm512_srli_epi16(res_lo, 8)); + let res_hi = _mm512_zextsi128_si512(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm512_add_epi16(accu_2, res_hi); + accu_3 = _mm512_add_epi16(accu_3, _mm512_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = _mm512_sub_epi16(accu_0, _mm512_slli_epi16(accu_1, 8)); + unsafe { + _mm256_storeu_si256( + result.as_mut_ptr().add(0).cast(), + combine4x2(accu_0, accu_1), + ); + } + + accu_2 = _mm512_sub_epi16(accu_2, _mm512_slli_epi16(accu_3, 8)); + unsafe { + _mm256_storeu_si256( + result.as_mut_ptr().add(16).cast(), + combine4x2(accu_2, accu_3), + ); + } + + result + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn scan_v4_test() { + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_v4(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn scan_v3(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::x86_64::*; + + #[inline] + #[crate::target_cpu(enable = "v3")] + fn combine2x2(x0x1: __m256i, y0y1: __m256i) -> __m256i { + let x1y0 = _mm256_permute2f128_si256(x0x1, y0y1, 0x21); + let x0y1 = _mm256_blend_epi32(x0x1, y0y1, 0xf0); + _mm256_add_epi16(x1y0, x0y1) + } + + let mut accu_0 = _mm256_setzero_si256(); + let mut accu_1 = _mm256_setzero_si256(); + let mut accu_2 = _mm256_setzero_si256(); + let mut accu_3 = _mm256_setzero_si256(); + + let mut i = 0_usize; + while i + 2 <= n { + let code = unsafe { _mm256_loadu_si256(code.as_ptr().add(i).cast()) }; + + let mask = _mm256_set1_epi8(0xf); + let clo = _mm256_and_si256(code, mask); + let chi = _mm256_and_si256(_mm256_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm256_loadu_si256(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm256_shuffle_epi8(lut, clo); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_shuffle_epi8(lut, chi); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + + i += 2; + } + if i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, clo)); + accu_0 = _mm256_add_epi16(accu_0, res_lo); + accu_1 = _mm256_add_epi16(accu_1, _mm256_srli_epi16(res_lo, 8)); + let res_hi = _mm256_zextsi128_si256(_mm_shuffle_epi8(lut, chi)); + accu_2 = _mm256_add_epi16(accu_2, res_hi); + accu_3 = _mm256_add_epi16(accu_3, _mm256_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = _mm256_sub_epi16(accu_0, _mm256_slli_epi16(accu_1, 8)); + unsafe { + _mm256_storeu_si256( + result.as_mut_ptr().add(0).cast(), + combine2x2(accu_0, accu_1), + ); + } + + accu_2 = _mm256_sub_epi16(accu_2, _mm256_slli_epi16(accu_3, 8)); + unsafe { + _mm256_storeu_si256( + result.as_mut_ptr().add(16).cast(), + combine2x2(accu_2, accu_3), + ); + } + + result + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn scan_v3_test() { + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_v3(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn scan_v2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::x86_64::*; + + let mut accu_0 = _mm_setzero_si128(); + let mut accu_1 = _mm_setzero_si128(); + let mut accu_2 = _mm_setzero_si128(); + let mut accu_3 = _mm_setzero_si128(); + + let mut i = 0_usize; + while i < n { + let code = unsafe { _mm_loadu_si128(code.as_ptr().add(i).cast()) }; + + let mask = _mm_set1_epi8(0xf); + let clo = _mm_and_si128(code, mask); + let chi = _mm_and_si128(_mm_srli_epi16(code, 4), mask); + + let lut = unsafe { _mm_loadu_si128(lut.as_ptr().add(i).cast()) }; + let res_lo = _mm_shuffle_epi8(lut, clo); + accu_0 = _mm_add_epi16(accu_0, res_lo); + accu_1 = _mm_add_epi16(accu_1, _mm_srli_epi16(res_lo, 8)); + let res_hi = _mm_shuffle_epi8(lut, chi); + accu_2 = _mm_add_epi16(accu_2, res_hi); + accu_3 = _mm_add_epi16(accu_3, _mm_srli_epi16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = _mm_sub_epi16(accu_0, _mm_slli_epi16(accu_1, 8)); + unsafe { + _mm_storeu_si128(result.as_mut_ptr().add(0).cast(), accu_0); + _mm_storeu_si128(result.as_mut_ptr().add(8).cast(), accu_1); + } + + accu_2 = _mm_sub_epi16(accu_2, _mm_slli_epi16(accu_3, 8)); + unsafe { + _mm_storeu_si128(result.as_mut_ptr().add(16).cast(), accu_2); + _mm_storeu_si128(result.as_mut_ptr().add(24).cast(), accu_3); + } + + result + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn scan_v2_test() { + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_v2(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn scan_a2(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::aarch64::*; + + let mut accu_0 = vdupq_n_u16(0); + let mut accu_1 = vdupq_n_u16(0); + let mut accu_2 = vdupq_n_u16(0); + let mut accu_3 = vdupq_n_u16(0); + + let mut i = 0_usize; + while i < n { + let code = unsafe { vld1q_u8(code.as_ptr().add(i).cast()) }; + + let clo = vandq_u8(code, vdupq_n_u8(0xf)); + let chi = vshrq_n_u8(code, 4); + + let lut = unsafe { vld1q_u8(lut.as_ptr().add(i).cast()) }; + let res_lo = vreinterpretq_u16_u8(vqtbl1q_u8(lut, clo)); + accu_0 = vaddq_u16(accu_0, res_lo); + accu_1 = vaddq_u16(accu_1, vshrq_n_u16(res_lo, 8)); + let res_hi = vreinterpretq_u16_u8(vqtbl1q_u8(lut, chi)); + accu_2 = vaddq_u16(accu_2, res_hi); + accu_3 = vaddq_u16(accu_3, vshrq_n_u16(res_hi, 8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = vsubq_u16(accu_0, vshlq_n_u16(accu_1, 8)); + unsafe { + vst1q_u16(result.as_mut_ptr().add(0).cast(), accu_0); + vst1q_u16(result.as_mut_ptr().add(8).cast(), accu_1); + } + + accu_2 = vsubq_u16(accu_2, vshlq_n_u16(accu_3, 8)); + unsafe { + vst1q_u16(result.as_mut_ptr().add(16).cast(), accu_2); + vst1q_u16(result.as_mut_ptr().add(24).cast(), accu_3); + } + + result + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn scan_a2_test() { + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_a2(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[cfg(target_arch = "s390x")] + #[crate::target_cpu(enable = "z13")] + fn scan_z13(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + unsafe { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::s390x::*; + use std::mem::transmute; + use vector_unsigned_char as u8x16; + use vector_unsigned_short as u16x8; + + let _0001_u16x8 = vec_splat_u16::<0x0001>(); + let _00ff_u16x8 = vec_splat_u16::<0x00ff>(); + let _ff00_u16x8 = vec_splat_u16::<{ 0xff00u16 as i16 }>(); + + let mut accu_0 = vec_splat_u16::<0>(); + let mut accu_1 = vec_splat_u16::<0>(); + let mut accu_2 = vec_splat_u16::<0>(); + let mut accu_3 = vec_splat_u16::<0>(); + + let mut i = 0_usize; + while i < n { + let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast()); + + let clo = vec_and(code, vec_splat_u8::<0xf>()); + let chi = vec_srl(code, vec_splat_u8::<4>()); + + let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast()); + let res_lo_r = transmute::(vec_perm(lut, lut, clo)); + let res_lo = vec_revb(res_lo_r); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); + let res_hi_r = transmute::(vec_perm(lut, lut, chi)); + let res_hi = vec_revb(res_hi_r); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + accu_0 = vec_sub(accu_0, vec_and(vec_revb(accu_1), _ff00_u16x8)); + vec_xst(accu_0, 0, result.as_mut_ptr().cast()); + vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + + accu_2 = vec_sub(accu_2, vec_and(vec_revb(accu_3), _ff00_u16x8)); + vec_xst(accu_2, 32, result.as_mut_ptr().cast()); + vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + + result + } + } + + #[cfg(all(target_arch = "s390x", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn scan_z13_test() { + if !crate::is_cpu_detected!("z13") { + println!("test {} ... skipped (z13)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_z13(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + /// The instructions required for byte order swapping differ across CPUs, so + /// different instructions needs to be generated for different CPUs for the + /// same code. + #[cfg(target_arch = "powerpc64")] + macro_rules! scan_powerpc64 { + ($name:ident, $cpu:literal) => { + #[crate::target_cpu(enable = $cpu)] + fn $name(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + unsafe { + // bounds checking is not enforced by compiler, so check it manually + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + use core::arch::powerpc64::*; + use std::mem::transmute; + use vector_unsigned_char as u8x16; + use vector_unsigned_short as u16x8; + + #[cfg(target_endian = "big")] + let revb = transmute::<[u8; 16], u8x16>([ + 1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 15, 14, + ]); + + let _0008_u16x8 = vec_splat_u16::<0x0008>(); + let _00ff_u16x8 = vec_splats(0x00ffu16); + let _ff00_u16x8 = vec_splats(0xff00u16); + + let mut accu_0 = vec_splat_u16::<0>(); + let mut accu_1 = vec_splat_u16::<0>(); + let mut accu_2 = vec_splat_u16::<0>(); + let mut accu_3 = vec_splat_u16::<0>(); + + let mut i = 0_usize; + while i < n { + let code: u8x16 = vec_xl((i as isize) * 16, code.as_ptr().cast::()); + + let clo = vec_and(code, vec_splat_u8::<0xf>()); + let chi = vec_srl(code, vec_splat_u8::<4>()); + + let lut: u8x16 = vec_xl((i as isize) * 16, lut.as_ptr().cast::()); + #[cfg(target_endian = "big")] + { + let res_lo_r = transmute::(vec_perm(lut, lut, clo)); + let res_lo = vec_perm(res_lo_r, res_lo_r, revb); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_and(res_lo_r, _00ff_u16x8)); + let res_hi_r = transmute::(vec_perm(lut, lut, chi)); + let res_hi = vec_perm(res_hi_r, res_hi_r, revb); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_and(res_hi_r, _00ff_u16x8)); + } + #[cfg(target_endian = "little")] + { + let res_lo = transmute::(vec_perm(lut, lut, clo)); + accu_0 = vec_add(accu_0, res_lo); + accu_1 = vec_add(accu_1, vec_sr(res_lo, _0008_u16x8)); + let res_hi = transmute::(vec_perm(lut, lut, chi)); + accu_2 = vec_add(accu_2, res_hi); + accu_3 = vec_add(accu_3, vec_sr(res_hi, _0008_u16x8)); + } + + i += 1; + } + debug_assert_eq!(i, n); + + let mut result = [0_u16; 32]; + + #[cfg(target_endian = "big")] + { + accu_0 = + vec_sub(accu_0, vec_and(vec_perm(accu_1, accu_1, revb), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_0 = vec_sub(accu_0, vec_sl(accu_1, _0008_u16x8)); + } + vec_xst(accu_0, 0, result.as_mut_ptr().cast()); + vec_xst(accu_1, 16, result.as_mut_ptr().cast()); + + #[cfg(target_endian = "big")] + { + accu_2 = + vec_sub(accu_2, vec_and(vec_perm(accu_3, accu_3, revb), _ff00_u16x8)); + } + #[cfg(target_endian = "little")] + { + accu_2 = vec_sub(accu_2, vec_sl(accu_3, _0008_u16x8)); + } + vec_xst(accu_2, 32, result.as_mut_ptr().cast()); + vec_xst(accu_3, 48, result.as_mut_ptr().cast()); + + result + } + } + }; + } + + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p9, "p9"); + + #[cfg(all(target_arch = "powerpc64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn scan_p9_test() { + if !crate::is_cpu_detected!("p9") { + println!("test {} ... skipped (p9)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_p9(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p8, "p8"); + + #[cfg(all(target_arch = "powerpc64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn scan_p8_test() { + if !crate::is_cpu_detected!("p8") { + println!("test {} ... skipped (p8)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_p8(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[cfg(target_arch = "powerpc64")] + scan_powerpc64!(scan_p7, "p7"); + + #[cfg(all(target_arch = "powerpc64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn scan_p7_test() { + if !crate::is_cpu_detected!("p7") { + println!("test {} ... skipped (p7)", module_path!()); + return; + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let code = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + let lut = (0..110) + .map(|_| std::array::from_fn(|_| rand::random())) + .collect::>(); + for n in 90..110 { + let code = &code[..n]; + let lut = &lut[..n]; + unsafe { + assert_eq!(scan_p7(&code, &lut), fallback(&code, &lut)); + } + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", @"z13", @"p9", @"p8", @"p7")] + pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + assert_eq!(code.len(), lut.len()); + let n = code.len(); + + let mut result = [0u16; 32]; + + for i in 0..n { + let code = code[i]; + let clo = code.map(|x| x & 0xf); + let chi = code.map(|x| x >> 4); + let lut = lut[i]; + + let res_lo: [u8; 16] = std::array::from_fn(|i| lut[clo[i] as usize]); + let res_hi: [u8; 16] = std::array::from_fn(|i| lut[chi[i] as usize]); + + result[0x00] += res_lo[0x0] as u16; + result[0x01] += res_lo[0x2] as u16; + result[0x02] += res_lo[0x4] as u16; + result[0x03] += res_lo[0x6] as u16; + result[0x04] += res_lo[0x8] as u16; + result[0x05] += res_lo[0xa] as u16; + result[0x06] += res_lo[0xc] as u16; + result[0x07] += res_lo[0xe] as u16; + result[0x08] += res_lo[0x1] as u16; + result[0x09] += res_lo[0x3] as u16; + result[0x0a] += res_lo[0x5] as u16; + result[0x0b] += res_lo[0x7] as u16; + result[0x0c] += res_lo[0x9] as u16; + result[0x0d] += res_lo[0xb] as u16; + result[0x0e] += res_lo[0xd] as u16; + result[0x0f] += res_lo[0xf] as u16; + result[0x10] += res_hi[0x0] as u16; + result[0x11] += res_hi[0x2] as u16; + result[0x12] += res_hi[0x4] as u16; + result[0x13] += res_hi[0x6] as u16; + result[0x14] += res_hi[0x8] as u16; + result[0x15] += res_hi[0xa] as u16; + result[0x16] += res_hi[0xc] as u16; + result[0x17] += res_hi[0xe] as u16; + result[0x18] += res_hi[0x1] as u16; + result[0x19] += res_hi[0x3] as u16; + result[0x1a] += res_hi[0x5] as u16; + result[0x1b] += res_hi[0x7] as u16; + result[0x1c] += res_hi[0x9] as u16; + result[0x1d] += res_hi[0xb] as u16; + result[0x1e] += res_hi[0xd] as u16; + result[0x1f] += res_hi[0xf] as u16; + } + + result + } +} + +#[inline(always)] +pub fn scan(code: &[[u8; 16]], lut: &[[u8; 16]]) -> [u16; 32] { + scan::scan(code, lut) +} + +mod accu { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { + for i in 0..32 { + sum[i] += delta[i] as u32; + } + } +} + +#[inline(always)] +pub fn accu(sum: &mut [u32; 32], delta: &[u16; 32]) { + accu::accu(sum, delta); +} diff --git a/crates/simd/src/fht.rs b/crates/simd/src/fht.rs new file mode 100644 index 00000000..2fec563e --- /dev/null +++ b/crates/simd/src/fht.rs @@ -0,0 +1,182 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[inline(always)] +fn basic_1(x: &mut [f32]) { + assert!(x.len() == (1 << (8 + 1))); + for i in 0..1 << (8 - Q) { + basic_2::(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } +} + +#[inline(always)] +fn basic_2(x: &mut [f32]) { + assert!(x.len() == (1 << (Q + 1))); + for j in 0..1 << Q { + let (l, r) = (x[j], x[j + (1 << Q)]); + (x[j], x[j + (1 << Q)]) = (l + r, l - r); + } +} + +mod step_1 { + seq_macro::seq!( + Q in 0..16 { + mod dispatch_~Q { + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn f(x: &mut [f32]) { + crate::fht::basic_1::(x); + } + } + #[allow(unused_imports)] + pub use dispatch_~Q::f as dispatch_~Q; + } + ); +} + +mod step_2 { + seq_macro::seq!( + Q in 0..16 { + mod dispatch_~Q { + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn f(x: &mut [f32]) { + crate::fht::basic_2::(x); + } + } + #[allow(unused_imports)] + pub use dispatch_~Q::f as dispatch_~Q; + } + ); +} + +macro_rules! fht { + ($p:literal, 0) => { + { + #[crate::multiversion("v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + fn walk(x: &mut [f32]) { + assert!(x.len() == (1 << $p)); + seq_macro::seq!( + Q in 0..$p { + for i in 0..1 << ($p - Q - 1) { + basic_2::(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } + } + ); + } + walk as fn(&mut [f32]) + } + }; + ($p:literal, 1) => { + { + fn walk(x: &mut [f32]) { + assert!(x.len() == (1 << $p)); + seq_macro::seq!( + Q in 0..8 { + for i in (0..(1 << ($p - Q - 1))).step_by(1 << (8 - Q)) { + step_1::dispatch_~Q(&mut x[i << (Q + 1)..][..1 << (8 + 1)]); + } + } + ); + seq_macro::seq!( + Q in 8..$p { + for i in 0..1 << ($p - Q - 1) { + step_2::dispatch_~Q(&mut x[i << (Q + 1)..][..1 << (Q + 1)]); + } + } + ); + } + walk as fn(&mut [f32]) + } + }; +} + +pub fn fht(x: &mut [f32]) { + const FHT: [fn(&mut [f32]); 1 + 16] = [ + |_| (), + fht!(1, 0), + fht!(2, 0), + fht!(3, 0), + fht!(4, 0), + fht!(5, 0), + fht!(6, 0), + fht!(7, 0), + fht!(8, 0), + fht!(9, 1), + fht!(10, 1), + fht!(11, 1), + fht!(12, 1), + fht!(13, 1), + fht!(14, 1), + fht!(15, 1), + fht!(16, 1), + ]; + let n = x.len(); + let Some(i) = n.checked_ilog2() else { + panic!("the dimension of the vector is 0") + }; + if n != (1 << i) { + panic!("the dimension of the vector is not a power of 2"); + } + if i > 16 { + panic!("the dimension of the vector is too large"); + } + FHT[i as usize](x) +} + +#[cfg(test)] +mod tests { + fn native(x: &mut [f32]) { + let n = x.len(); + assert!(n.is_power_of_two()); + let mut h = 1; + while h < n { + for i in (0..n).step_by(h * 2) { + for j in i..i + h { + (x[j], x[j + h]) = (x[j] + x[j + h], x[j] - x[j + h]); + } + } + h *= 2; + } + } + + #[test] + fn fht() { + use rand::RngExt; + use std::iter::zip; + const EPSILON: f32 = 1e-6; + let mut rng = rand::rng(); + let mut n = 1_usize; + while n <= if cfg!(not(miri)) { 65536 } else { 4096 } { + let x = (0..n) + .map(|_| rng.random_range(-1.0_f32..=1.0_f32)) + .collect::>(); + let x_expected = { + let mut x = x.clone(); + native(x.as_mut_slice()); + x + }; + let x_got = { + let mut x = x.clone(); + crate::fht::fht(x.as_mut_slice()); + x + }; + let mse = zip(x_expected, x_got) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + / n as f32; + eprintln!("n = {n}, mse = {mse:.12}"); + assert!(mse <= EPSILON); + n *= 2; + } + } +} diff --git a/crates/simd/src/floating_f16.rs b/crates/simd/src/floating_f16.rs new file mode 100644 index 00000000..1707955d --- /dev/null +++ b/crates/simd/src/floating_f16.rs @@ -0,0 +1,2741 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{F16, Floating, f16}; + +impl Floating for f16 { + #[inline(always)] + fn zero() -> Self { + f16::_ZERO + } + + #[inline(always)] + fn infinity() -> Self { + f16::INFINITY + } + + #[inline(always)] + fn mask(self, m: bool) -> Self { + f16::from_bits(self.to_bits() & (m as u16).wrapping_neg()) + } + + #[inline(always)] + fn scalar_neg(this: Self) -> Self { + -this + } + + #[inline(always)] + fn scalar_add(lhs: Self, rhs: Self) -> Self { + lhs + rhs + } + + #[inline(always)] + fn scalar_sub(lhs: Self, rhs: Self) -> Self { + lhs - rhs + } + + #[inline(always)] + fn scalar_mul(lhs: Self, rhs: Self) -> Self { + lhs * rhs + } + + #[inline(always)] + fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { + reduce_or_of_is_zero_x::reduce_or_of_is_zero_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x(this: &[f16]) -> f32 { + reduce_sum_of_x::reduce_sum_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { + reduce_sum_of_abs_x::reduce_sum_of_abs_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x2(this: &[f16]) -> f32 { + reduce_sum_of_x2::reduce_sum_of_x2(this) + } + + #[inline(always)] + fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { + reduce_min_max_of_x::reduce_min_max_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_xy::reduce_sum_of_xy(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { + reduce_sum_of_d2::reduce_sum_of_d2(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + reduce_sum_of_xy_sparse::reduce_sum_of_xy_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + reduce_sum_of_d2_sparse::reduce_sum_of_d2_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn vector_from_f32(this: &[f32]) -> Vec { + vector_from_f32::vector_from_f32(this) + } + + #[inline(always)] + fn vector_to_f32(this: &[Self]) -> Vec { + vector_to_f32::vector_to_f32(this) + } + + #[inline(always)] + fn vector_add(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_add::vector_add(lhs, rhs) + } + + #[inline(always)] + fn vector_add_inplace(lhs: &mut [Self], rhs: &[Self]) { + vector_add_inplace::vector_add_inplace(lhs, rhs) + } + + #[inline(always)] + fn vector_sub(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_sub::vector_sub(lhs, rhs) + } + + #[inline(always)] + fn vector_mul(lhs: &[Self], rhs: &[Self]) -> Vec { + vector_mul::vector_mul(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar(lhs: &[Self], rhs: f32) -> Vec { + vector_mul_scalar::vector_mul_scalar(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar_inplace(lhs: &mut [Self], rhs: f32) { + vector_mul_scalar_inplace::vector_mul_scalar_inplace(lhs, rhs) + } + + #[inline(always)] + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { + Self::vector_to_f32(this) + } + + #[inline(always)] + fn vector_abs_inplace(this: &mut [Self]) { + vector_abs_inplace::vector_abs_inplace(this); + } +} + +mod reduce_or_of_is_zero_x { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_or_of_is_zero_x(this: &[f16]) -> bool { + for &x in this { + if x == f16::_ZERO { + return true; + } + } + false + } +} + +mod reduce_sum_of_x { + use crate::{F16, f16}; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_x_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_add_ph(_0, x_0); + _1 = _mm512_add_ph(_1, x_1); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_add_ph(_0, x_0); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_add_ph(_1, x_1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_add_ps(sum, x); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_add_ps(sum, x); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_x_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, x); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, x); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_x_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vaddq_f16(sum_0, x_0); + sum_1 = vaddq_f16(sum_1, x_1); + sum_2 = vaddq_f16(sum_2, x_2); + sum_3 = vaddq_f16(sum_3, x_3); + sum_4 = vaddq_f16(sum_4, x_4); + sum_5 = vaddq_f16(sum_5, x_5); + sum_6 = vaddq_f16(sum_6, x_6); + sum_7 = vaddq_f16(sum_7, x_7); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vaddq_f16(sum_0, x_0); + sum_1 = vaddq_f16(sum_1, x_1); + sum_2 = vaddq_f16(sum_2, x_2); + sum_3 = vaddq_f16(sum_3, x_3); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vaddq_f16(sum_4, x_4); + sum_5 = vaddq_f16(sum_5, x_5); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vaddq_f16(sum_6, x_6); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vaddq_f16(sum_7, x_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[f16]) -> f32 { + use crate::emulate::emulate_vreinterpret_f16_u16; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let mut _a = [f16::_ZERO; 4]; + let mut _b = [f16::_ZERO; 4]; + unsafe { + std::ptr::copy_nonoverlapping(a.cast(), _a.as_mut_ptr(), n); + } + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, x); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion( + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_sum_of_x(this: &[f16]) -> f32 { + let n = this.len(); + let mut x = 0.0f32; + for i in 0..n { + x += this[i]._to_f32(); + } + x + } +} + +mod reduce_sum_of_abs_x { + use crate::{F16, f16}; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_abs_x_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_add_ph(_0, _mm512_abs_ph(x_0)); + _1 = _mm512_add_ph(_1, _mm512_abs_ph(x_1)); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + while n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_add_ph(_0, _mm512_abs_ph(x_0)); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_add_ph(_1, _mm512_abs_ph(x_1)); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_abs_x_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_add_ps(sum, _mm512_abs_ps(x)); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_add_ps(sum, _mm512_abs_ps(x)); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_abs_x_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_abs_x_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, _mm256_and_ps(abs, x)); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_add_ps(sum, _mm256_and_ps(abs, x)); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_abs_x_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vaddq_f16(sum_0, vabsq_f16(x_0)); + sum_1 = vaddq_f16(sum_1, vabsq_f16(x_1)); + sum_2 = vaddq_f16(sum_2, vabsq_f16(x_2)); + sum_3 = vaddq_f16(sum_3, vabsq_f16(x_3)); + sum_4 = vaddq_f16(sum_4, vabsq_f16(x_4)); + sum_5 = vaddq_f16(sum_5, vabsq_f16(x_5)); + sum_6 = vaddq_f16(sum_6, vabsq_f16(x_6)); + sum_7 = vaddq_f16(sum_7, vabsq_f16(x_7)); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vaddq_f16(sum_0, vabsq_f16(x_0)); + sum_1 = vaddq_f16(sum_1, vabsq_f16(x_1)); + sum_2 = vaddq_f16(sum_2, vabsq_f16(x_2)); + sum_3 = vaddq_f16(sum_3, vabsq_f16(x_3)); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vaddq_f16(sum_4, vabsq_f16(x_4)); + sum_5 = vaddq_f16(sum_5, vabsq_f16(x_5)); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vaddq_f16(sum_6, vabsq_f16(x_6)); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vaddq_f16(sum_7, vabsq_f16(x_7)); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_abs_x_a2(this: &[f16]) -> f32 { + use crate::emulate::emulate_vreinterpret_f16_u16; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, vabsq_f32(x)); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let mut _a = [f16::_ZERO; 4]; + let mut _b = [f16::_ZERO; 4]; + unsafe { + std::ptr::copy_nonoverlapping(a.cast(), _a.as_mut_ptr(), n); + } + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vaddq_f32(sum, vabsq_f32(x)); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion( + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_sum_of_abs_x(this: &[f16]) -> f32 { + let n = this.len(); + let mut x = 0.0f32; + for i in 0..n { + x += this[i]._to_f32().abs(); + } + x + } +} + +mod reduce_sum_of_x2 { + use crate::{F16, f16}; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_x2_v4_avx512fp16(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + _0 = _mm512_fmadd_ph(x_0, x_0, _0); + _1 = _mm512_fmadd_ph(x_1, x_1, _1); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + _0 = _mm512_fmadd_ph(x_0, x_0, _0); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + _1 = _mm512_fmadd_ph(x_1, x_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v4_avx512fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x2_v4(this: &[f16]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + sum = _mm512_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + sum = _mm512_fmadd_ps(x, x, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_x2_v4(&this) }; + let fallback = fallback(&this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x2_v3(this: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + sum = _mm256_fmadd_ps(x, x, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_x2_a2_fp16(this: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, x_0); + sum_1 = vfmaq_f16(sum_1, x_1, x_1); + sum_2 = vfmaq_f16(sum_2, x_2, x_2); + sum_3 = vfmaq_f16(sum_3, x_3, x_3); + sum_4 = vfmaq_f16(sum_4, x_4, x_4); + sum_5 = vfmaq_f16(sum_5, x_5, x_5); + sum_6 = vfmaq_f16(sum_6, x_6, x_6); + sum_7 = vfmaq_f16(sum_7, x_7, x_7); + (n, a) = unsafe { (n - 64, a.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, x_0); + sum_1 = vfmaq_f16(sum_1, x_1, x_1); + sum_2 = vfmaq_f16(sum_2, x_2, x_2); + sum_3 = vfmaq_f16(sum_3, x_3, x_3); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + sum_4 = vfmaq_f16(sum_4, x_4, x_4); + sum_5 = vfmaq_f16(sum_5, x_5, x_5); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + sum_6 = vfmaq_f16(sum_6, x_6, x_6); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + sum_7 = vfmaq_f16(sum_7, x_7, x_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_a2_fp16(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x2_a2(this: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vfmaq_f32(sum, x, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + sum = vfmaq_f32(sum, x, x); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion( + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_sum_of_x2(this: &[f16]) -> f32 { + let n = this.len(); + let mut x2 = 0.0f32; + for i in 0..n { + x2 += this[i]._to_f32() * this[i]._to_f32(); + } + x2 + } +} + +mod reduce_min_max_of_x { + use crate::{F16, f16}; + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_min_max_of_x_v4_avx512fp16(this: &[f16]) -> (f32, f32) { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_cvtepi16_ph(_mm512_set1_epi16(f16::INFINITY.to_bits() as _)); + let mut max = _mm512_cvtepi16_ph(_mm512_set1_epi16(f16::NEG_INFINITY.to_bits() as _)); + while n >= 32 { + let x = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.cast())) }; + min = _mm512_min_ph(x, min); + max = _mm512_max_ph(x, max); + (n, a) = unsafe { (n - 32, a.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + min = _mm512_mask_min_ph(min, mask, x, min); + max = _mm512_mask_max_ph(max, mask, x, max); + } + let min = { + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(min), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(min), 1)); + _mm512_reduce_min_ps(_mm512_min_ps(s_0, s_1)) + }; + let max = { + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(max), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(max), 1)); + _mm512_reduce_max_ps(_mm512_max_ps(s_0, s_1)) + }; + (min, max) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_v4_avx512fp16_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4_avx512fp16(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_min_max_of_x_v4(this: &[f16]) -> (f32, f32) { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_set1_ps(f32::INFINITY); + let mut max = _mm512_set1_ps(f32::NEG_INFINITY); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + min = _mm512_min_ps(x, min); + max = _mm512_max_ps(x, max); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + min = _mm512_mask_min_ps(min, mask, x, min); + max = _mm512_mask_max_ps(max, mask, x, max); + } + let min = _mm512_reduce_min_ps(min); + let max = _mm512_reduce_max_ps(max); + (min, max) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_min_max_of_x_v3(this: &[f16]) -> (f32, f32) { + use crate::emulate::{ + emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps, partial_load, + }; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm256_set1_ps(f32::INFINITY); + let mut max = _mm256_set1_ps(f32::NEG_INFINITY); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + ( + emulate_mm256_reduce_min_ps(min), + emulate_mm256_reduce_max_ps(max), + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v3(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_min_max_of_x_a2_fp16(this: &[f16]) -> (f32, f32) { + use crate::emulate::{emulate_vreinterpretq_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = emulate_vreinterpretq_f16_u16(vdupq_n_u16(0x7C00u16)); + let mut max = emulate_vreinterpretq_f16_u16(vdupq_n_u16(0xFC00u16)); + while n >= 8 { + let x = emulate_vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + min = vminnmq_f16(x, min); + max = vmaxnmq_f16(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + min = vminnmq_f16(x, min); + max = vmaxnmq_f16(x, max); + } + ( + vminnmvq_f32(vcvt_f32_f16(vminnm_f16( + vget_low_f16(min), + vget_high_f16(min), + ))), + vmaxnmvq_f32(vcvt_f32_f16(vmaxnm_f16( + vget_low_f16(max), + vget_high_f16(max), + ))), + ) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + fn reduce_min_max_of_x_a2_fp16_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a2_fp16(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_min_max_of_x_a2(this: &[f16]) -> (f32, f32) { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = vdupq_n_f32(f32::INFINITY); + let mut max = vdupq_n_f32(f32::NEG_INFINITY); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f16::NAN) }; + (a,) = (_a.as_ptr(),); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + } + (vminnmvq_f32(min), vmaxnmvq_f32(max)) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + fn reduce_min_max_of_x_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + (x[0], x[1]) = (f16::NAN, -f16::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a2(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[crate::multiversion( + @"v4:avx512fp16", @"v4", @"v3", "v2", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_min_max_of_x(this: &[f16]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let n = this.len(); + for i in 0..n { + min = min.min(this[i]._to_f32()); + max = max.max(this[i]._to_f32()); + } + (min, max) + } +} + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + use crate::{F16, f16}; + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_xy_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(32).cast())) }; + _0 = _mm512_fmadd_ph(x_0, y_0, _0); + _1 = _mm512_fmadd_ph(x_1, y_1, _1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + _0 = _mm512_fmadd_ph(x_0, y_0, _0); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + _1 = _mm512_fmadd_ph(x_1, y_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; + sum = _mm512_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; + sum = _mm512_fmadd_ps(x, y, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let specialized = unsafe { reduce_sum_of_xy_v4(&lhs, &rhs) }; + let fallback = fallback(&lhs, &rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + sum = _mm256_fmadd_ps(x, y, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] + fn reduce_sum_of_xy_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_xy_a3_512"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a3_512_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a3_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_xy_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(32).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(40).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(48).cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(56).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + sum_0 = vfmaq_f16(sum_0, x_0, y_0); + sum_1 = vfmaq_f16(sum_1, x_1, y_1); + sum_2 = vfmaq_f16(sum_2, x_2, y_2); + sum_3 = vfmaq_f16(sum_3, x_3, y_3); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + sum_4 = vfmaq_f16(sum_4, x_4, y_4); + sum_5 = vfmaq_f16(sum_5, x_5, y_5); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + sum_6 = vfmaq_f16(sum_6, x_6, y_6); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.cast()) }); + sum_7 = vfmaq_f16(sum_7, x_7, y_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2_fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + sum = vfmaq_f32(sum, x, y); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + sum = vfmaq_f32(sum, x, y); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += lhs[i]._to_f32() * rhs[i]._to_f32(); + } + sum + } +} + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_d2 { + use crate::{F16, f16}; + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512fp16")] + fn reduce_sum_of_d2_v4_avx512fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = _mm512_setzero_ph(); + let mut _1 = _mm512_setzero_ph(); + while n >= 64 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(32).cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(32).cast())) }; + let d_0 = _mm512_sub_ph(x_0, y_0); + let d_1 = _mm512_sub_ph(x_1, y_1); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(a.add(0).cast())) }; + let y_0 = unsafe { _mm512_castsi512_ph(_mm512_loadu_epi16(b.add(0).cast())) }; + let d_0 = _mm512_sub_ph(x_0, y_0); + _0 = _mm512_fmadd_ph(d_0, d_0, _0); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffffffff, n as u32); + let x_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, a.cast())) }; + let y_1 = unsafe { _mm512_castsi512_ph(_mm512_maskz_loadu_epi16(mask, b.cast())) }; + let d_1 = _mm512_sub_ph(x_1, y_1); + _1 = _mm512_fmadd_ph(d_1, d_1, _1); + } + let s_0 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 0)); + let s_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_0), 1)); + let s_2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 0)); + let s_3 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(_mm512_castph_si512(_1), 1)); + _mm512_reduce_add_ps(_mm512_add_ps( + _mm512_add_ps(s_0, s_2), + _mm512_add_ps(s_1, s_3), + )) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_v4_avx512fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512fp16") { + println!("test {} ... skipped (v4:avx512fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4_avx512fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_v4(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_loadu_epi16(b.cast())) }; + let d = _mm512_sub_ps(x, y); + sum = _mm512_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, a.cast())) }; + let y = unsafe { _mm512_cvtph_ps(_mm256_maskz_loadu_epi16(mask, b.cast())) }; + let d = _mm512_sub_ps(x, y); + sum = _mm512_fmadd_ps(d, d, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_d2_v3(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + + let x = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(a.cast())) }; + let y = unsafe { _mm256_cvtph_ps(_mm_loadu_si128(b.cast())) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.512")] + fn reduce_sum_of_d2_a3_512(lhs: &[f16], rhs: &[f16]) -> f32 { + unsafe extern "C" { + #[link_name = "fp16_reduce_sum_of_d2_a3_512"] + unsafe fn f(n: usize, a: *const f16, b: *const f16) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a3_512_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a3.512") { + println!("test {} ... skipped (a3.512)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a3_512(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "fp16")] + fn reduce_sum_of_d2_a2_fp16(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum_0 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_1 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_2 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_3 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_4 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_5 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_6 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + let mut sum_7 = vreinterpretq_f16_u16(vdupq_n_u16(0)); + while n >= 64 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(32).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(40).cast()) }); + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(48).cast()) }); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(56).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(32).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(40).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(48).cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(56).cast()) }); + let d_0 = vsubq_f16(x_0, y_0); + let d_1 = vsubq_f16(x_1, y_1); + let d_2 = vsubq_f16(x_2, y_2); + let d_3 = vsubq_f16(x_3, y_3); + let d_4 = vsubq_f16(x_4, y_4); + let d_5 = vsubq_f16(x_5, y_5); + let d_6 = vsubq_f16(x_6, y_6); + let d_7 = vsubq_f16(x_7, y_7); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n >= 32 { + let x_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let x_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(16).cast()) }); + let x_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(24).cast()) }); + let y_0 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_1 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let y_2 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(16).cast()) }); + let y_3 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(24).cast()) }); + let d_0 = vsubq_f16(x_0, y_0); + let d_1 = vsubq_f16(x_1, y_1); + let d_2 = vsubq_f16(x_2, y_2); + let d_3 = vsubq_f16(x_3, y_3); + sum_0 = vfmaq_f16(sum_0, d_0, d_0); + sum_1 = vfmaq_f16(sum_1, d_1, d_1); + sum_2 = vfmaq_f16(sum_2, d_2, d_2); + sum_3 = vfmaq_f16(sum_3, d_3, d_3); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n >= 16 { + let x_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let x_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(8).cast()) }); + let y_4 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let y_5 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(8).cast()) }); + let d_4 = vsubq_f16(x_4, y_4); + let d_5 = vsubq_f16(x_5, y_5); + sum_4 = vfmaq_f16(sum_4, d_4, d_4); + sum_5 = vfmaq_f16(sum_5, d_5, d_5); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n >= 8 { + let x_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.add(0).cast()) }); + let y_6 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.add(0).cast()) }); + let d_6 = vsubq_f16(x_6, y_6); + sum_6 = vfmaq_f16(sum_6, d_6, d_6); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(8, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(a.cast()) }); + let y_7 = vreinterpretq_f16_u16(unsafe { vld1q_u16(b.cast()) }); + let d_7 = vsubq_f16(x_7, y_7); + sum_7 = vfmaq_f16(sum_7, d_7, d_7); + } + let s_0 = vcvt_f32_f16(vget_low_f16(sum_0)); + let s_1 = vcvt_f32_f16(vget_high_f16(sum_0)); + let s_2 = vcvt_f32_f16(vget_low_f16(sum_1)); + let s_3 = vcvt_f32_f16(vget_high_f16(sum_1)); + let s_4 = vcvt_f32_f16(vget_low_f16(sum_2)); + let s_5 = vcvt_f32_f16(vget_high_f16(sum_2)); + let s_6 = vcvt_f32_f16(vget_low_f16(sum_3)); + let s_7 = vcvt_f32_f16(vget_high_f16(sum_3)); + let s_8 = vcvt_f32_f16(vget_low_f16(sum_4)); + let s_9 = vcvt_f32_f16(vget_high_f16(sum_4)); + let s_a = vcvt_f32_f16(vget_low_f16(sum_5)); + let s_b = vcvt_f32_f16(vget_high_f16(sum_5)); + let s_c = vcvt_f32_f16(vget_low_f16(sum_6)); + let s_d = vcvt_f32_f16(vget_high_f16(sum_6)); + let s_e = vcvt_f32_f16(vget_low_f16(sum_7)); + let s_f = vcvt_f32_f16(vget_high_f16(sum_7)); + let s = vpaddq_f32( + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_0, s_1), vpaddq_f32(s_2, s_3)), + vpaddq_f32(vpaddq_f32(s_4, s_5), vpaddq_f32(s_6, s_7)), + ), + vpaddq_f32( + vpaddq_f32(vpaddq_f32(s_8, s_9), vpaddq_f32(s_a, s_b)), + vpaddq_f32(vpaddq_f32(s_c, s_d), vpaddq_f32(s_e, s_f)), + ), + ); + vaddvq_f32(s) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a2_fp16_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("fp16") { + println!("test {} ... skipped (a2:fp16)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a2_fp16(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_d2_a2(lhs: &[f16], rhs: &[f16]) -> f32 { + use crate::emulate::{emulate_vreinterpret_f16_u16, partial_load}; + use core::arch::aarch64::*; + assert!(lhs.len() == rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(a.cast()) }); + let x = vcvt_f32_f16(x); + let y = emulate_vreinterpret_f16_u16(unsafe { vld1_u16(b.cast()) }); + let y = vcvt_f32_f16(y); + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 2.0; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + let rhs = (0..n) + .map(|_| f16::_from_f32(rng.random_range(-1.0..=1.0))) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512fp16", @"v4", @"v3", #[cfg(target_endian = "little")] @"a3.512", @"a2:fp16", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_d2(lhs: &[f16], rhs: &[f16]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut sum = 0.0_f32; + for i in 0..n { + let d = lhs[i]._to_f32() - rhs[i]._to_f32(); + sum += d * d; + } + sum + } +} + +mod reduce_sum_of_xy_sparse { + // There is no manually-implemented SIMD version. + // Add it if `svecf16` is supported. + + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut sum = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + sum += lval[lp]._to_f32() * rval[rp]._to_f32(); + lp += 1; + rp += 1; + } + Ordering::Less => { + lp += 1; + } + Ordering::Greater => { + rp += 1; + } + } + } + sum + } +} + +mod reduce_sum_of_d2_sparse { + // There is no manually-implemented SIMD version. + // Add it if `svecf16` is supported. + + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f16], ridx: &[u32], rval: &[f16]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut sum = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + let d = lval[lp]._to_f32() - rval[rp]._to_f32(); + sum += d * d; + lp += 1; + rp += 1; + } + Ordering::Less => { + sum += lval[lp]._to_f32() * lval[lp]._to_f32(); + lp += 1; + } + Ordering::Greater => { + sum += rval[rp]._to_f32() * rval[rp]._to_f32(); + rp += 1; + } + } + } + for i in lp..ln { + sum += lval[i]._to_f32() * lval[i]._to_f32(); + } + for i in rp..rn { + sum += rval[i]._to_f32() * rval[i]._to_f32(); + } + sum + } +} + +mod vector_add { + use crate::f16; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_add(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] + rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_add_inplace { + use crate::f16; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_add_inplace(lhs: &mut [f16], rhs: &[f16]) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + for i in 0..n { + lhs[i] += rhs[i]; + } + } +} + +mod vector_sub { + use crate::f16; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_sub(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] - rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul { + use crate::f16; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul(lhs: &[f16], rhs: &[f16]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul_scalar(lhs: &[f16], rhs: f32) -> Vec { + let rhs = f16::_from_f32(rhs); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar_inplace { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul_scalar_inplace(lhs: &mut [f16], rhs: f32) { + let rhs = f16::_from_f32(rhs); + let n = lhs.len(); + for i in 0..n { + lhs[i] *= rhs; + } + } +} + +mod vector_abs_inplace { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_abs_inplace(this: &mut [f16]) { + let n = this.len(); + for i in 0..n { + this[i] = f16::_from_f32(this[i]._to_f32().abs()); + } + } +} + +mod vector_from_f32 { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_from_f32(this: &[f32]) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(f16::_from_f32(this[i])); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_to_f32 { + use crate::{F16, f16}; + + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_to_f32(this: &[f16]) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(this[i]._to_f32()); + } + } + unsafe { + r.set_len(n); + } + r + } +} diff --git a/crates/simd/src/floating_f32.rs b/crates/simd/src/floating_f32.rs new file mode 100644 index 00000000..b298c107 --- /dev/null +++ b/crates/simd/src/floating_f32.rs @@ -0,0 +1,2063 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Floating; + +impl Floating for f32 { + #[inline(always)] + fn zero() -> Self { + 0.0f32 + } + + #[inline(always)] + fn infinity() -> Self { + f32::INFINITY + } + + #[inline(always)] + fn mask(self, m: bool) -> Self { + f32::from_bits(self.to_bits() & (m as u32).wrapping_neg()) + } + + #[inline(always)] + fn scalar_neg(this: Self) -> Self { + -this + } + + #[inline(always)] + fn scalar_add(lhs: Self, rhs: Self) -> Self { + lhs + rhs + } + + #[inline(always)] + fn scalar_sub(lhs: Self, rhs: Self) -> Self { + lhs - rhs + } + + #[inline(always)] + fn scalar_mul(lhs: Self, rhs: Self) -> Self { + lhs * rhs + } + + #[inline(always)] + fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { + reduce_or_of_is_zero_x::reduce_or_of_is_zero_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x(this: &[f32]) -> f32 { + reduce_sum_of_x::reduce_sum_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { + reduce_sum_of_abs_x::reduce_sum_of_abs_x(this) + } + + #[inline(always)] + fn reduce_sum_of_x2(this: &[f32]) -> f32 { + reduce_sum_of_x2::reduce_sum_of_x2(this) + } + + #[inline(always)] + fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { + reduce_min_max_of_x::reduce_min_max_of_x(this) + } + + #[inline(always)] + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_xy::reduce_sum_of_xy(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_d2(lhs: &[Self], rhs: &[Self]) -> f32 { + reduce_sum_of_d2::reduce_sum_of_d2(lhs, rhs) + } + + #[inline(always)] + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + reduce_sum_of_xy_sparse::reduce_sum_of_xy_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + reduce_sum_of_d2_sparse::reduce_sum_of_d2_sparse(lidx, lval, ridx, rval) + } + + #[inline(always)] + fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_add::vector_add(lhs, rhs) + } + + #[inline(always)] + fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { + vector_add_inplace::vector_add_inplace(lhs, rhs) + } + + #[inline(always)] + fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_sub::vector_sub(lhs, rhs) + } + + #[inline(always)] + fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { + vector_mul::vector_mul(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { + vector_mul_scalar::vector_mul_scalar(lhs, rhs) + } + + #[inline(always)] + fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { + vector_mul_scalar_inplace::vector_mul_scalar_inplace(lhs, rhs); + } + + #[inline(always)] + fn vector_from_f32(this: &[f32]) -> Vec { + this.to_vec() + } + + #[inline(always)] + fn vector_to_f32(this: &[f32]) -> Vec { + this.to_vec() + } + + #[inline(always)] + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]> { + this + } + + #[inline(always)] + fn vector_abs_inplace(this: &mut [Self]) { + vector_abs_inplace::vector_abs_inplace(this); + } +} + +mod reduce_or_of_is_zero_x { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn reduce_or_of_is_zero_x(this: &[f32]) -> bool { + for &x in this { + if x == 0.0f32 { + return true; + } + } + false + } +} + +mod reduce_sum_of_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x_v4(this: &[f32]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + sum = _mm512_add_ps(x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + sum = _mm512_add_ps(x, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v4(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x_v3(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + sum = _mm256_add_ps(x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_add_ps(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_add_ps(x, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_x_v2(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_add_ps(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_add_ps(x, sum); + } + emulate_mm_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x_v2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_v2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x_a2(this: &[f32]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + sum = vaddq_f32(x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + sum = vaddq_f32(x, sum); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_x(this: &[f32]) -> f32 { + let n = this.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += this[i]; + } + sum + } +} + +mod reduce_sum_of_abs_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_abs_x_v4(this: &[f32]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let abs_x = _mm512_abs_ps(x); + sum = _mm512_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let abs_x = _mm512_abs_ps(x); + sum = _mm512_add_ps(abs_x, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_abs_x_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.009; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v4(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_abs_x_v3(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let abs = _mm256_castsi256_ps(_mm256_srli_epi32(_mm256_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let abs_x = _mm256_and_ps(abs, x); + sum = _mm256_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let abs_x = _mm256_and_ps(abs, x); + sum = _mm256_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let abs_x = _mm256_and_ps(abs, x); + sum = _mm256_add_ps(abs_x, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_abs_x_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 0.009; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_abs_x_v2(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let abs = _mm_castsi128_ps(_mm_srli_epi32(_mm_set1_epi32(-1), 1)); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let abs_x = _mm_and_ps(abs, x); + sum = _mm_add_ps(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + let abs_x = _mm_and_ps(abs, x); + sum = _mm_add_ps(abs_x, sum); + } + emulate_mm_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_abs_x_v2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.009; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_v2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_abs_x_a2(this: &[f32]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let abs_x = vabsq_f32(x); + sum = vaddq_f32(abs_x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + let abs_x = vabsq_f32(x); + sum = vaddq_f32(abs_x, sum); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_abs_x_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.009; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_abs_x_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_abs_x(this: &[f32]) -> f32 { + let n = this.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += this[i].abs(); + } + sum + } +} + +mod reduce_sum_of_x2 { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_x2_v4(this: &[f32]) -> f32 { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + sum = _mm512_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + sum = _mm512_fmadd_ps(x, x, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x2_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v4(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_x2_v3(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + sum = _mm256_fmadd_ps(x, x, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x2_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v3(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_x2_v2_fma(this: &[f32]) -> f32 { + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_fmadd_ps(x, x, sum); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + sum = _mm_fmadd_ps(x, x, sum); + } + emulate_mm_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_x2_v2_fma_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_v2_fma(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_x2_a2(this: &[f32]) -> f32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + sum = vfmaq_f32(sum, x, x); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + sum = vfmaq_f32(sum, x, x); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_x2_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.008; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let this = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let this = &this[..z]; + let specialized = unsafe { reduce_sum_of_x2_a2(this) }; + let fallback = fallback(this); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_x2(this: &[f32]) -> f32 { + let n = this.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += this[i] * this[i]; + } + sum + } +} + +mod reduce_min_max_of_x { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_min_max_of_x_v4(this: &[f32]) -> (f32, f32) { + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm512_set1_ps(f32::INFINITY); + let mut max = _mm512_set1_ps(f32::NEG_INFINITY); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + min = _mm512_min_ps(x, min); + max = _mm512_max_ps(x, max); + (n, a) = unsafe { (n - 16, a.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + min = _mm512_mask_min_ps(min, mask, x, min); + max = _mm512_mask_max_ps(max, mask, x, max); + } + let min = _mm512_reduce_min_ps(min); + let max = _mm512_reduce_max_ps(max); + (min, max) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v4(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0); + assert_eq!(specialized.1, fallback.1); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_min_max_of_x_v3(this: &[f32]) -> (f32, f32) { + use crate::emulate::{ + emulate_mm256_reduce_max_ps, emulate_mm256_reduce_min_ps, partial_load, + }; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm256_set1_ps(f32::INFINITY); + let mut max = _mm256_set1_ps(f32::NEG_INFINITY); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 8, a.add(8)) }; + } + if n >= 4 { + let x = _mm256_setr_m128(unsafe { _mm_loadu_ps(a) }, _mm_set1_ps(f32::NAN)); + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = _mm256_setr_m128(unsafe { _mm_loadu_ps(a) }, _mm_set1_ps(f32::NAN)); + min = _mm256_min_ps(x, min); + max = _mm256_max_ps(x, max); + } + ( + emulate_mm256_reduce_min_ps(min), + emulate_mm256_reduce_max_ps(max), + ) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v3(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_min_max_of_x_v2(this: &[f32]) -> (f32, f32) { + use crate::emulate::{emulate_mm_reduce_max_ps, emulate_mm_reduce_min_ps, partial_load}; + use core::arch::x86_64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = _mm_set1_ps(f32::INFINITY); + let mut max = _mm_set1_ps(f32::NEG_INFINITY); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + min = _mm_min_ps(x, min); + max = _mm_max_ps(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + min = _mm_min_ps(x, min); + max = _mm_max_ps(x, max); + } + (emulate_mm_reduce_min_ps(min), emulate_mm_reduce_max_ps(max)) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_min_max_of_x_v2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_v2(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_min_max_of_x_a2(this: &[f32]) -> (f32, f32) { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut min = vdupq_n_f32(f32::INFINITY); + let mut max = vdupq_n_f32(f32::NEG_INFINITY); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + (n, a) = unsafe { (n - 4, a.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a = f32::NAN) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + min = vminnmq_f32(x, min); + max = vmaxnmq_f32(x, max); + } + (vminnmvq_f32(min), vmaxnmvq_f32(max)) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_min_max_of_x_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 200; + let mut x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + (x[0], x[1]) = (f32::NAN, -f32::NAN); + for z in 50..200 { + let x = &x[..z]; + let specialized = unsafe { reduce_min_max_of_x_a2(x) }; + let fallback = fallback(x); + assert_eq!(specialized.0, fallback.0,); + assert_eq!(specialized.1, fallback.1,); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_min_max_of_x(this: &[f32]) -> (f32, f32) { + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let n = this.len(); + for i in 0..n { + min = min.min(this[i]); + max = max.max(this[i]); + } + (min, max) + } +} + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let y = unsafe { _mm512_loadu_ps(b) }; + sum = _mm512_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; + sum = _mm512_fmadd_ps(x, y, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let y = unsafe { _mm256_loadu_ps(b) }; + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + sum = _mm256_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + sum = _mm256_fmadd_ps(x, y, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_xy_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + sum = _mm_fmadd_ps(x, y, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + sum = _mm_fmadd_ps(x, y, sum); + } + emulate_mm_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v2_fma_test() { + use rand::RngExt; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2_fma(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] + fn reduce_sum_of_xy_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { + unsafe extern "C" { + #[link_name = "fp32_reduce_sum_of_xy_a3_256"] + unsafe fn f(n: usize, a: *const f32, b: *const f32) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a3_256_test() { + use rand::RngExt; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + sum = vfmaq_f32(sum, x, y); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + sum = vfmaq_f32(sum, x, y); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.004; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut sum = 0.0f32; + for i in 0..n { + sum += lhs[i] * rhs[i]; + } + sum + } +} + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_d2 { + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_v4(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm512_setzero_ps(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let y = unsafe { _mm512_loadu_ps(b) }; + let d = _mm512_sub_ps(x, y); + sum = _mm512_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let y = unsafe { _mm512_maskz_loadu_ps(mask, b) }; + let d = _mm512_sub_ps(x, y); + sum = _mm512_fmadd_ps(d, d, sum); + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_d2_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_d2_v3(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::{emulate_mm256_reduce_add_ps, partial_load}; + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm256_setzero_ps(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let y = unsafe { _mm256_loadu_ps(b) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 8, a.add(8), b.add(8)) }; + } + if n >= 4 { + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(a)) }; + let y = unsafe { _mm256_zextps128_ps256(_mm_loadu_ps(b)) }; + let d = _mm256_sub_ps(x, y); + sum = _mm256_fmadd_ps(d, d, sum); + } + emulate_mm256_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_d2_v3_test() { + use rand::RngExt; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn reduce_sum_of_d2_v2_fma(lhs: &[f32], rhs: &[f32]) -> f32 { + use crate::emulate::{emulate_mm_reduce_add_ps, partial_load}; + assert!(lhs.len() == rhs.len()); + use core::arch::x86_64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = _mm_setzero_ps(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + let d = _mm_sub_ps(x, y); + sum = _mm_fmadd_ps(d, d, sum); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_ps(a) }; + let y = unsafe { _mm_loadu_ps(b) }; + let d = _mm_sub_ps(x, y); + sum = _mm_fmadd_ps(d, d, sum); + } + emulate_mm_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_d2_v2_fma_test() { + use rand::RngExt; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_v2_fma(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[crate::target_cpu(enable = "a3.256")] + fn reduce_sum_of_d2_a3_256(lhs: &[f32], rhs: &[f32]) -> f32 { + unsafe extern "C" { + #[link_name = "fp32_reduce_sum_of_d2_a3_256"] + unsafe fn f(n: usize, a: *const f32, b: *const f32) -> f32; + } + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a3_256_test() { + use rand::RngExt; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("a3.256") { + println!("test {} ... skipped (a3.256)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a3_256(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_d2_a2(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + use crate::emulate::partial_load; + use core::arch::aarch64::*; + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut sum = vdupq_n_f32(0.0); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + (n, a, b) = unsafe { (n - 4, a.add(4), b.add(4)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(4, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_f32(a) }; + let y = unsafe { vld1q_f32(b) }; + let d = vsubq_f32(x, y); + sum = vfmaq_f32(sum, d, d); + } + vaddvq_f32(sum) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_a2_test() { + use rand::RngExt; + const EPSILON: f32 = 0.02; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rhs = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_d2_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", #[cfg(target_endian = "little")] @"a3.256", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_d2(lhs: &[f32], rhs: &[f32]) -> f32 { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let mut sum = 0.0f32; + for i in 0..n { + let d = lhs[i] - rhs[i]; + sum += d * d; + } + sum + } +} + +mod reduce_sum_of_xy_sparse { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + use crate::emulate::emulate_mm512_2intersect_epi32; + assert_eq!(li.len(), lv.len()); + assert_eq!(ri.len(), rv.len()); + let (mut lp, ln) = (0, li.len()); + let (mut rp, rn) = (0, ri.len()); + let (li, lv) = (li.as_ptr(), lv.as_ptr()); + let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); + use core::arch::x86_64::*; + let mut sum = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; + let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; + sum = _mm512_fmadd_ps(lv, rv, sum); + let lt = unsafe { li.add(lp + 16 - 1).read() }; + let rt = unsafe { ri.add(rp + 16 - 1).read() }; + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; + } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()) }; + let rx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; + sum = _mm512_fmadd_ps(lv, rv, sum); + let lt = unsafe { li.add(lp + lw - 1).read() }; + let rt = unsafe { ri.add(rp + rw - 1).read() }; + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_sparse_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.000001; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + pub fn sample_u32_sorted( + rng: &mut (impl RngExt + ?Sized), + length: u32, + amount: u32, + ) -> Vec { + let mut x = match rand::seq::index::sample(rng, length as usize, amount as usize) { + rand::seq::index::IndexVec::U32(x) => x, + _ => unreachable!(), + }; + x.sort(); + x + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lm = 300; + let lidx = sample_u32_sorted(&mut rng, 10000, lm); + let lval = (0..lm) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rm = 350; + let ridx = sample_u32_sorted(&mut rng, 10000, rm); + let rval = (0..rm) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let specialized = unsafe { reduce_sum_of_xy_sparse_v4(&lidx, &lval, &ridx, &rval) }; + let fallback = fallback(&lidx, &lval, &ridx, &rval); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut sum = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + sum += lval[lp] * rval[rp]; + lp += 1; + rp += 1; + } + Ordering::Less => { + lp += 1; + } + Ordering::Greater => { + rp += 1; + } + } + } + sum + } +} + +mod reduce_sum_of_d2_sparse { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_d2_sparse_v4(li: &[u32], lv: &[f32], ri: &[u32], rv: &[f32]) -> f32 { + use crate::emulate::emulate_mm512_2intersect_epi32; + assert_eq!(li.len(), lv.len()); + assert_eq!(ri.len(), rv.len()); + let (mut lp, ln) = (0, li.len()); + let (mut rp, rn) = (0, ri.len()); + let (li, lv) = (li.as_ptr(), lv.as_ptr()); + let (ri, rv) = (ri.as_ptr(), rv.as_ptr()); + use core::arch::x86_64::*; + let mut sum = _mm512_setzero_ps(); + while lp + 16 <= ln && rp + 16 <= rn { + let lx = unsafe { _mm512_loadu_epi32(li.add(lp).cast()) }; + let rx = unsafe { _mm512_loadu_epi32(ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_loadu_ps(lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_loadu_ps(rv.add(rp))) }; + let d = _mm512_sub_ps(lv, rv); + sum = _mm512_fmadd_ps(d, d, sum); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(lv, lv)); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(rv, rv)); + let lt = unsafe { li.add(lp + 16 - 1).read() }; + let rt = unsafe { ri.add(rp + 16 - 1).read() }; + lp += (lt <= rt) as usize * 16; + rp += (lt >= rt) as usize * 16; + } + while lp < ln && rp < rn { + let lw = 16.min(ln - lp); + let rw = 16.min(rn - rp); + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let lx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), lm, li.add(lp).cast()) }; + let rx = + unsafe { _mm512_mask_loadu_epi32(_mm512_set1_epi32(-1), rm, ri.add(rp).cast()) }; + let (lk, rk) = emulate_mm512_2intersect_epi32(lx, rx); + let lv = unsafe { _mm512_maskz_compress_ps(lk, _mm512_maskz_loadu_ps(lm, lv.add(lp))) }; + let rv = unsafe { _mm512_maskz_compress_ps(rk, _mm512_maskz_loadu_ps(rm, rv.add(rp))) }; + let d = _mm512_sub_ps(lv, rv); + sum = _mm512_fmadd_ps(d, d, sum); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(lv, lv)); + sum = _mm512_sub_ps(sum, _mm512_mul_ps(rv, rv)); + let lt = unsafe { li.add(lp + lw - 1).read() }; + let rt = unsafe { ri.add(rp + rw - 1).read() }; + lp += (lt <= rt) as usize * lw; + rp += (lt >= rt) as usize * rw; + } + { + let mut lp = 0; + while lp + 16 <= ln { + let d = unsafe { _mm512_loadu_ps(lv.add(lp)) }; + sum = _mm512_fmadd_ps(d, d, sum); + lp += 16; + } + if lp < ln { + let lw = ln - lp; + let lm = _bzhi_u32(0xffff, lw as _) as u16; + let d = unsafe { _mm512_maskz_loadu_ps(lm, lv.add(lp)) }; + sum = _mm512_fmadd_ps(d, d, sum); + } + } + { + let mut rp = 0; + while rp + 16 <= rn { + let d = unsafe { _mm512_loadu_ps(rv.add(rp)) }; + sum = _mm512_fmadd_ps(d, d, sum); + rp += 16; + } + if rp < rn { + let rw = rn - rp; + let rm = _bzhi_u32(0xffff, rw as _) as u16; + let d = unsafe { _mm512_maskz_loadu_ps(rm, rv.add(rp)) }; + sum = _mm512_fmadd_ps(d, d, sum); + } + } + _mm512_reduce_add_ps(sum) + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_d2_sparse_v4_test() { + use rand::RngExt; + const EPSILON: f32 = 0.0004; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + pub fn sample_u32_sorted( + rng: &mut (impl RngExt + ?Sized), + length: u32, + amount: u32, + ) -> Vec { + let mut x = match rand::seq::index::sample(rng, length as usize, amount as usize) { + rand::seq::index::IndexVec::U32(x) => x, + _ => unreachable!(), + }; + x.sort(); + x + } + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let lm = 300; + let lidx = sample_u32_sorted(&mut rng, 10000, lm); + let lval = (0..lm) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let rm = 350; + let ridx = sample_u32_sorted(&mut rng, 10000, rm); + let rval = (0..rm) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + let specialized = unsafe { reduce_sum_of_d2_sparse_v4(&lidx, &lval, &ridx, &rval) }; + let fallback = fallback(&lidx, &lval, &ridx, &rval); + assert!( + (specialized - fallback).abs() < EPSILON, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + + #[crate::multiversion(@"v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[f32], ridx: &[u32], rval: &[f32]) -> f32 { + use std::cmp::Ordering; + assert_eq!(lidx.len(), lval.len()); + assert_eq!(ridx.len(), rval.len()); + let (mut lp, ln) = (0, lidx.len()); + let (mut rp, rn) = (0, ridx.len()); + let mut sum = 0.0f32; + while lp < ln && rp < rn { + match Ord::cmp(&lidx[lp], &ridx[rp]) { + Ordering::Equal => { + let d = lval[lp] - rval[rp]; + sum += d * d; + lp += 1; + rp += 1; + } + Ordering::Less => { + sum += lval[lp] * lval[lp]; + lp += 1; + } + Ordering::Greater => { + sum += rval[rp] * rval[rp]; + rp += 1; + } + } + } + for i in lp..ln { + sum += lval[i] * lval[i]; + } + for i in rp..rn { + sum += rval[i] * rval[i]; + } + sum + } +} + +mod vector_add { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_add(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] + rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_add_inplace { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_add_inplace(lhs: &mut [f32], rhs: &[f32]) { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + for i in 0..n { + lhs[i] += rhs[i]; + } + } +} + +mod vector_sub { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_sub(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] - rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul(lhs: &[f32], rhs: &[f32]) -> Vec { + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs[i]); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul_scalar(lhs: &[f32], rhs: f32) -> Vec { + let n = lhs.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + unsafe { + r.as_mut_ptr().add(i).write(lhs[i] * rhs); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +mod vector_mul_scalar_inplace { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_mul_scalar_inplace(lhs: &mut [f32], rhs: f32) { + let n = lhs.len(); + for i in 0..n { + lhs[i] *= rhs; + } + } +} + +mod vector_abs_inplace { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn vector_abs_inplace(this: &mut [f32]) { + let n = this.len(); + for i in 0..n { + this[i] = this[i].abs(); + } + } +} diff --git a/crates/simd/src/halfbyte.rs b/crates/simd/src/halfbyte.rs new file mode 100644 index 00000000..e1aa0163 --- /dev/null +++ b/crates/simd/src/halfbyte.rs @@ -0,0 +1,430 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[cfg_attr(feature = "internal", simd_macros::public)] +mod reduce_sum_of_xy { + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + #[target_feature(enable = "avx512vnni")] + fn reduce_sum_of_xy_v4_avx512vnni(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let lo = _mm512_set1_epi8(0x0f_i8); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + _0 = _mm512_dpbusd_epi32(_0, x_0, y_0); + _1 = _mm512_dpbusd_epi32(_1, x_1, y_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + _0 = _mm512_dpbusd_epi32(_0, x_0, y_0); + _1 = _mm512_dpbusd_epi32(_1, x_1, y_1); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_avx512vnni_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") || !crate::is_feature_detected!("avx512vnni") { + println!("test {} ... skipped (v4:avx512vnni)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4_avx512vnni(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn reduce_sum_of_xy_v4(lhs: &[u8], rhs: &[u8]) -> u32 { + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let i16_1 = _mm512_set1_epi16(1); + let lo = _mm512_set1_epi8(0x0f_i8); + let mut _0 = _mm512_setzero_si512(); + let mut _1 = _mm512_setzero_si512(); + while n >= 64 { + let x = unsafe { _mm512_loadu_epi8(a.cast()) }; + let y = unsafe { _mm512_loadu_epi8(b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let z_0 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 64, a.add(64), b.add(64)) }; + } + if n > 0 { + let mask = _bzhi_u64(0xffffffffffffffff, n as u32); + let x = unsafe { _mm512_maskz_loadu_epi8(mask, a.cast()) }; + let y = unsafe { _mm512_maskz_loadu_epi8(mask, b.cast()) }; + let x_0 = _mm512_and_si512(x, lo); + let x_1 = _mm512_and_si512(_mm512_srli_epi16(x, 4), lo); + let y_0 = _mm512_and_si512(y, lo); + let y_1 = _mm512_and_si512(_mm512_srli_epi16(y, 4), lo); + let z_0 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm512_madd_epi16(_mm512_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm512_add_epi32(_0, z_0); + _1 = _mm512_add_epi32(_1, z_1); + } + _mm512_reduce_add_epi32(_mm512_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v4(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn reduce_sum_of_xy_v3(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm256_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let i16_1 = _mm256_set1_epi16(1); + let lo = _mm256_set1_epi8(0x0f_i8); + let mut _0 = _mm256_setzero_si256(); + let mut _1 = _mm256_setzero_si256(); + while n >= 32 { + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_and_si256(_mm256_srli_epi16(x, 4), lo); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_and_si256(_mm256_srli_epi16(y, 4), lo); + let z_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 32, a.add(32), b.add(32)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(32, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm256_loadu_si256(a.cast()) }; + let y = unsafe { _mm256_loadu_si256(b.cast()) }; + let x_0 = _mm256_and_si256(x, lo); + let x_1 = _mm256_and_si256(_mm256_srli_epi16(x, 4), lo); + let y_0 = _mm256_and_si256(y, lo); + let y_1 = _mm256_and_si256(_mm256_srli_epi16(y, 4), lo); + let z_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm256_add_epi32(_0, z_0); + _1 = _mm256_add_epi32(_1, z_1); + } + emulate_mm256_reduce_add_epi32(_mm256_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v3(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + fn reduce_sum_of_xy_v2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::{emulate_mm_reduce_add_epi32, partial_load}; + use core::arch::x86_64::*; + assert_eq!(lhs.len(), rhs.len()); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let i16_1 = _mm_set1_epi16(1); + let lo = _mm_set1_epi8(0x0f_i8); + let mut _0 = _mm_setzero_si128(); + let mut _1 = _mm_setzero_si128(); + while n >= 16 { + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_and_si128(_mm_srli_epi16(x, 4), lo); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_and_si128(_mm_srli_epi16(y, 4), lo); + let z_0 = _mm_madd_epi16(_mm_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm_madd_epi16(_mm_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { _mm_loadu_si128(a.cast()) }; + let y = unsafe { _mm_loadu_si128(b.cast()) }; + let x_0 = _mm_and_si128(x, lo); + let x_1 = _mm_and_si128(_mm_srli_epi16(x, 4), lo); + let y_0 = _mm_and_si128(y, lo); + let y_1 = _mm_and_si128(_mm_srli_epi16(y, 4), lo); + let z_0 = _mm_madd_epi16(_mm_maddubs_epi16(x_0, y_0), i16_1); + let z_1 = _mm_madd_epi16(_mm_maddubs_epi16(x_1, y_1), i16_1); + _0 = _mm_add_epi32(_0, z_0); + _1 = _mm_add_epi32(_1, z_1); + } + emulate_mm_reduce_add_epi32(_mm_add_epi32(_0, _1)) as u32 + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn reduce_sum_of_xy_v2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v2") { + println!("test {} ... skipped (v2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_v2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + #[target_feature(enable = "dotprod")] + fn reduce_sum_of_xy_a2_dotprod(lhs: &[u8], rhs: &[u8]) -> u32 { + unsafe extern "C" { + #[link_name = "halfbyte_reduce_sum_of_xy_a2_dotprod"] + unsafe fn f(n: usize, a: *const u8, b: *const u8) -> u32; + } + assert_eq!(lhs.len(), rhs.len()); + let n = lhs.len(); + let a = lhs.as_ptr(); + let b = rhs.as_ptr(); + unsafe { f(n, a, b) } + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_dotprod_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") || !crate::is_feature_detected!("dotprod") { + println!("test {} ... skipped (a2:dotprod)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2_dotprod(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[cfg_attr(feature = "internal", simd_macros::public)] + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn reduce_sum_of_xy_a2(lhs: &[u8], rhs: &[u8]) -> u32 { + use crate::emulate::partial_load; + use core::arch::aarch64::*; + assert_eq!(lhs.len(), rhs.len()); + let lo = vdupq_n_u8(0x0f_u8); + let mut n = lhs.len(); + let mut a = lhs.as_ptr(); + let mut b = rhs.as_ptr(); + let mut _0 = vdupq_n_u32(0); + let mut _1 = vdupq_n_u32(0); + while n >= 16 { + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let x_0 = vandq_u8(x, lo); + let x_1 = vshrq_n_u8(x, 4); + let y_0 = vandq_u8(y, lo); + let y_1 = vshrq_n_u8(y, 4); + let z_0 = vmulq_u8(x_0, y_0); + let z_1 = vmulq_u8(x_1, y_1); + let lo = vaddl_u8(vget_low_u8(z_0), vget_low_u8(z_1)); + let hi = vaddl_high_u8(z_0, z_1); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + (n, a, b) = unsafe { (n - 16, a.add(16), b.add(16)) }; + } + if n > 0 { + let (_a, _b) = unsafe { partial_load!(16, n, a, b) }; + (a, b) = (_a.as_ptr(), _b.as_ptr()); + let x = unsafe { vld1q_u8(a.cast()) }; + let y = unsafe { vld1q_u8(b.cast()) }; + let x_0 = vandq_u8(x, lo); + let x_1 = vshrq_n_u8(x, 4); + let y_0 = vandq_u8(y, lo); + let y_1 = vshrq_n_u8(y, 4); + let z_0 = vmulq_u8(x_0, y_0); + let z_1 = vmulq_u8(x_1, y_1); + let lo = vaddl_u8(vget_low_u8(z_0), vget_low_u8(z_1)); + let hi = vaddl_high_u8(z_0, z_1); + _0 = vaddq_u32(_0, vpaddlq_u16(lo)); + _1 = vaddq_u32(_1, vpaddlq_u16(hi)); + } + vaddvq_u32(vaddq_u32(_0, _1)) + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn reduce_sum_of_xy_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4016; + let lhs = (0..n).map(|_| rng.random()).collect::>(); + let rhs = (0..n).map(|_| rng.random()).collect::>(); + for z in 3984..4016 { + let lhs = &lhs[..z]; + let rhs = &rhs[..z]; + let specialized = unsafe { reduce_sum_of_xy_a2(lhs, rhs) }; + let fallback = fallback(lhs, rhs); + assert!( + specialized == fallback, + "specialized = {specialized}, fallback = {fallback}." + ); + } + } + } + + #[crate::multiversion(@"v4:avx512vnni", @"v4", @"v3", @"v2", @"a2:dotprod", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + assert_eq!(s.len(), t.len()); + let n = s.len(); + let mut result = 0; + for i in 0..n { + result += (s[i] as u32 & 0xf) * (t[i] as u32 & 0xf); + result += (s[i] as u32 >> 4) * (t[i] as u32 >> 4); + } + result + } +} + +#[inline(always)] +pub fn reduce_sum_of_xy(s: &[u8], t: &[u8]) -> u32 { + reduce_sum_of_xy::reduce_sum_of_xy(s, t) +} diff --git a/crates/simd/src/lib.rs b/crates/simd/src/lib.rs new file mode 100644 index 00000000..8f6f87ef --- /dev/null +++ b/crates/simd/src/lib.rs @@ -0,0 +1,372 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#![allow(unsafe_code)] +#![cfg_attr(feature = "nightly_f16", feature(f16))] +#![cfg_attr(target_arch = "s390x", feature(stdarch_s390x))] +#![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))] +#![cfg_attr(target_arch = "powerpc64", feature(powerpc_target_feature))] +#![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc))] +#![cfg_attr(target_arch = "riscv64", feature(stdarch_riscv_feature_detection))] +#![cfg_attr(target_arch = "riscv64", feature(riscv_target_feature))] + +mod aligned; +mod emulate; + +#[cfg(not(feature = "internal"))] +mod floating_f16; + +#[cfg(feature = "internal")] +pub mod floating_f16; + +#[cfg(not(feature = "internal"))] +mod floating_f32; + +#[cfg(feature = "internal")] +pub mod floating_f32; + +pub mod bit; +pub mod byte; +pub mod fast_scan; +pub mod fht; +pub mod halfbyte; +pub mod quantize; +pub mod rotate; + +#[cfg(not(feature = "nightly_f16"))] +pub use half::f16; + +#[cfg(feature = "nightly_f16")] +pub use f16; + +pub trait F16: Sized { + const _ZERO: Self; + + fn _from_f32(x: f32) -> Self; + + fn _to_f32(self) -> f32; +} + +#[cfg(not(feature = "nightly_f16"))] +impl F16 for f16 { + const _ZERO: Self = f16::ZERO; + + #[inline(always)] + fn _from_f32(x: f32) -> Self { + #[cfg(not(miri))] + { + f16::from_f32(x) + } + #[cfg(miri)] + { + f16::from_f32_const(x) + } + } + + #[inline(always)] + fn _to_f32(self) -> f32 { + #[cfg(not(miri))] + { + self.to_f32() + } + #[cfg(miri)] + { + self.to_f32_const() + } + } +} + +#[cfg(feature = "nightly_f16")] +impl F16 for f16 { + const _ZERO: Self = 0.0; + + #[inline(always)] + fn _from_f32(x: f32) -> Self { + x as _ + } + + #[inline(always)] + fn _to_f32(self) -> f32 { + self as _ + } +} + +pub trait Floating: + Copy + Send + Sync + std::fmt::Debug + Default + 'static + PartialEq + PartialOrd +{ + fn zero() -> Self; + fn infinity() -> Self; + fn mask(self, m: bool) -> Self; + + fn scalar_neg(this: Self) -> Self; + fn scalar_add(lhs: Self, rhs: Self) -> Self; + fn scalar_sub(lhs: Self, rhs: Self) -> Self; + fn scalar_mul(lhs: Self, rhs: Self) -> Self; + + fn reduce_or_of_is_zero_x(this: &[Self]) -> bool; + fn reduce_sum_of_x(this: &[Self]) -> f32; + fn reduce_sum_of_abs_x(this: &[Self]) -> f32; + fn reduce_sum_of_x2(this: &[Self]) -> f32; + fn reduce_min_max_of_x(this: &[Self]) -> (f32, f32); + fn reduce_sum_of_xy(lhs: &[Self], rhs: &[Self]) -> f32; + fn reduce_sum_of_d2(lhs: &[Self], rhs: &[Self]) -> f32; + fn reduce_sum_of_xy_sparse(lidx: &[u32], lval: &[Self], ridx: &[u32], rval: &[Self]) -> f32; + fn reduce_sum_of_d2_sparse(lidx: &[u32], lval: &[Self], ridx: &[u32], rval: &[Self]) -> f32; + + fn vector_from_f32(this: &[f32]) -> Vec; + fn vector_to_f32(this: &[Self]) -> Vec; + fn vector_to_f32_borrowed(this: &[Self]) -> impl AsRef<[f32]>; + fn vector_add(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_add_inplace(lhs: &mut [Self], rhs: &[Self]); + fn vector_sub(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_mul(lhs: &[Self], rhs: &[Self]) -> Vec; + fn vector_mul_scalar(lhs: &[Self], rhs: f32) -> Vec; + fn vector_mul_scalar_inplace(lhs: &mut [Self], rhs: f32); + fn vector_abs_inplace(this: &mut [Self]); +} + +#[doc(hidden)] +pub mod internal { + #[cfg(target_arch = "x86_64")] + simd_macros::define_is_cpu_detected!("x86_64"); + + #[cfg(target_arch = "aarch64")] + simd_macros::define_is_cpu_detected!("aarch64"); + + #[cfg(target_arch = "s390x")] + simd_macros::define_is_cpu_detected!("s390x"); + + #[cfg(target_arch = "powerpc64")] + simd_macros::define_is_cpu_detected!("powerpc64"); + + #[cfg(target_arch = "riscv64")] + simd_macros::define_is_cpu_detected!("riscv64"); + + #[cfg(target_arch = "x86_64")] + #[allow(unused_imports)] + pub use is_x86_64_cpu_detected; + + #[cfg(target_arch = "aarch64")] + #[allow(unused_imports)] + pub use is_aarch64_cpu_detected; + + #[cfg(target_arch = "s390x")] + #[allow(unused_imports)] + pub use is_s390x_cpu_detected; + + #[cfg(target_arch = "powerpc64")] + #[allow(unused_imports)] + pub use is_powerpc64_cpu_detected; + + #[cfg(target_arch = "riscv64")] + #[allow(unused_imports)] + pub use is_riscv64_cpu_detected; + + #[cfg(target_arch = "x86_64")] + pub fn is_v4_detected() -> bool { + std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("avx512cd") + && std::arch::is_x86_feature_detected!("avx512dq") + && std::arch::is_x86_feature_detected!("avx512vl") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("movbe") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "x86_64")] + pub fn is_v3_detected() -> bool { + std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("f16c") + && std::arch::is_x86_feature_detected!("fma") + && std::arch::is_x86_feature_detected!("bmi1") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("lzcnt") + && std::arch::is_x86_feature_detected!("movbe") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "x86_64")] + pub fn is_v2_detected() -> bool { + std::arch::is_x86_feature_detected!("sse4.2") + && std::arch::is_x86_feature_detected!("popcnt") + } + + #[cfg(target_arch = "aarch64")] + #[cfg_attr(target_endian = "big", expect(dead_code))] + pub fn is_a3_512_detected() -> bool { + #[target_feature(enable = "sve")] + fn is_512_detected() -> bool { + let vl: u64; + unsafe { + core::arch::asm!( + "rdvl {0}, #8", + out(reg) vl + ); + } + vl >= 512 + } + std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_512_detected() } + } + + #[cfg(target_arch = "aarch64")] + #[cfg_attr(target_endian = "big", expect(dead_code))] + pub fn is_a3_256_detected() -> bool { + #[target_feature(enable = "sve")] + fn is_256_detected() -> bool { + let vl: u64; + unsafe { + core::arch::asm!( + "rdvl {0}, #8", + out(reg) vl + ); + } + vl >= 256 + } + std::arch::is_aarch64_feature_detected!("sve") && unsafe { is_256_detected() } + } + + #[cfg(target_arch = "aarch64")] + pub fn is_a3_128_detected() -> bool { + std::arch::is_aarch64_feature_detected!("sve") + } + + #[cfg(target_arch = "aarch64")] + pub fn is_a2_detected() -> bool { + std::arch::is_aarch64_feature_detected!("neon") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z17_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-enhancements-3") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-3") + && std::arch::is_s390x_feature_detected!("nnp-assist") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-4") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z16_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement-2") + && std::arch::is_s390x_feature_detected!("nnp-assist") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z15_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-enhancements-2") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal-enhancement") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-3") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z14_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + && std::arch::is_s390x_feature_detected!("vector-enhancements-1") + && std::arch::is_s390x_feature_detected!("vector-packed-decimal") + && std::arch::is_s390x_feature_detected!("miscellaneous-extensions-2") + } + + #[cfg(target_arch = "s390x")] + pub fn is_z13_detected() -> bool { + std::arch::is_s390x_feature_detected!("vector") + } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p9_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + && std::arch::is_powerpc64_feature_detected!("power8-altivec") + && std::arch::is_powerpc64_feature_detected!("power8-crypto") + && std::arch::is_powerpc64_feature_detected!("power8-vector") + && std::arch::is_powerpc64_feature_detected!("power9-altivec") + && std::arch::is_powerpc64_feature_detected!("power9-vector") + } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p8_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + && std::arch::is_powerpc64_feature_detected!("power8-altivec") + && std::arch::is_powerpc64_feature_detected!("power8-crypto") + && std::arch::is_powerpc64_feature_detected!("power8-vector") + } + + #[cfg(target_arch = "powerpc64")] + pub fn is_p7_detected() -> bool { + std::arch::is_powerpc64_feature_detected!("altivec") + && std::arch::is_powerpc64_feature_detected!("vsx") + } + + #[cfg(target_arch = "riscv64")] + pub fn is_r1_detected() -> bool { + std::arch::is_riscv_feature_detected!("v") + } +} + +pub use simd_macros::{multiversion, target_cpu}; + +#[cfg(target_arch = "x86_64")] +#[allow(unused_imports)] +pub use std::arch::is_x86_feature_detected as is_feature_detected; + +#[cfg(target_arch = "aarch64")] +#[allow(unused_imports)] +pub use std::arch::is_aarch64_feature_detected as is_feature_detected; + +#[cfg(target_arch = "s390x")] +#[allow(unused_imports)] +pub use std::arch::is_s390x_feature_detected as is_feature_detected; + +#[cfg(target_arch = "powerpc64")] +#[allow(unused_imports)] +pub use std::arch::is_powerpc64_feature_detected as is_feature_detected; + +#[cfg(target_arch = "x86_64")] +#[allow(unused_imports)] +pub use internal::is_x86_64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "aarch64")] +#[allow(unused_imports)] +pub use internal::is_aarch64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "s390x")] +#[allow(unused_imports)] +pub use internal::is_s390x_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "powerpc64")] +#[allow(unused_imports)] +pub use internal::is_powerpc64_cpu_detected as is_cpu_detected; + +#[cfg(target_arch = "riscv64")] +#[allow(unused_imports)] +pub use internal::is_riscv64_cpu_detected as is_cpu_detected; diff --git a/crates/simd/src/quantize.rs b/crates/simd/src/quantize.rs new file mode 100644 index 00000000..40df3b39 --- /dev/null +++ b/crates/simd/src/quantize.rs @@ -0,0 +1,313 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod mul_add_round { + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v4")] + fn mul_add_round_v4(this: &[f32], k: f32, b: f32) -> Vec { + let mut r = Vec::::with_capacity(this.len()); + use core::arch::x86_64::*; + let lk = _mm512_set1_ps(k); + let lb = _mm512_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 16 { + let x = unsafe { _mm512_loadu_ps(a) }; + let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + unsafe { _mm_storeu_si128(p.cast(), vfl) }; + (n, a, p) = unsafe { (n - 16, a.add(16), p.add(16)) }; + } + if n > 0 { + let mask = _bzhi_u32(0xffff, n as u32) as u16; + let x = unsafe { _mm512_maskz_loadu_ps(mask, a) }; + let v = _mm512_fmadd_round_ps(x, lk, lb, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + let v = _mm512_cvtps_epi32(v); + let vfl = _mm512_cvtepi32_epi8(v); + unsafe { _mm_mask_storeu_epi8(p.cast(), mask, vfl) }; + } + unsafe { + r.set_len(this.len()); + } + r + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn mul_add_round_v4_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v4") { + println!("test {} ... skipped (v4)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 127.0; + let b = 127.0; + let specialized = unsafe { mul_add_round_v4(x, k, b) }; + let fallback = fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v3")] + fn mul_add_round_v3(this: &[f32], k: f32, b: f32) -> Vec { + let mut r = Vec::::with_capacity(this.len().next_multiple_of(8)); + use crate::emulate::partial_load; + use core::arch::x86_64::*; + let cons = _mm256_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + 0, 4, 8, 12, -1, -1, -1, -1, // 16..24 + -1, -1, -1, -1, -1, -1, -1, -1, // 24..32 + ); + let lk = _mm256_set1_ps(k); + let lb = _mm256_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 8 { + let x = unsafe { _mm256_loadu_ps(a) }; + let v = _mm256_fmadd_ps(x, lk, lb); + let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); + let vs = _mm256_shuffle_epi8(v, cons); + let vlo = _mm256_extract_epi32::<0>(vs) as u32; + let vhi = _mm256_extract_epi32::<4>(vs) as u32; + let vfl = vlo as u64 | ((vhi as u64) << 32); + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 8, a.add(8), p.add(8)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(8, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm256_loadu_ps(a) }; + let v = _mm256_fmadd_ps(x, lk, lb); + let v = _mm256_cvtps_epi32(_mm256_round_ps(v, 0x00)); + let vs = _mm256_shuffle_epi8(v, cons); + let vlo = _mm256_extract_epi32::<0>(vs) as u32; + let vhi = _mm256_extract_epi32::<4>(vs) as u32; + let vfl = vlo as u64 | ((vhi as u64) << 32); + unsafe { p.cast::().write_unaligned(vfl) }; + } + unsafe { + r.set_len(this.len()); + } + r + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn mul_add_round_v3_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v3") { + println!("test {} ... skipped (v3)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 127.0; + let b = 127.0; + let specialized = unsafe { mul_add_round_v3(x, k, b) }; + let fallback = fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "x86_64")] + #[crate::target_cpu(enable = "v2")] + #[target_feature(enable = "fma")] + fn mul_add_round_v2_fma(this: &[f32], k: f32, b: f32) -> Vec { + let mut r = Vec::::with_capacity(this.len().next_multiple_of(4)); + use crate::emulate::partial_load; + use core::arch::x86_64::*; + let cons = _mm_setr_epi8( + 0, 4, 8, 12, -1, -1, -1, -1, // 0..8 + -1, -1, -1, -1, -1, -1, -1, -1, // 8..15 + ); + let lk = _mm_set1_ps(k); + let lb = _mm_set1_ps(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 4 { + let x = unsafe { _mm_loadu_ps(a) }; + let v = _mm_fmadd_ps(x, lk, lb); + let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); + let vs = _mm_shuffle_epi8(v, cons); + let vfl = _mm_extract_epi32::<0>(vs) as u32; + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 4, a.add(4), p.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { _mm_loadu_ps(a) }; + let v = _mm_fmadd_ps(x, lk, lb); + let v = _mm_cvtps_epi32(_mm_round_ps(v, 0x00)); + let vs = _mm_shuffle_epi8(v, cons); + let vfl = _mm_extract_epi32::<0>(vs) as u32; + unsafe { p.cast::().write_unaligned(vfl) }; + } + unsafe { + r.set_len(this.len()); + } + r + } + + #[cfg(all(target_arch = "x86_64", test))] + #[test] + fn mul_add_round_v2_fma_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("v2") || !crate::is_feature_detected!("fma") { + println!("test {} ... skipped (v2:fma)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 127.0; + let b = 127.0; + let specialized = unsafe { mul_add_round_v2_fma(x, k, b) }; + let fallback = fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[inline] + #[cfg(target_arch = "aarch64")] + #[crate::target_cpu(enable = "a2")] + fn mul_add_round_a2(this: &[f32], k: f32, b: f32) -> Vec { + let mut r = Vec::::with_capacity(this.len().next_multiple_of(4)); + use crate::emulate::partial_load; + use core::arch::aarch64::*; + #[cfg(target_endian = "little")] + const CONS: [u8; 16] = [ + 0, 4, 8, 12, 0xff, 0xff, 0xff, 0xff, // 0..8 + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 + ]; + #[cfg(target_endian = "big")] + const CONS: [u8; 16] = [ + 12, 8, 4, 0, 0xff, 0xff, 0xff, 0xff, // 0..8 + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 8..15 + ]; + let cons = unsafe { vld1q_u8(CONS.as_ptr()) }; + let lk = vdupq_n_f32(k); + let lb = vdupq_n_f32(b); + let mut n = this.len(); + let mut a = this.as_ptr(); + let mut p = r.as_mut_ptr(); + while n >= 4 { + let x = unsafe { vld1q_f32(a) }; + let v = vfmaq_f32(lb, x, lk); + let v = vcvtnq_u32_f32(v); + let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); + let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); + unsafe { p.cast::().write_unaligned(vfl) }; + (n, a, p) = unsafe { (n - 4, a.add(4), p.add(4)) }; + } + if n > 0 { + let (_a,) = unsafe { partial_load!(4, n, a) }; + (a,) = (_a.as_ptr(),); + let x = unsafe { vld1q_f32(a) }; + let v = vfmaq_f32(lb, x, lk); + let v = vcvtnq_u32_f32(v); + let vs = vqtbl1q_u8(vreinterpretq_u8_u32(v), cons); + let vfl = vgetq_lane_u32::<0>(vreinterpretq_u32_u8(vs)); + unsafe { p.cast::().write_unaligned(vfl) }; + } + unsafe { + r.set_len(this.len()); + } + r + } + + #[cfg(all(target_arch = "aarch64", test))] + #[test] + #[cfg_attr(miri, ignore)] + fn mul_add_round_a2_test() { + use rand::RngExt; + if !crate::is_cpu_detected!("a2") { + println!("test {} ... skipped (a2)", module_path!()); + return; + } + let mut rng = rand::rng(); + for _ in 0..if cfg!(not(miri)) { 256 } else { 1 } { + let n = 4010; + let x = (0..n) + .map(|_| rng.random_range(-1.0..=1.0)) + .collect::>(); + for z in 3990..4010 { + let x = &x[..z]; + let k = 127.0; + let b = 127.0; + let specialized = unsafe { mul_add_round_a2(x, k, b) }; + let fallback = fallback(x, k, b); + assert_eq!(specialized, fallback); + } + } + } + + #[crate::multiversion(@"v4", @"v3", @"v2:fma", @"a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1")] + pub fn mul_add_round(this: &[f32], k: f32, b: f32) -> Vec { + let n = this.len(); + let mut r = Vec::::with_capacity(n); + for i in 0..n { + let x = this[i]; + let v = x.mul_add(k, b).round_ties_even() as u8; + unsafe { + r.as_mut_ptr().add(i).write(v); + } + } + unsafe { + r.set_len(n); + } + r + } +} + +#[inline(always)] +pub fn quantize(lut: &[f32], n: f32) -> (f32, f32, Vec) { + use crate::Floating; + let (min, max) = f32::reduce_min_max_of_x(lut); + let k = 0.0f32.max((max - min) / n); + let b = min; + (k, b, mul_add_round::mul_add_round(lut, 1.0 / k, -b / k)) +} diff --git a/crates/simd/src/rotate.rs b/crates/simd/src/rotate.rs new file mode 100644 index 00000000..5ed0de39 --- /dev/null +++ b/crates/simd/src/rotate.rs @@ -0,0 +1,56 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod givens { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { + assert!(lhs.len() == rhs.len()); + let n = lhs.len(); + let scale = 1.0 / (2.0_f32).sqrt(); + for i in 0..n { + (lhs[i], rhs[i]) = ((lhs[i] + rhs[i]) * scale, (lhs[i] - rhs[i]) * scale); + } + } +} + +pub fn givens(lhs: &mut [f32], rhs: &mut [f32]) { + givens::givens(lhs, rhs) +} + +mod flip { + #[crate::multiversion( + "v4", "v3", "v2", "a2", "z17", "z16", "z15", "z14", "z13", "p9", "p8", "p7", "r1" + )] + pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { + use std::hint::select_unpredictable; + let result: &mut [u32] = zerocopy::transmute_mut!(result); + let (arrays, remainder) = result.as_chunks_mut::<64>(); + let n = arrays.len(); + assert!(n <= 1024); + for i in 0..n { + for j in 0..64 { + arrays[i][j] ^= select_unpredictable((bits[i] & (1 << j)) != 0, 0x80000000, 0); + } + } + for j in 0..remainder.len() { + remainder[j] ^= select_unpredictable((bits[n] & (1 << j)) != 0, 0x80000000, 0); + } + } +} + +pub fn flip(bits: &[u64; 1024], result: &mut [f32]) { + flip::flip(bits, result) +} diff --git a/crates/simd_macros/Cargo.toml b/crates/simd_macros/Cargo.toml new file mode 100644 index 00000000..f206bff0 --- /dev/null +++ b/crates/simd_macros/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "simd_macros" +version.workspace = true +edition.workspace = true +publish = false + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = { version = "1", features = ["proc-macro"] } +quote = "1" +syn = { version = "2", default-features = false, features = [ + "clone-impls", + "full", + "parsing", + "printing", + "proc-macro", +] } diff --git a/crates/simd_macros/src/lib.rs b/crates/simd_macros/src/lib.rs new file mode 100644 index 00000000..ed6f436b --- /dev/null +++ b/crates/simd_macros/src/lib.rs @@ -0,0 +1,279 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod target; + +struct MultiversionVersion { + attrs: Vec, + target: String, + import: bool, +} + +impl syn::parse::Parse for MultiversionVersion { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let attrs = if input.lookahead1().peek(syn::Token![#]) { + syn::Attribute::parse_outer(input)? + } else { + Vec::new() + }; + let import = if input.lookahead1().peek(syn::Token![@]) { + let _: syn::Token![@] = input.parse()?; + true + } else { + false + }; + let target = input.parse::()?.value(); + Ok(Self { + attrs, + target, + import, + }) + } +} + +struct Multiversion { + versions: syn::punctuated::Punctuated, +} + +impl syn::parse::Parse for Multiversion { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + Ok(Multiversion { + versions: syn::punctuated::Punctuated::parse_terminated(input)?, + }) + } +} + +#[proc_macro_attribute] +pub fn multiversion( + attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + let attr = syn::parse_macro_input!(attr as Multiversion); + let item_fn = syn::parse::(item).expect("not a function item"); + let syn::ItemFn { + attrs, + vis, + sig, + block, + } = item_fn; + let name = sig.ident.to_string(); + if sig.constness.is_some() { + panic!("const functions are not supported"); + } + if sig.asyncness.is_some() { + panic!("async functions are not supported"); + } + let generics_params = sig.generics.params.clone(); + for generic_param in generics_params.iter() { + if !matches!(generic_param, syn::GenericParam::Lifetime(_)) { + panic!("generic parameters are not supported"); + } + } + let generics_where = sig.generics.where_clause.clone(); + let inputs = sig.inputs.clone(); + let arguments = { + let mut list = vec![]; + for x in sig.inputs.iter() { + if let syn::FnArg::Typed(y) = x { + if let syn::Pat::Ident(ident) = *y.pat.clone() { + list.push(ident); + } else { + panic!("patterns on parameters are not supported") + } + } else { + panic!("receiver parameters are not supported") + } + } + list + }; + if sig.variadic.is_some() { + panic!("variadic parameters are not supported"); + } + let output = sig.output.clone(); + let mut versions = quote::quote! {}; + let mut cold = quote::quote! {}; + for version in attr.versions { + let attrs = version.attrs.clone(); + let target = version.target.clone(); + let name = syn::Ident::new( + &format!("{name}_{}", target.replace(":", "_").replace(".", "_")), + proc_macro2::Span::mixed_site(), + ); + let s = target.split(":").collect::>(); + let target_cpu = target::TARGET_CPUS + .iter() + .find(|target_cpu| target_cpu.target_cpu == s[0]) + .expect("unknown target_cpu"); + let additional_target_features = s[1..].to_vec(); + let target_arch = target_cpu.target_arch; + let target_cpu = target_cpu.target_cpu; + if !version.import { + versions.extend(quote::quote! { + #(#attrs)* + #[inline] + #[cfg(target_arch = #target_arch)] + #[crate::target_cpu(enable = #target_cpu)] + #(#[target_feature(enable = #additional_target_features)])* + fn #name < #generics_params > (#inputs) #output #generics_where { #block } + }); + } + cold.extend(quote::quote! { + #(#attrs)* + #[cfg(target_arch = #target_arch)] + if crate::is_cpu_detected!(#target_cpu) #(&& crate::is_feature_detected!(#additional_target_features))* { + let ptr = unsafe { std::mem::transmute::(#name) }; + CACHE.store(ptr as *mut (), core::sync::atomic::Ordering::Relaxed); + return ptr; + } + }); + } + cold.extend(quote::quote! { + let ptr = unsafe { std::mem::transmute::(fallback)} ; + CACHE.store(ptr as *mut (), core::sync::atomic::Ordering::Relaxed); + ptr + }); + quote::quote! { + #versions + fn fallback < #generics_params > (#inputs) #output #generics_where { #block } + #[must_use] + pub(crate) fn pointer() -> fn(#inputs) #output { + static CACHE: core::sync::atomic::AtomicPtr<()> = core::sync::atomic::AtomicPtr::new(core::ptr::null_mut()); + #[must_use] + #[cold] + fn cold() -> fn(#inputs) #output { + #cold + } + #[cfg(all(feature = "init", not(miri)))] + #[cfg(all(target_os = "linux", target_env = "gnu"))] + { + #[used] + #[unsafe(link_section = ".init_array")] + static INIT: extern "C" fn() -> usize = { + #[unsafe(link_section = ".text.startup")] + extern "C" fn f() -> usize { + let _ = cold(); + 0 + } + f + }; + let cache = unsafe { CACHE.as_ptr().read() }; + if !cache.is_null() { + let ptr = unsafe { std::mem::transmute::<*mut (), fn(#inputs) #output>(cache) }; + return ptr; + } + panic!("feature `init` is not supported on this platform") + } + #[allow(unreachable_code)] + { + let cache = CACHE.load(core::sync::atomic::Ordering::Relaxed); + if !cache.is_null() { + let ptr = unsafe { std::mem::transmute::<*mut (), fn(#inputs) #output>(cache) }; + return ptr; + } + cold() + } + } + #[inline(always)] + #(#attrs)* #vis #sig { + pointer()(#(#arguments,)*) + } + } + .into() +} + +struct TargetCpu { + enable: String, +} + +impl syn::parse::Parse for TargetCpu { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let _: syn::Ident = input.parse()?; + let _: syn::Token![=] = input.parse()?; + let enable: syn::LitStr = input.parse()?; + Ok(Self { + enable: enable.value(), + }) + } +} + +#[proc_macro_attribute] +pub fn target_cpu( + attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + let attr = syn::parse_macro_input!(attr as TargetCpu); + let mut result = quote::quote! {}; + for s in attr.enable.split(',') { + let target_cpu = target::TARGET_CPUS + .iter() + .find(|target_cpu| target_cpu.target_cpu == s) + .expect("unknown target_cpu"); + let target_features = target_cpu.target_features; + result.extend(quote::quote!( + #(#[target_feature(enable = #target_features)])* + )); + } + result.extend(proc_macro2::TokenStream::from(item)); + result.into() +} + +#[proc_macro] +pub fn define_is_cpu_detected(input: proc_macro::TokenStream) -> proc_macro::TokenStream { + let target_arch = syn::parse_macro_input!(input as syn::LitStr).value(); + let mut arms = quote::quote! { + () => { compile_error!("cpu is not provided") }; + }; + for target_cpu in target::TARGET_CPUS { + if target_cpu.target_arch != target_arch { + continue; + } + let target_cpu = target_cpu.target_cpu; + let ident = syn::Ident::new( + &format!("is_{}_detected", target_cpu.replace('.', "_")), + proc_macro2::Span::mixed_site(), + ); + arms.extend(quote::quote! { + (#target_cpu) => { $crate::internal::#ident() }; + }); + } + let ident = syn::Ident::new( + &format!("is_{target_arch}_cpu_detected"), + proc_macro2::Span::mixed_site(), + ); + quote::quote! { + #[macro_export] + macro_rules! #ident { + #arms + } + } + .into() +} + +#[proc_macro_attribute] +pub fn public( + _attr: proc_macro::TokenStream, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + use quote::ToTokens; + + let mut input: syn::Item = syn::parse_macro_input!(item); + + if let syn::Item::Fn(syn::ItemFn { ref mut vis, .. }) + | syn::Item::Mod(syn::ItemMod { ref mut vis, .. }) = input + { + *vis = syn::Visibility::Public(Default::default()); + } + + input.to_token_stream().into() +} diff --git a/crates/simd_macros/src/target.rs b/crates/simd_macros/src/target.rs new file mode 100644 index 00000000..0f3c60c7 --- /dev/null +++ b/crates/simd_macros/src/target.rs @@ -0,0 +1,161 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub struct TargetCpu { + pub target_cpu: &'static str, + pub target_arch: &'static str, + pub target_features: &'static [&'static str], +} + +pub const TARGET_CPUS: &[TargetCpu] = &[ + TargetCpu { + target_cpu: "v4", + target_arch: "x86_64", + target_features: &[ + "avx512bw", "avx512cd", "avx512dq", "avx512vl", // simd + "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations + ], + }, + TargetCpu { + target_cpu: "v3", + target_arch: "x86_64", + target_features: &[ + "avx2", "f16c", "fma", // simd + "bmi1", "bmi2", "lzcnt", "movbe", "popcnt", // bit-operations + ], + }, + TargetCpu { + target_cpu: "v2", + target_arch: "x86_64", + target_features: &[ + "sse4.2", // simd + "popcnt", // bit-operations + ], + }, + TargetCpu { + target_cpu: "a3.512", + target_arch: "aarch64", + target_features: &["sve"], + }, + TargetCpu { + target_cpu: "a3.256", + target_arch: "aarch64", + target_features: &["sve"], + }, + TargetCpu { + target_cpu: "a3.128", + target_arch: "aarch64", + target_features: &["sve"], + }, + TargetCpu { + target_cpu: "a2", + target_arch: "aarch64", + target_features: &["neon"], + }, + TargetCpu { + target_cpu: "z17", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-enhancements-3", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "vector-packed-decimal-enhancement-2", + "vector-packed-decimal-enhancement-3", + "nnp-assist", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", + "miscellaneous-extensions-4", + ], + }, + TargetCpu { + target_cpu: "z16", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "vector-packed-decimal-enhancement-2", + "nnp-assist", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", + ], + }, + TargetCpu { + target_cpu: "z15", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-enhancements-2", + "vector-packed-decimal", + "vector-packed-decimal-enhancement", + "miscellaneous-extensions-2", + "miscellaneous-extensions-3", + ], + }, + TargetCpu { + target_cpu: "z14", + target_arch: "s390x", + target_features: &[ + "vector", + "vector-enhancements-1", + "vector-packed-decimal", + "miscellaneous-extensions-2", + ], + }, + TargetCpu { + target_cpu: "z13", + target_arch: "s390x", + target_features: &["vector"], + }, + TargetCpu { + target_cpu: "p9", + target_arch: "powerpc64", + target_features: &[ + "altivec", + "vsx", + "power8-altivec", + "power8-crypto", + "power8-vector", + "power9-altivec", + "power9-vector", + ], + }, + TargetCpu { + target_cpu: "p8", + target_arch: "powerpc64", + target_features: &[ + "altivec", + "vsx", + "power8-altivec", + "power8-crypto", + "power8-vector", + ], + }, + TargetCpu { + target_cpu: "p7", + target_arch: "powerpc64", + target_features: &["altivec", "vsx"], + }, + TargetCpu { + target_cpu: "r1", + target_arch: "riscv64", + target_features: &["v"], + }, +]; diff --git a/crates/small_iter/Cargo.toml b/crates/small_iter/Cargo.toml new file mode 100644 index 00000000..011d90b9 --- /dev/null +++ b/crates/small_iter/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "small_iter" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] + +[lints] +workspace = true diff --git a/crates/small_iter/src/borrowed.rs b/crates/small_iter/src/borrowed.rs new file mode 100644 index 00000000..ff90cc8d --- /dev/null +++ b/crates/small_iter/src/borrowed.rs @@ -0,0 +1,145 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::stack::StackIter; +use std::marker::PhantomData; +use std::ptr::NonNull; + +#[derive(Clone, Copy)] +pub struct HeapIter<'a, T: Copy> { + off: u16, + len: u16, + pointer: NonNull, + _phantom: PhantomData<&'a T>, +} + +impl<'a, T: Copy> HeapIter<'a, T> { + #[cold] + pub(crate) fn from_slice(slice: &'a [T]) -> Self { + assert!(slice.len() <= 65535_usize); + let c = NonNull::from_ref(slice).cast::(); + Self { + off: 0, + len: slice.len() as u16, + pointer: c, + _phantom: PhantomData, + } + } +} + +impl<'a, T: Copy> Iterator for HeapIter<'a, T> { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + #[allow(unsafe_code)] + unsafe { + if self.off < self.len { + let r = self.pointer.as_ptr().add(self.off as _).read(); + self.off += 1; + Some(r) + } else { + None + } + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + let size = (self.len - self.off) as usize; + (size, Some(size)) + } +} + +impl<'a, T: Copy> ExactSizeIterator for HeapIter<'a, T> {} + +#[derive(Clone, Copy)] +pub enum Iter<'a, T: Copy, const N: usize> { + Stack(StackIter), + Heap(HeapIter<'a, T>), +} + +impl<'a, T: Copy, const N: usize> Iter<'a, T, N> { + #[inline(always)] + pub fn from_slice(slice: &[T], alloc: impl Fn(&[T]) -> &'a [T]) -> Self { + assert!(slice.len() <= 65535); + if slice.len() <= N && N <= 65535 { + Iter::Stack(StackIter::from_slice(slice)) + } else { + Iter::Heap(HeapIter::from_slice(alloc(slice))) + } + } +} + +impl<'a, T: Copy, const N: usize> Iterator for Iter<'a, T, N> { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + match self { + Iter::Stack(iter) => iter.next(), + Iter::Heap(iter) => iter.next(), + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + match self { + Iter::Stack(iter) => iter.size_hint(), + Iter::Heap(iter) => iter.size_hint(), + } + } +} + +impl<'a, T: Copy, const N: usize> ExactSizeIterator for Iter<'a, T, N> {} + +#[cfg(target_pointer_width = "64")] +const _: () = { + assert!(size_of::>() == 16); +}; + +#[test] +fn tests() { + let alloc = |slice: &[u32]| -> &'static [u32] { + // make miri happy + static GLOBAL: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + let pointer = Vec::leak(slice.to_vec()); + GLOBAL.lock().expect("failed to lock").push(pointer); + pointer + }; + for x in Iter::::from_slice(&[1; 0], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 1], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 2], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 3], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 4], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 5], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 6], alloc) { + assert_eq!(x, 1); + } + for x in Iter::::from_slice(&[1; 7], alloc) { + assert_eq!(x, 1); + } +} diff --git a/crates/small_iter/src/lib.rs b/crates/small_iter/src/lib.rs new file mode 100644 index 00000000..60b4aa34 --- /dev/null +++ b/crates/small_iter/src/lib.rs @@ -0,0 +1,16 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod borrowed; +pub mod stack; diff --git a/crates/small_iter/src/stack.rs b/crates/small_iter/src/stack.rs new file mode 100644 index 00000000..1bab4a4a --- /dev/null +++ b/crates/small_iter/src/stack.rs @@ -0,0 +1,66 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::mem::MaybeUninit; + +#[derive(Clone, Copy)] +pub struct StackIter { + off: u16, + len: u16, + buffer: [MaybeUninit; N], +} + +impl StackIter { + #[inline(always)] + pub(crate) fn from_slice(slice: &[T]) -> Self { + assert!(slice.len() <= N && N <= 65535); + Self { + off: 0, + len: slice.len() as u16, + buffer: { + let mut buffer = [const { MaybeUninit::uninit() }; N]; + for i in 0..slice.len() { + buffer[i] = MaybeUninit::new(slice[i]); + } + buffer + }, + } + } +} + +impl Iterator for StackIter { + type Item = T; + + #[inline(always)] + fn next(&mut self) -> Option { + #[allow(unsafe_code)] + unsafe { + if self.off < self.len { + let r = self.buffer[self.off as usize].assume_init(); + self.off += 1; + Some(r) + } else { + None + } + } + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + let size = (self.len - self.off) as usize; + (size, Some(size)) + } +} + +impl ExactSizeIterator for StackIter {} diff --git a/crates/vchordg/Cargo.toml b/crates/vchordg/Cargo.toml new file mode 100644 index 00000000..c1772834 --- /dev/null +++ b/crates/vchordg/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "vchordg" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } +index = { path = "../index" } +index_accessor = { path = "../index_accessor" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +min-max-heap = "1.3.0" +rand.workspace = true +serde.workspace = true +validator.workspace = true +zerocopy.workspace = true + +[lints] +workspace = true diff --git a/crates/vchordg/src/build.rs b/crates/vchordg/src/build.rs new file mode 100644 index 00000000..54b69d2e --- /dev/null +++ b/crates/vchordg/src/build.rs @@ -0,0 +1,68 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use crate::tuples::{MetaTuple, OptionPointer, Tuple}; +use crate::types::{VchordgIndexOptions, VectorOptions}; +use index::relation::{Page, PageGuard, RelationWrite}; + +pub fn build( + vector_options: VectorOptions, + index_options: VchordgIndexOptions, + index: &R, +) where + R::Page: Page, +{ + let mut meta_guard = index.extend( + Opaque { + next: u32::MAX, + link: 1, + }, + false, + ); + assert_eq!(meta_guard.id(), 0); + let vertex_guard = index.extend( + Opaque { + next: u32::MAX, + link: 2, + }, + true, + ); + assert_eq!(vertex_guard.id(), 1); + drop(vertex_guard); + let vector_guard = index.extend( + Opaque { + next: u32::MAX, + link: u32::MAX, + }, + false, + ); + assert_eq!(vector_guard.id(), 2); + drop(vector_guard); + let serialized = MetaTuple::serialize(&MetaTuple { + dim: vector_options.dim, + bits: index_options.bits, + m: index_options.m, + alpha: index_options.alpha, + ef_construction: index_options.ef_construction, + beam_construction: index_options.beam_construction, + start: OptionPointer::NONE, + skip: 1, + }); + let i = meta_guard + .alloc(&serialized) + .expect("implementation: a free page cannot accommodate a single tuple"); + assert_eq!(i, 1); +} diff --git a/crates/vchordg/src/bulkdelete.rs b/crates/vchordg/src/bulkdelete.rs new file mode 100644 index 00000000..2e619837 --- /dev/null +++ b/crates/vchordg/src/bulkdelete.rs @@ -0,0 +1,70 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use crate::tuples::{MetaTuple, VertexTuple, WithReader, WithWriter}; +use index::relation::{Page, RelationRead, RelationWrite}; +use std::num::NonZero; + +pub fn bulkdelete( + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let start = meta_tuple.start(); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + let Some(_) = start.into_inner() else { + return; + }; + let mut current = link; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(bytes) = read.get(i) { + let tuple = VertexTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + if let Some(bytes) = write.get_mut(i) { + let mut tuple = VertexTuple::deserialize_mut(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + *p = None; + } + } + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } +} diff --git a/crates/vchordg/src/candidates.rs b/crates/vchordg/src/candidates.rs new file mode 100644 index 00000000..2c32019f --- /dev/null +++ b/crates/vchordg/src/candidates.rs @@ -0,0 +1,73 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use index::fetch::Fetch; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::RelationRead; +use std::collections::{BinaryHeap, VecDeque}; +use std::marker::PhantomData; + +pub struct Candidates<'a, P, F, R> +where + P: Prefetcher<'a>, + P::Item: Fetch<'a>, +{ + beam: usize, + front: Option

, + heap: BinaryHeap, + prefetch: F, + _phantom: PhantomData &'a R>, +} + +impl<'a, P, F, R> Candidates<'a, P, F, R> +where + P: Prefetcher<'a, R = R>, + P::Item: Fetch<'a> + Ord, + F: PrefetcherSequenceFamily<'a, R, P> = P>, + R: RelationRead, +{ + pub fn new(beam: usize, prefetch: F) -> Self { + assert_ne!(beam, 0); + Self { + beam, + front: None, + heap: BinaryHeap::new(), + prefetch, + _phantom: PhantomData, + } + } + pub fn pop(&mut self) -> Option<(P::Item, P::Guards)> { + if let Some(front) = self.front.as_mut() + && let Some(item) = front.next() + { + return Some(item); + } + self.front = Some( + self.prefetch.prefetch( + (0..self.beam) + .flat_map(|_| self.heap.pop()) + .collect::>(), + ), + ); + if let Some(front) = self.front.as_mut() + && let Some(item) = front.next() + { + return Some(item); + } + None + } + pub fn push(&mut self, item: P::Item) { + self.heap.push(item); + } +} diff --git a/crates/vchordg/src/insert.rs b/crates/vchordg/src/insert.rs new file mode 100644 index 00000000..ff11ee9f --- /dev/null +++ b/crates/vchordg/src/insert.rs @@ -0,0 +1,395 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::candidates::Candidates; +use crate::operator::{CloneAccessor, Operator, Vector}; +use crate::results::Results; +use crate::tuples::*; +use crate::types::DistanceKind; +use crate::vectors::{by_prefetch, by_read, copy_all, copy_nothing, copy_outs, update}; +use crate::visited::Visited; +use always_equal::AlwaysEqual; +use index::bump::Bump; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::{DefaultWithDimension, LAccess}; +use rabitq::bits::Bits; +use std::cmp::Reverse; +use std::collections::VecDeque; +use std::num::{NonZero, Wrapping}; +use vector::{VectorBorrowed, VectorOwned}; + +pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + vector: ::Borrowed<'b>, + payload: NonZero, + bump: &'b impl Bump, + mut prefetch_vertices: impl PrefetcherSequenceFamily<'b, R> + 'b, + prefetch_vectors: impl PrefetcherSequenceFamily<'b, R> + 'b, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + let start = meta_tuple.start(); + let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); + let m = meta_tuple.m(); + let alpha = meta_tuple.alpha().to_vec(); + let ef = meta_tuple.ef_construction(); + let beam = meta_tuple.beam_construction(); + let skip = meta_tuple.skip(); + drop(meta_guard); + let version_t = Wrapping(rand::random()); + let (pointers_t, t) = { + let list_of_vector_bytes = { + let (left, right) = O::Vector::split(vector, m as _); + let left = left.into_iter().enumerate().map(|(index, elements)| { + VectorTuple::serialize(&VectorTuple::::_1 { + payload: Some(payload), + elements: elements.to_vec(), + index: index as u32, + }) + }); + let right = VectorTuple::serialize(&VectorTuple::::_0 { + payload: Some(payload), + elements: right.0.to_vec(), + metadata: right.1, + neighbours: vec![OptionNeighbour::NONE; m as usize], + version: version_t, + }); + left.chain(std::iter::once(right)).collect::>() + }; + let mut vertex_bytes = { + let code = O::Vector::code(bits, vector); + VertexTuple::serialize(&VertexTuple { + metadata: code.0.into_array(), + payload: Some(payload), + elements: rabitq::bits::pack_code(bits, &code.1), + pointers: vec![Pointer::new((u32::MAX, 0)); list_of_vector_bytes.len()], // a sentinel value + }) + }; + let mut vertex_guard = if let Some(guard) = index.search(vertex_bytes.len()) { + guard + } else { + append_vertex_tuple(index, skip, vertex_bytes.len()) + }; + let link = vertex_guard.get_opaque().link; + let pointers_t = list_of_vector_bytes + .iter() + .map(|vector_bytes| { + let mut vector_guard = append_vector_tuple(index, link, vector_bytes.len()); + let i = vector_guard + .alloc(vector_bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + Pointer::new((vector_guard.id(), i)) + }) + .collect::>(); + { + let mut vertex_tuple = VertexTuple::deserialize_mut(&mut vertex_bytes); + vertex_tuple.pointers().copy_from_slice(&pointers_t); + } + let t = { + let i = vertex_guard + .alloc(&vertex_bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + (vertex_guard.id(), i) + }; + drop(vertex_guard); + if skip < t.0 { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + *meta_tuple.skip() = (*meta_tuple.skip()).max(t.0); + } + (pointers_t, t) + }; + let start = if start.into_inner().is_none() { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + let dst = meta_tuple.start(); + if dst.into_inner().is_none() { + *dst = OptionPointer::some(t); + return; + } else { + *dst + } + } else { + start + }; + let lut = O::Vector::preprocess(vector); + let mut visited = Visited::new(); + let mut candidates = Candidates::new(beam as usize, prefetch_vectors); + let Some(s) = start.into_inner() else { + return; + }; + { + visited.insert(s); + let vertex_guard = index.read(s.0); + let Some(vertex_bytes) = vertex_guard.get(s.1) else { + // the link is broken + return; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); + let score_s = O::process( + bits, + dim, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); + candidates.push((Reverse(score_s), AlwaysEqual((pointers_s, s)))); + } + let mut iter = std::iter::from_fn(|| { + while let Some(((_, AlwaysEqual((pointers_u, u))), guards)) = candidates.pop() { + let Ok((dis_u, outs_u, _, _)) = crate::vectors::read::( + by_prefetch::(guards, pointers_u.iter().copied()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + copy_outs, + ) else { + // the link is broken + continue; + }; + let mut iterator = prefetch_vertices.prefetch( + outs_u + .into_iter() + .filter(|&x| !visited.contains(x)) + .collect::>(), + ); + while let Some((v, guards)) = iterator.next() { + visited.insert(v); + let vertex_guard = { + let mut guards = guards; + let r = guards.next().expect("internal"); + assert!(guards.next().is_none(), "internal"); + drop(guards); + r + }; + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v: &[_] = bump.alloc_slice(vertex_tuple.pointers()); + let score_v = O::process( + bits, + dim, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); + candidates.push((Reverse(score_v), AlwaysEqual((pointers_v, v)))); + } + return Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))); + } + None + }); + let mut results = Results::new(ef as _); + for element @ (Reverse(dis_c), _) in iter.by_ref() { + results.push(element); + if results + .peek_ef_th() + .map(|dis_e| dis_e < dis_c) + .unwrap_or_default() + { + break; + } + } + let trace = { + let (left, right) = results.into_inner(); + left.into_iter() + .map(|(dis_v, v)| (Reverse(dis_v), v)) + .chain(right) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let Ok((vector_u, _, _, _)) = crate::vectors::read::( + by_read(index, pointers_u.iter().copied()), + CloneAccessor::::default_with_dimension(dim), + copy_nothing, + ) else { + // the link is broken + return None; + }; + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }) + .collect::>() + }; + let outs = crate::prune::prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (bump.alloc_slice(&pointers_t) as &[_], t), + trace.into_iter(), + m, + &alpha, + |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, + ); + let _ = update::( + (index, pointers_t.as_slice()), + (version_t, VecDeque::new()), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ); + for (Reverse(dis_t), AlwaysEqual((pointers_u, u))) in outs { + 'occ: loop { + let Ok((neighbours_u, _, version)) = + crate::vectors::read_without_accessor::((index, pointers_u), copy_all) + else { + // the link is broken + break 'occ; + }; + let trace = neighbours_u + .iter() + .copied() + .chain(std::iter::once((t, dis_t))) + .map(|(v, dis_v)| (Reverse(dis_v), AlwaysEqual(v))) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual(u)) = item; + let vertex_guard = index.read(u.0); + let Some(vertex_bytes) = vertex_guard.get(u.1) else { + // the link is broken + return None; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_u = vertex_tuple.pointers().to_vec(); + Some((Reverse(dis_u), AlwaysEqual((pointers_u, u)))) + }) + .flat_map(|item| { + let (Reverse(dis_u), AlwaysEqual((pointers_u, u))) = item; + let Ok((vector_u, _, _, _)) = crate::vectors::read::( + by_read(index, pointers_u.iter().copied()), + CloneAccessor::::default_with_dimension(dim), + copy_nothing, + ) else { + // the link is broken + return None; + }; + Some(((Reverse(dis_u), AlwaysEqual((pointers_u, u))), vector_u)) + }) + .collect::>(); + let outs = crate::prune::prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (pointers_u.to_vec(), u), + trace.into_iter(), + m, + &alpha, + |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, + ); + if update::( + (index, pointers_u), + (version, neighbours_u), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ) != Ok(false) + { + break 'occ; + } + } + } +} + +fn append_vertex_tuple<'b, R: RelationRead + RelationWrite>( + index: &'b R, + first: u32, + size: usize, +) -> R::WriteGuard<'b> +where + R::Page: Page, +{ + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= size || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, true); + if write.freespace() as usize >= size { + return write; + } + if write.get_opaque().next == u32::MAX { + let link = index + .extend( + Opaque { + next: u32::MAX, + link: u32::MAX, + }, + false, + ) + .id(); + let extend = index.extend( + Opaque { + next: u32::MAX, + link, + }, + true, + ); + { write }.get_opaque_mut().next = extend.id(); + if extend.freespace() as usize >= size { + return extend; + } else { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } +} + +fn append_vector_tuple<'b, R: RelationRead + RelationWrite>( + index: &'b R, + first: u32, + size: usize, +) -> R::WriteGuard<'b> +where + R::Page: Page, +{ + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= size || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, false); + if write.freespace() as usize >= size { + return write; + } + if write.get_opaque().next == u32::MAX { + let extend = index.extend( + Opaque { + next: u32::MAX, + link: u32::MAX, + }, + false, + ); + { write }.get_opaque_mut().next = extend.id(); + if extend.freespace() as usize >= size { + return extend; + } else { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + } + current = write.get_opaque().next + } else { + current = read.get_opaque().next + } + } +} diff --git a/crates/vchordg/src/lib.rs b/crates/vchordg/src/lib.rs new file mode 100644 index 00000000..a16936c0 --- /dev/null +++ b/crates/vchordg/src/lib.rs @@ -0,0 +1,56 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod build; +mod bulkdelete; +mod candidates; +mod insert; +mod maintain; +mod prewarm; +mod prune; +mod results; +mod search; +mod tuples; +mod vectors; +mod visited; + +pub mod operator; +pub mod types; + +pub use build::build; +pub use bulkdelete::bulkdelete; +pub use insert::insert; +pub use maintain::maintain; +pub use prewarm::prewarm; +pub use search::search; + +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, FromBytes, IntoBytes, Immutable, KnownLayout)] +pub struct Opaque { + pub next: u32, + pub link: u32, +} + +#[allow(unsafe_code)] +unsafe impl index::relation::Opaque for Opaque {} + +pub type Id = (u32, u16); + +impl index::fetch::Fetch1 for tuples::Pointer { + fn fetch_1(&self) -> u32 { + self.into_inner().0 + } +} diff --git a/crates/vchordg/src/maintain.rs b/crates/vchordg/src/maintain.rs new file mode 100644 index 00000000..db3303cc --- /dev/null +++ b/crates/vchordg/src/maintain.rs @@ -0,0 +1,189 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::operator::{CloneAccessor, Operator}; +use crate::tuples::{MetaTuple, VertexTuple, WithReader}; +use crate::types::DistanceKind; +use crate::vectors::{by_read, copy_all, copy_nothing, copy_outs, update}; +use always_equal::AlwaysEqual; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::DefaultWithDimension; +use std::cmp::Reverse; +use vector::VectorOwned; + +pub fn maintain(index: &R, check: impl Fn()) +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let m = meta_tuple.m(); + let alpha = meta_tuple.alpha().to_vec(); + let start = meta_tuple.start(); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + let Some(s) = start.into_inner() else { + return; + }; + // do it's best to remove broken edges + { + let mut current = link; + while current != u32::MAX { + check(); + let vertex_guard = index.read(current); + let mut members = Vec::new(); + for i in 1..=vertex_guard.len() { + if let Some(vertex_bytes) = vertex_guard.get(i) { + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + if vertex_tuple.payload().is_some() { + let pointers_u = vertex_tuple.pointers().to_vec(); + members.push((pointers_u, (vertex_guard.id(), i))); + } + } + } + let next = { vertex_guard }.get_opaque().next; + for (pointers_u, u) in members { + 'occ: loop { + let Ok((vector_u, neighbours_u, _, version)) = crate::vectors::read::( + by_read::(index, pointers_u.iter().copied()), + CloneAccessor::::default_with_dimension(dim), + copy_all, + ) else { + // the link is broken + break 'occ; + }; + let trace = { + let mut trace = Vec::new(); + let mut extend = Vec::new(); + for &(v, dis_v) in neighbours_u.iter() { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = vertex_tuple.pointers().to_vec(); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, outs_v, _, _)) = crate::vectors::read::( + by_read::(index, pointers_v.iter().copied()), + CloneAccessor::::default_with_dimension(dim), + copy_outs, + ) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + trace.push(( + (Reverse(dis_v), AlwaysEqual((pointers_v, v))), + vector_v, + )); + } else { + extend.extend(outs_v.iter().copied()); + } + } + for v in extend { + let vertex_guard = index.read(v.0); + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = vertex_tuple.pointers().to_vec(); + let payload_v = vertex_tuple.payload(); + drop(vertex_guard); + let Ok((vector_v, _, _, _)) = crate::vectors::read::( + by_read::(index, pointers_v.iter().copied()), + CloneAccessor::::default_with_dimension(dim), + copy_nothing, + ) else { + // the link is broken + continue; + }; + if payload_v.is_some() { + let dis_v = + O::distance(vector_u.as_borrowed(), vector_v.as_borrowed()); + trace.push(( + (Reverse(dis_v), AlwaysEqual((pointers_v, v))), + vector_v, + )); + } + } + trace + }; + let outs = crate::prune::prune( + |x, y| O::distance(x.as_borrowed(), y.as_borrowed()), + (pointers_u.to_vec(), u), + trace.into_iter(), + m, + &alpha, + |(_, u)| *u, + O::DISTANCE == DistanceKind::L2S, + ); + if update::( + (index, pointers_u.as_slice()), + (version, neighbours_u), + outs.iter() + .map(|&(Reverse(dis_u), AlwaysEqual((_, u)))| (u, dis_u)), + ) != Ok(false) + { + break 'occ; + } + } + } + current = next; + } + } + // remove vertices and vectors + { + let mut current = link; + while current != u32::MAX { + check(); + let mut vertex_guard = index.write(current, true); + let mut reachable_set = Vec::<(u32, u16)>::new(); + for i in 1..=vertex_guard.len() { + if let Some(bytes) = vertex_guard.get(i) { + let tuple = VertexTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if p.is_none() && (current, i) != s { + vertex_guard.free(i); + } else { + let iter = tuple.pointers().iter().map(|pointer| pointer.into_inner()); + reachable_set.extend(iter); + } + } + } + reachable_set.sort_unstable(); + { + let mut current = vertex_guard.get_opaque().link; + while current != u32::MAX { + check(); + let mut vector_guard = index.write(current, false); + for i in 1..=vector_guard.len() { + if vector_guard.get(i).is_some() { + if reachable_set.binary_search(&(current, i)).is_err() { + vector_guard.free(i); + } + } + } + current = vector_guard.get_opaque().next; + } + } + current = vertex_guard.get_opaque().next; + } + } +} diff --git a/crates/vchordg/src/operator.rs b/crates/vchordg/src/operator.rs new file mode 100644 index 00000000..d5961c7d --- /dev/null +++ b/crates/vchordg/src/operator.rs @@ -0,0 +1,572 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::types::DistanceKind; +use distance::Distance; +use index_accessor::{ + Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, + HalfbyteDistanceAccessor, L2S, +}; +use rabitq::bits::Bits; +use simd::{Floating, f16}; +use std::fmt::Debug; +use std::marker::PhantomData; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub trait Vector: VectorOwned { + type Element: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + type Metadata: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)); + fn pack(dim: u32, elements: Vec, metadata: Self::Metadata) -> Self; + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code; + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut; +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f32; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let slice = vector.slice(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= (8000 - tailing) / size_of::() { + return (vec![], (slice, ())); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); + ( + l.chunks(8000 / size_of::()).collect::>(), + (r, ()), + ) + } + + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + rabitq::bits::code(bits, vector.slice()) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + rabitq::bits::binary::preprocess(vector.slice()) + } +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f16; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let slice = vector.slice(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= (8000 - tailing) / size_of::() { + return (vec![], (slice, ())); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing) / size_of::()); + ( + l.chunks(8000 / size_of::()).collect::>(), + (r, ()), + ) + } + + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + rabitq::bits::code(bits, &f16::vector_to_f32(vector.slice())) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + rabitq::bits::binary::preprocess(&f16::vector_to_f32(vector.slice())) + } +} + +impl Vector for Rabitq8Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let metadata = [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ]; + let slice = vector.packed_code(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= 8000 - tailing { + return (vec![], (slice, metadata)); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing)); + (l.chunks(8000).collect::>(), (r, metadata)) + } + + fn pack(dim: u32, code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(dim, _0, _1, _2, _3, code) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::code(bits, &result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::binary::preprocess(&result) + } +} + +impl Vector for Rabitq4Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn split( + vector: Self::Borrowed<'_>, + m: usize, + ) -> (Vec<&[Self::Element]>, (&[Self::Element], Self::Metadata)) { + let metadata = [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ]; + let slice = vector.packed_code(); + let tailing = (size_of::() * m) + .next_multiple_of(crate::tuples::ALIGN); + assert!(tailing <= 8000); + if slice.len() <= 8000 - tailing { + return (vec![], (slice, metadata)); + } + let (l, r) = slice.split_at(slice.len() - (8000 - tailing)); + (l.chunks(8000).collect::>(), (r, metadata)) + } + + fn pack(dim: u32, code: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq4Owned::new(dim, _0, _1, _2, _3, code) + } + + fn code(bits: Bits, vector: Self::Borrowed<'_>) -> rabitq::bits::Code { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::code(bits, &result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> rabitq::bits::binary::BinaryLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bits::binary::preprocess(&result) + } +} + +pub trait Operator: 'static + Debug + Copy { + const DISTANCE: DistanceKind; + + type Vector: Vector; + + type DistanceAccessor: DefaultWithDimension + + Accessor2< + ::Element, + ::Element, + ::Metadata, + ::Metadata, + Output = Distance, + >; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance; + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance; +} + +#[derive(Debug)] +pub struct Op(PhantomData V>, PhantomData D>); + +impl Clone for Op { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Op {} + +impl Operator for Op, L2S> { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op, Dot> { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +impl Operator for Op, L2S> { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op, Dot> { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = Rabitq8Owned; + + type DistanceAccessor = ByteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = Rabitq8Owned; + + type DistanceAccessor = ByteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::L2S; + + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_l2s( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_l2s(rhs) + } +} + +impl Operator for Op { + const DISTANCE: DistanceKind = DistanceKind::Dot; + + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn process( + bits: Bits, + dim: u32, + code: ([f32; 3], &[u64]), + lut: &rabitq::bits::binary::BinaryLut, + ) -> Distance { + use rabitq::bits::CodeMetadata; + let sum = rabitq::bits::binary::accumulate(bits, code.1, &lut.1); + let (distance,) = rabitq::bits::binary::half_process_dot( + bits, + dim, + sum, + CodeMetadata::from_array(code.0), + lut.0, + ); + Distance::from_f32(distance) + } + + fn distance( + lhs: ::Borrowed<'_>, + rhs: ::Borrowed<'_>, + ) -> Distance { + lhs.operator_dot(rhs) + } +} + +#[derive(Debug, Clone)] +pub struct CloneAccessor(u32, Vec); + +impl DefaultWithDimension for CloneAccessor { + fn default_with_dimension(dim: u32) -> Self { + Self(dim, Vec::new()) + } +} + +impl Accessor1 for CloneAccessor { + type Output = V; + + fn push(&mut self, input: &[V::Element]) { + self.1.extend(input); + } + + fn finish(self, metadata: V::Metadata) -> Self::Output { + V::pack(self.0, self.1, metadata) + } +} diff --git a/crates/vchordg/src/prewarm.rs b/crates/vchordg/src/prewarm.rs new file mode 100644 index 00000000..9ccc076f --- /dev/null +++ b/crates/vchordg/src/prewarm.rs @@ -0,0 +1,46 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::operator::Operator; +use index::relation::{Page, RelationRead}; + +pub fn prewarm(index: &R) -> String +where + R::Page: Page, +{ + use std::fmt::Write; + let meta_guard = index.read(0); + let link = meta_guard.get_opaque().link; + drop(meta_guard); + let mut number_of_vertex_pages = 0_u64; + let mut number_of_vertices = 0_u64; + { + let mut current = link; + while current != u32::MAX { + let vertex_guard = index.read(current); + number_of_vertex_pages += 1; + for i in 1..=vertex_guard.len() { + if vertex_guard.get(i).is_some() { + number_of_vertices += 1; + } + } + current = vertex_guard.get_opaque().next; + } + } + let mut message = String::new(); + writeln!(message, "number of vertex pages: {number_of_vertex_pages}").unwrap(); + writeln!(message, "number of vertices: {number_of_vertices}").unwrap(); + message +} diff --git a/crates/vchordg/src/prune.rs b/crates/vchordg/src/prune.rs new file mode 100644 index 00000000..138ebc1f --- /dev/null +++ b/crates/vchordg/src/prune.rs @@ -0,0 +1,72 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; + +pub fn prune( + mut d: impl FnMut(&V, &V) -> Distance, + u: T, + outs: impl Iterator, AlwaysEqual), V)>, + m: u32, + alpha: &[f32], + key: impl Fn(&T) -> K, + is_l2s: bool, +) -> Vec<(Reverse, AlwaysEqual)> { + // V ← (V ∪ Nout(p)) \ {p} + let mut trace = outs.collect::>(); + trace.sort_by_key(|((_, AlwaysEqual(v)), _)| key(v)); + trace.dedup_by_key(|((_, AlwaysEqual(v)), _)| key(v)); + trace.retain(|((_, AlwaysEqual(v)), _)| key(v) != key(&u)); + trace.sort_by_key(|&((Reverse(d), _), _)| d); + // Nout(p) ← ∅ + let mut result = Vec::new(); + for &alpha in if is_l2s { alpha } else { &[1.0] } { + if result.len() == m as usize { + break; + } + trace = robust_prune(&mut d, m, alpha, &mut result, trace); + } + if !(result.len() == m as usize) { + result.extend(trace); + result.truncate(m as usize); + } + result.into_iter().map(|(x, _)| x).collect() +} + +fn robust_prune( + mut d: impl FnMut(&V, &V) -> Distance, + m: u32, + alpha: f32, + result: &mut Vec<((Reverse, AlwaysEqual), V)>, + trace: Vec<((Reverse, AlwaysEqual), V)>, +) -> Vec<((Reverse, AlwaysEqual), V)> { + let mut pruned = Vec::new(); + for ((Reverse(dis_u), AlwaysEqual(u)), vector_u) in trace { + if result.len() == m as usize { + break; + } + let check = result + .iter() + .map(|(_, vector_v)| d(&vector_u, vector_v)) + .all(|dis| dis_u.to_f32() < alpha * dis.to_f32()); + if check { + result.push(((Reverse(dis_u), AlwaysEqual(u)), vector_u)); + } else { + pruned.push(((Reverse(dis_u), AlwaysEqual(u)), vector_u)); + } + } + pruned +} diff --git a/crates/vchordg/src/results.rs b/crates/vchordg/src/results.rs new file mode 100644 index 00000000..8a31396a --- /dev/null +++ b/crates/vchordg/src/results.rs @@ -0,0 +1,73 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use always_equal::AlwaysEqual; +use distance::Distance; +use min_max_heap::MinMaxHeap; +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +pub struct Results { + ef: usize, + front: MinMaxHeap<(Distance, AlwaysEqual)>, + back: BinaryHeap<(Reverse, AlwaysEqual)>, + // len(front) < ef, then len(back) == 0 + // len(front) == ef, then max(front) < min(back) +} + +impl Results { + pub fn new(ef: usize) -> Self { + assert!(ef > 0, "ef must be positive integer"); + Self { + ef, + front: MinMaxHeap::with_capacity(ef), + back: BinaryHeap::new(), + } + } + pub fn push(&mut self, item: (Reverse, AlwaysEqual)) { + if self.front.len() < self.ef { + self.front.push((item.0.0, item.1)); + } else { + let item = self.front.push_pop_max((item.0.0, item.1)); + self.back.push((Reverse(item.0), item.1)); + } + } + pub fn peek_ef_th(&mut self) -> Option { + if self.front.len() < self.ef { + None + } else { + self.front.peek_max().map(|&(d, _)| d) + } + } + #[allow(clippy::collapsible_else_if)] + pub fn pop_min(&mut self) -> Option<(Distance, AlwaysEqual)> { + if self.front.len() < self.ef { + self.front.pop_min() + } else { + if let Some(item) = self.back.pop() { + Some(self.front.push_pop_min((item.0.0, item.1))) + } else { + self.front.pop_min() + } + } + } + pub fn into_inner( + self, + ) -> ( + Vec<(Distance, AlwaysEqual)>, + Vec<(Reverse, AlwaysEqual)>, + ) { + (self.front.into_vec(), self.back.into_vec()) + } +} diff --git a/crates/vchordg/src/search.rs b/crates/vchordg/src/search.rs new file mode 100644 index 00000000..23a6230d --- /dev/null +++ b/crates/vchordg/src/search.rs @@ -0,0 +1,140 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::candidates::Candidates; +use crate::operator::{Operator, Vector}; +use crate::results::Results; +use crate::tuples::*; +use crate::vectors::{by_prefetch, copy_outs}; +use crate::visited::Visited; +use always_equal::AlwaysEqual; +use distance::Distance; +use index::bump::Bump; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, RelationRead}; +use index_accessor::{DefaultWithDimension, LAccess}; +use rabitq::bits::Bits; +use std::cmp::Reverse; +use std::collections::VecDeque; +use std::num::NonZero; +use vector::{VectorBorrowed, VectorOwned}; + +pub fn search<'b, R: RelationRead, O: Operator>( + index: &'b R, + vector: ::Borrowed<'b>, + ef_search: u32, + beam_search: u32, + bump: &'b impl Bump, + mut prefetch_vertices: impl PrefetcherSequenceFamily<'b, R> + 'b, + prefetch_vectors: impl PrefetcherSequenceFamily<'b, R> + 'b, +) -> Box)> + 'b> +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let start = meta_tuple.start(); + let bits = Bits::try_from(meta_tuple.bits()).expect("data corruption"); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + let ef = ef_search; + let beam = beam_search; + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + let mut visited = Visited::new(); + let mut candidates = Candidates::new(beam as usize, prefetch_vectors); + let Some(s) = start.into_inner() else { + return Box::new(std::iter::empty()); + }; + { + visited.insert(s); + let vertex_guard = index.read(s.0); + let Some(vertex_bytes) = vertex_guard.get(s.1) else { + // the link is broken + return Box::new(std::iter::empty()); + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_s: &[_] = bump.alloc_slice(vertex_tuple.pointers()); + let score_s = O::process( + bits, + dim, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); + candidates.push((Reverse(score_s), AlwaysEqual(pointers_s))); + } + let mut iter = std::iter::from_fn(move || { + while let Some(((_, AlwaysEqual(pointers_u)), guards)) = candidates.pop() { + let Ok((dis_u, outs_u, payload_u, _)) = crate::vectors::read::( + by_prefetch::(guards, pointers_u.iter().copied()), + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + copy_outs, + ) else { + // the link is broken + continue; + }; + let mut iterator = prefetch_vertices.prefetch( + outs_u + .into_iter() + .filter(|&x| !visited.contains(x)) + .collect::>(), + ); + while let Some((v, guards)) = iterator.next() { + visited.insert(v); + let vertex_guard = { + let mut guards = guards; + let r = guards.next().expect("internal"); + assert!(guards.next().is_none(), "internal"); + drop(guards); + r + }; + let Some(vertex_bytes) = vertex_guard.get(v.1) else { + // the link is broken + continue; + }; + let vertex_tuple = VertexTuple::deserialize_ref(vertex_bytes); + let pointers_v = bump.alloc_slice(vertex_tuple.pointers()); + let score_v = O::process( + bits, + dim, + (vertex_tuple.metadata(), vertex_tuple.elements()), + &lut, + ); + candidates.push((Reverse(score_v), AlwaysEqual(pointers_v))); + } + return Some((Reverse(dis_u), AlwaysEqual(payload_u))); + } + None + }); + let mut results = Results::new(ef as _); + let search = std::iter::from_fn(move || { + for element @ (Reverse(dis_c), _) in iter.by_ref() { + results.push(element); + if results + .peek_ef_th() + .map(|dis_e| dis_e < dis_c) + .unwrap_or_default() + { + break; + } + } + results.pop_min() + }); + Box::new(search.filter_map(|(dis_u, AlwaysEqual(payload_u))| Some((dis_u, payload_u?)))) +} diff --git a/crates/vchordg/src/tuples.rs b/crates/vchordg/src/tuples.rs new file mode 100644 index 00000000..42efddf8 --- /dev/null +++ b/crates/vchordg/src/tuples.rs @@ -0,0 +1,764 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::Vector; +use distance::Distance; +use index::tuples::{Bool, MutChecker, Padding, RefChecker}; +use std::num::{NonZero, Wrapping}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; +pub type Tag = u64; +const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordg\0"); +const VERSION: u64 = 1001; + +#[inline(always)] +fn tag(source: &[u8]) -> Tag { + assert!(source.len() >= size_of::()); + #[allow(unsafe_code)] + unsafe { + source.as_ptr().cast::().read_unaligned() + } +} + +pub trait Tuple: 'static { + fn serialize(&self) -> Vec; +} + +pub trait WithReader: Tuple { + type Reader<'a>; + fn deserialize_ref(source: &[u8]) -> Self::Reader<'_>; +} + +pub trait WithWriter: Tuple { + type Writer<'a>; + fn deserialize_mut(source: &mut [u8]) -> Self::Writer<'_>; +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct MetaTupleHeader { + version: u64, + dim: u32, + bits: u8, + _padding: [Padding; 7], + m: u32, + alpha_s: u16, + alpha_e: u16, + ef_construction: u32, + beam_construction: u32, + start: OptionPointer, + skip: u32, +} + +pub struct MetaTuple { + pub dim: u32, + pub bits: u8, + pub m: u32, + pub alpha: Vec, + pub ef_construction: u32, + pub beam_construction: u32, + pub start: OptionPointer, + pub skip: u32, +} + +impl Tuple for MetaTuple { + #[allow(clippy::match_single_binding)] + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + MetaTuple { + dim, + bits, + m, + alpha, + ef_construction, + beam_construction, + start, + skip, + } => { + buffer.extend((MAGIC as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let alpha_s = buffer.len() as u16; + buffer.extend(alpha.as_bytes()); + let alpha_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + MetaTupleHeader { + version: VERSION, + dim: *dim, + bits: *bits, + m: *m, + alpha_s, + alpha_e, + ef_construction: *ef_construction, + beam_construction: *beam_construction, + start: *start, + skip: *skip, + _padding: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for MetaTuple { + type Reader<'a> = MetaTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!("deserialization: bad version number"); + } + let header: &MetaTupleHeader = checker.prefix(size_of::()); + let alpha = checker.bytes(header.alpha_s, header.alpha_e); + MetaTupleReader { header, alpha } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct MetaTupleReader<'a> { + header: &'a MetaTupleHeader, + alpha: &'a [f32], +} + +impl<'a> MetaTupleReader<'a> { + pub fn dim(self) -> u32 { + self.header.dim + } + pub fn bits(self) -> u8 { + self.header.bits + } + pub fn m(self) -> u32 { + self.header.m + } + pub fn alpha(self) -> &'a [f32] { + self.alpha + } + pub fn ef_construction(self) -> u32 { + self.header.ef_construction + } + pub fn beam_construction(self) -> u32 { + self.header.beam_construction + } + pub fn start(self) -> OptionPointer { + self.header.start + } + pub fn skip(self) -> u32 { + self.header.skip + } +} + +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!("deserialization: bad version number"); + } + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl<'a> MetaTupleWriter<'a> { + pub fn start(&mut self) -> &mut OptionPointer { + &mut self.header.start + } + pub fn skip(&mut self) -> &mut u32 { + &mut self.header.skip + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VertexTupleHeader { + metadata: [f32; 3], + elements_s: u16, + elements_e: u16, + payload: Option>, + pointers_s: u16, + pointers_e: u16, + _padding_0: [Padding; 4], +} + +#[derive(Debug, Clone)] +pub struct VertexTuple { + pub metadata: [f32; 3], + pub elements: Vec, + pub payload: Option>, + pub pointers: Vec, +} + +impl Tuple for VertexTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(self.elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // pointers + let pointers_s = buffer.len() as u16; + buffer.extend(self.pointers.as_bytes()); + let pointers_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[..size_of::()].copy_from_slice( + VertexTupleHeader { + metadata: self.metadata, + elements_s, + elements_e, + payload: self.payload, + pointers_s, + pointers_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + } +} + +impl WithReader for VertexTuple { + type Reader<'a> = VertexTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> VertexTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &VertexTupleHeader = checker.prefix(0_u16); + let elements = checker.bytes(header.elements_s, header.elements_e); + let pointers = checker.bytes(header.pointers_s, header.pointers_e); + VertexTupleReader { + header, + elements, + pointers, + } + } +} + +impl WithWriter for VertexTuple { + type Writer<'a> = VertexTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> VertexTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut VertexTupleHeader = checker.prefix(0_u16); + let elements = checker.bytes(header.elements_s, header.elements_e); + let pointers = checker.bytes(header.pointers_s, header.pointers_e); + VertexTupleWriter { + header, + elements, + pointers, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct VertexTupleReader<'a> { + header: &'a VertexTupleHeader, + elements: &'a [u64], + pointers: &'a [Pointer], +} + +impl<'a> VertexTupleReader<'a> { + pub fn metadata(self) -> [f32; 3] { + self.header.metadata + } + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [u64] { + self.elements + } + pub fn pointers(self) -> &'a [Pointer] { + self.pointers + } +} + +#[derive(Debug)] +pub struct VertexTupleWriter<'a> { + header: &'a mut VertexTupleHeader, + elements: &'a mut [u64], + pointers: &'a mut [Pointer], +} + +impl VertexTupleWriter<'_> { + #[expect(dead_code)] + pub fn metadata(&mut self) -> &mut [f32; 3] { + &mut self.header.metadata + } + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [u64] { + self.elements + } + pub fn pointers(&mut self) -> &mut [Pointer] { + self.pointers + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader0 { + payload: Option>, + elements_s: u16, + elements_e: u16, + metadata_s: u16, + neighbours_s: u16, + neighbours_e: u16, + _padding_0: [Padding; 2], + version: Wrapping, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader1 { + payload: Option>, + elements_s: u16, + elements_e: u16, + index: u32, +} + +#[derive(Debug, Clone)] +pub enum VectorTuple { + _0 { + payload: Option>, + metadata: V::Metadata, + elements: Vec, + neighbours: Vec, + version: Wrapping, + }, + _1 { + payload: Option>, + elements: Vec, + index: u32, + }, +} + +impl Tuple for VectorTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + VectorTuple::_0 { + payload, + metadata, + elements, + neighbours, + version, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // metadata + let metadata_s = buffer.len() as u16; + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // neighbours + let neighbours_s = buffer.len() as u16; + buffer.extend(neighbours.as_bytes()); + let neighbours_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader0 { + metadata_s, + elements_s, + elements_e, + neighbours_s, + neighbours_e, + payload: *payload, + version: *version, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + VectorTuple::_1 { + payload, + elements, + index, + } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader1 { + elements_s, + elements_e, + payload: *payload, + index: *index, + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for VectorTuple { + type Reader<'a> = VectorTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + let metadata = checker.prefix(header.metadata_s); + let neighbours = checker.bytes(header.neighbours_s, header.neighbours_e); + VectorTupleReader::_0(VectorTupleReader0 { + header, + elements, + metadata, + neighbours, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_1(VectorTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +impl WithWriter for VectorTuple { + type Writer<'a> = VectorTupleWriter<'a, V>; + + fn deserialize_mut(source: &mut [u8]) -> VectorTupleWriter<'_, V> { + let tag = tag(source); + match tag { + 0 => { + let mut checker = MutChecker::new(source); + let header: &mut VectorTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + let metadata = checker.prefix(header.metadata_s); + let neighbours = checker.bytes(header.neighbours_s, header.neighbours_e); + VectorTupleWriter::_0(VectorTupleWriter0 { + header, + elements, + metadata, + neighbours, + }) + } + 1 => { + let mut checker = MutChecker::new(source); + let header: &mut VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleWriter::_1(VectorTupleWriter1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug)] +pub enum VectorTupleReader<'a, V: Vector> { + _0(VectorTupleReader0<'a, V>), + _1(VectorTupleReader1<'a, V>), +} + +impl<'a, V: Vector> Copy for VectorTupleReader<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +#[derive(Debug)] +pub struct VectorTupleReader0<'a, V: Vector> { + header: &'a VectorTupleHeader0, + elements: &'a [V::Element], + metadata: &'a V::Metadata, + neighbours: &'a [OptionNeighbour], +} + +impl<'a, V: Vector> Copy for VectorTupleReader0<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader0<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +impl<'a, V: Vector> VectorTupleReader0<'a, V> { + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [V::Element] { + self.elements + } + pub fn metadata(self) -> &'a V::Metadata { + self.metadata + } + pub fn neighbours(self) -> &'a [OptionNeighbour] { + self.neighbours + } + pub fn version(self) -> Wrapping { + self.header.version + } +} + +#[derive(Debug)] +pub struct VectorTupleReader1<'a, V: Vector> { + header: &'a VectorTupleHeader1, + elements: &'a [V::Element], +} + +impl<'a, V: Vector> Copy for VectorTupleReader1<'a, V> {} + +impl<'a, V: Vector> Clone for VectorTupleReader1<'a, V> { + fn clone(&self) -> Self { + *self + } +} + +impl<'a, V: Vector> VectorTupleReader1<'a, V> { + #[allow(dead_code)] + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [V::Element] { + self.elements + } + pub fn index(self) -> u32 { + self.header.index + } +} + +#[derive(Debug)] +pub enum VectorTupleWriter<'a, V: Vector> { + _0(VectorTupleWriter0<'a, V>), + #[expect(dead_code)] + _1(VectorTupleWriter1<'a, V>), +} + +#[derive(Debug)] +pub struct VectorTupleWriter0<'a, V: Vector> { + header: &'a mut VectorTupleHeader0, + elements: &'a mut [V::Element], + metadata: &'a mut V::Metadata, + neighbours: &'a mut [OptionNeighbour], +} + +impl VectorTupleWriter0<'_, V> { + #[expect(dead_code)] + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [V::Element] { + self.elements + } + #[expect(dead_code)] + pub fn metadata(&mut self) -> &mut V::Metadata { + self.metadata + } + pub fn neighbours(&mut self) -> &mut [OptionNeighbour] { + self.neighbours + } + pub fn version(&mut self) -> &mut Wrapping { + &mut self.header.version + } +} + +#[derive(Debug)] +pub struct VectorTupleWriter1<'a, V: Vector> { + header: &'a mut VectorTupleHeader1, + elements: &'a mut [V::Element], +} + +impl VectorTupleWriter1<'_, V> { + #[expect(dead_code)] + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } + #[expect(dead_code)] + pub fn elements(&mut self) -> &mut [V::Element] { + self.elements + } + #[expect(dead_code)] + pub fn index(&mut self) -> &mut u32 { + &mut self.header.index + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct Pointer { + x: u32, + y: u16, + _padding_0: [Padding; 2], +} + +impl Pointer { + pub fn new((x, y): (u32, u16)) -> Self { + Self { + x, + y, + _padding_0: Default::default(), + } + } + pub fn into_inner(self) -> (u32, u16) { + (self.x, self.y) + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct OptionPointer { + x: u32, + y: u16, + _padding_0: [Padding; 1], + validity: Bool, +} + +impl OptionPointer { + pub const NONE: Self = Self { + x: 0, + y: 0, + _padding_0: [Padding::ZERO; 1], + validity: Bool::FALSE, + }; + pub fn some((x, y): (u32, u16)) -> Self { + Self { + x, + y, + _padding_0: Default::default(), + validity: Bool::TRUE, + } + } + pub fn into_inner(self) -> Option<(u32, u16)> { + if self.validity.into() { + Some((self.x, self.y)) + } else { + None + } + } +} + +#[repr(C)] +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + IntoBytes, + FromBytes, + Immutable, + KnownLayout, +)] +pub struct OptionNeighbour { + pointer: OptionPointer, + distance: Distance, +} + +impl OptionNeighbour { + pub const NONE: Self = Self { + pointer: OptionPointer::NONE, + distance: Distance::ZERO, + }; + pub fn some(pointer: (u32, u16), distance: Distance) -> Self { + Self { + pointer: OptionPointer::some(pointer), + distance, + } + } + pub fn into_inner(self) -> Option<((u32, u16), Distance)> { + let inner = self.pointer.into_inner()?; + Some((inner, self.distance)) + } +} diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs new file mode 100644 index 00000000..5a5a6169 --- /dev/null +++ b/crates/vchordg/src/types.rs @@ -0,0 +1,168 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use serde::{Deserialize, Serialize}; +use simd::f16; +use validator::{Validate, ValidationError}; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; +use vector::vect::{VectBorrowed, VectOwned}; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordgIndexOptions { + #[serde(default = "VchordgIndexOptions::default_bits")] + #[validate(range(min = 1, max = 2))] + pub bits: u8, + #[serde(default = "VchordgIndexOptions::default_m")] + #[validate(range(min = 1, max = 512))] + pub m: u32, + #[serde(default = "VchordgIndexOptions::default_alpha")] + #[validate(length(min = 1, max = 8), custom(function = VchordgIndexOptions::validate_alpha))] + pub alpha: Vec, + #[serde(default = "VchordgIndexOptions::default_ef_construction")] + #[validate(range(min = 1, max = 65535))] + pub ef_construction: u32, + #[serde(default = "VchordgIndexOptions::default_beam_construction")] + #[validate(range(min = 1, max = 65535))] + pub beam_construction: u32, +} + +impl VchordgIndexOptions { + fn default_bits() -> u8 { + 2 + } + fn default_m() -> u32 { + 32 + } + fn default_alpha() -> Vec { + vec![1.0, 1.2] + } + fn validate_alpha(alpha: &[f32]) -> Result<(), ValidationError> { + if !alpha.is_sorted() { + return Err(ValidationError::new("`alpha` should be in ascending order")); + } + if !alpha.iter().all(|x| (1.0..2.0).contains(x)) { + return Err(ValidationError::new("alpha is too large or too small")); + } + if alpha.first() != Some(&1.0) { + return Err(ValidationError::new("`alpha` should contain `1.0`")); + } + Ok(()) + } + fn default_ef_construction() -> u32 { + 64 + } + fn default_beam_construction() -> u32 { + 1 + } +} + +impl Default for VchordgIndexOptions { + fn default() -> Self { + Self { + bits: Self::default_bits(), + m: Self::default_m(), + alpha: Self::default_alpha(), + ef_construction: Self::default_ef_construction(), + beam_construction: Self::default_beam_construction(), + } + } +} + +#[derive(Debug, Clone)] +pub enum OwnedVector { + Vecf32(VectOwned), + Vecf16(VectOwned), + Rabitq8(Rabitq8Owned), + Rabitq4(Rabitq4Owned), +} + +#[derive(Debug, Clone, Copy)] +pub enum BorrowedVector<'a> { + Vecf32(VectBorrowed<'a, f32>), + Vecf16(VectBorrowed<'a, f16>), + Rabitq8(Rabitq8Borrowed<'a>), + Rabitq4(Rabitq4Borrowed<'a>), +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum DistanceKind { + L2S, + Dot, +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum VectorKind { + Vecf32, + Vecf16, + Rabitq8, + Rabitq4, +} + +#[derive(Debug, Clone, Validate)] +#[validate(schema(function = "Self::validate_self"))] +pub struct VectorOptions { + #[validate(range(min = 1))] + pub dim: u32, + pub v: VectorKind, + pub d: DistanceKind, +} + +impl VectorOptions { + pub fn validate_self(&self) -> Result<(), ValidationError> { + match (self.v, self.d, self.dim) { + (_, _, 1..=30000) => Ok(()), + _ => Err(ValidationError::new("invalid vector options")), + } + } +} + +pub struct Structure { + pub centroids: Vec, + pub children: Vec>, +} + +impl Structure { + pub fn len(&self) -> usize { + self.children.len() + } + pub fn is_empty(&self) -> bool { + self.children.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::VchordgIndexOptions; + + #[test] + fn validate_alpha_handles_empty_without_panicking() { + let empty: &[f32] = &[]; + assert!(VchordgIndexOptions::validate_alpha(empty).is_err()); + } + + #[test] + fn validate_alpha_accepts_default() { + assert!(VchordgIndexOptions::validate_alpha(&[1.0, 1.2]).is_ok()); + } + + #[test] + fn validate_alpha_rejects_unsorted_or_out_of_range() { + assert!(VchordgIndexOptions::validate_alpha(&[1.2, 1.0]).is_err()); + assert!(VchordgIndexOptions::validate_alpha(&[2.0]).is_err()); + } +} diff --git a/crates/vchordg/src/vectors.rs b/crates/vchordg/src/vectors.rs new file mode 100644 index 00000000..e2933fba --- /dev/null +++ b/crates/vchordg/src/vectors.rs @@ -0,0 +1,170 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::{Operator, Vector}; +use crate::tuples::*; +use distance::Distance; +use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::Accessor1; +use std::collections::VecDeque; +use std::num::{NonZero, Wrapping}; + +pub fn by_prefetch<'b, R: RelationRead>( + guards: impl ExactSizeIterator>, + pointers_u: impl ExactSizeIterator, +) -> impl ExactSizeIterator, u16)> { + assert!(guards.len() == pointers_u.len() && pointers_u.len() > 0); + pointers_u + .map(Pointer::into_inner) + .zip(guards) + .map(|(pointer, guard)| (guard, pointer.1)) +} + +pub fn by_read<'b, R: RelationRead>( + index: &'b R, + pointers_u: impl ExactSizeIterator, +) -> impl ExactSizeIterator, u16)> { + assert!(pointers_u.len() > 0); + pointers_u + .map(Pointer::into_inner) + .map(|pointer| (pointer, index.read(pointer.0))) + .map(|(pointer, guard)| (guard, pointer.1)) +} + +pub fn copy_nothing(_: &[OptionNeighbour]) {} + +pub fn copy_outs(x: &[OptionNeighbour]) -> VecDeque<(u32, u16)> { + x.iter() + .flat_map(|neighbour| neighbour.into_inner()) + .map(|(v, _)| v) + .collect::>() +} + +pub fn copy_all(x: &[OptionNeighbour]) -> VecDeque<((u32, u16), Distance)> { + x.iter() + .flat_map(|neighbour| neighbour.into_inner()) + .collect::>() +} + +pub fn read< + 'b, + R: RelationRead, + O: Operator, + A: Accessor1<::Element, ::Metadata>, + Output, +>( + mut iterator: impl ExactSizeIterator, u16)>, + accessor: A, + copy: impl FnOnce(&[OptionNeighbour]) -> Output, +) -> Result<(A::Output, Output, Option>, Wrapping), ()> { + let m = iterator.len().strict_sub(1); + let mut result = accessor; + for index in 0..m { + let (vector_guard, i) = iterator.next().expect("internal: bad size"); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_1(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if vector_tuple.index() as usize != index { + // the link is broken + return Err(()); + } + result.push(vector_tuple.elements()); + } + let value_u; + let payload_u; + let neighbours_u; + let version_u; + { + let (vector_guard, i) = iterator.next().expect("internal: bad size"); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + result.push(vector_tuple.elements()); + value_u = result.finish(*vector_tuple.metadata()); + neighbours_u = copy(vector_tuple.neighbours()); + payload_u = vector_tuple.payload(); + version_u = vector_tuple.version(); + } + Ok((value_u, neighbours_u, payload_u, version_u)) +} + +pub fn read_without_accessor( + (index, pointers_u): (&R, &[Pointer]), + copy: impl FnOnce(&[OptionNeighbour]) -> Output, +) -> Result<(Output, Option>, Wrapping), ()> { + let payload_u; + let neighbours_u; + let version_u; + { + let (id, i) = pointers_u.last().expect("internal: bad size").into_inner(); + let vector_guard = index.read(id); + let Some(vector_bytes) = vector_guard.get(i) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_ref(vector_bytes); + let VectorTupleReader::_0(vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + neighbours_u = copy(vector_tuple.neighbours()); + payload_u = vector_tuple.payload(); + version_u = vector_tuple.version(); + } + Ok((neighbours_u, payload_u, version_u)) +} + +pub fn update( + (index, pointers_u): (&R, &[Pointer]), + (version, neighbours_u): (Wrapping, VecDeque<((u32, u16), Distance)>), + outs: impl Iterator + Clone, +) -> Result { + if outs.clone().eq(neighbours_u) { + return Ok(true); + } + let mut vector_guard = index.write(pointers_u.last().unwrap().into_inner().0, false); + let Some(vector_bytes) = vector_guard.get_mut(pointers_u.last().unwrap().into_inner().1) else { + // the link is broken + return Err(()); + }; + let vector_tuple = VectorTuple::::deserialize_mut(vector_bytes); + let VectorTupleWriter::_0(mut vector_tuple) = vector_tuple else { + // the link is broken + return Err(()); + }; + if *vector_tuple.version() != version { + return Ok(false); + } else { + *vector_tuple.version() += 1; + } + let filling = outs + .map(|(v, dis_v)| OptionNeighbour::some(v, dis_v)) + .chain(std::iter::repeat(OptionNeighbour::NONE)); + for (hole, fill) in std::iter::zip(vector_tuple.neighbours().iter_mut(), filling) { + *hole = fill; + } + Ok(true) +} diff --git a/crates/vchordg/src/visited.rs b/crates/vchordg/src/visited.rs new file mode 100644 index 00000000..8599dff2 --- /dev/null +++ b/crates/vchordg/src/visited.rs @@ -0,0 +1,34 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Id; +use std::collections::HashSet; + +pub struct Visited { + inner: HashSet, +} + +impl Visited { + pub fn new() -> Self { + Self { + inner: HashSet::new(), + } + } + pub fn insert(&mut self, x: Id) { + self.inner.insert(x); + } + pub fn contains(&mut self, x: Id) -> bool { + self.inner.contains(&x) + } +} diff --git a/crates/vchordrq/Cargo.toml b/crates/vchordrq/Cargo.toml new file mode 100644 index 00000000..531118a4 --- /dev/null +++ b/crates/vchordrq/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "vchordrq" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +always_equal = { path = "../always_equal" } +distance = { path = "../distance" } +index = { path = "../index" } +index_accessor = { path = "../index_accessor" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } +vector = { path = "../vector" } + +pin-project = "1.1.11" +serde.workspace = true +validator.workspace = true +zerocopy.workspace = true + +[dev-dependencies] +rand.workspace = true + +[lints] +workspace = true diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs new file mode 100644 index 00000000..d622fbc6 --- /dev/null +++ b/crates/vchordrq/src/build.rs @@ -0,0 +1,151 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::{Operator, Vector}; +use crate::tape::TapeWriter; +use crate::tape_writer::{DirectoryTapeWriter, H1TapeWriter}; +use crate::tuples::*; +use crate::types::*; +use crate::{Branch, Opaque}; +use index::relation::{Page, RelationWrite}; +use vector::{VectorBorrowed, VectorOwned}; + +pub fn build( + vector_options: VectorOptions, + vchordrq_options: VchordrqIndexOptions, + index: &R, + structures: Vec>, +) where + R::Page: Page, +{ + let dim = vector_options.dim; + let is_residual = vchordrq_options.residual_quantization; + let mut meta = TapeWriter::<_, MetaTuple>::create(index, false); + assert_eq!(meta.first(), 0); + let mut freepages = TapeWriter::<_, FreepagesTuple>::create(index, false); + freepages.push(FreepagesTuple {}); + let mut centroids = TapeWriter::<_, CentroidTuple>::create(index, false); + let vectors = (0..vchordrq_options.degree_of_parallelism) + .map(|_| TapeWriter::<_, VectorTuple>::create(index, true).first()) + .collect::>(); + let mut pointer_of_centroids = Vec::, u16)>>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + let vector = structures[i].centroids[j].as_borrowed(); + let (slices, metadata) = O::Vector::split(vector); + let mut chain = Ok(metadata); + let mut prefetch = Vec::new(); + for i in (0..slices.len()).rev() { + let (id, head) = centroids.push(match chain { + Ok(metadata) => CentroidTuple::_0 { + elements: slices[i].to_vec(), + metadata, + }, + Err(head) => CentroidTuple::_1 { + elements: slices[i].to_vec(), + head, + }, + }); + chain = Err(head); + prefetch.push(id); + } + prefetch.reverse(); + level.push(( + prefetch, + chain.expect_err("internal error: 0-dimensional vector"), + )); + } + pointer_of_centroids.push(level); + } + let mut pointer_of_firsts = Vec::>::new(); + for i in 0..structures.len() { + let mut level = Vec::new(); + for j in 0..structures[i].len() { + if i == 0 { + let frozen_tape = TapeWriter::<_, FrozenTuple>::create(index, false); + let appendable_tape = TapeWriter::<_, AppendableTuple>::create(index, false); + let frozen_first = { frozen_tape }.first(); + + let mut directory_tape = DirectoryTapeWriter::create(index, false); + directory_tape.push(&[frozen_first]); + let directory_tape = directory_tape.into_inner(); + + let mut jump = TapeWriter::<_, JumpTuple>::create(index, false); + jump.push(JumpTuple { + directory_first: { directory_tape }.first(), + frozen_first, + appendable_first: { appendable_tape }.first(), + centroid_prefetch: pointer_of_centroids[i][j].0.clone(), + centroid_head: pointer_of_centroids[i][j].1, + tuples: 0, + }); + level.push(jump.first()); + } else { + let mut tape = H1TapeWriter::create(index, O::Vector::count(dim) as _, false); + let centroid = structures[i].centroids[j].as_borrowed(); + for child in structures[i].children[j].iter().copied() { + let vector = structures[i - 1].centroids[child as usize].as_borrowed(); + let (code, delta) = O::build(vector, is_residual.then_some(centroid.own())); + tape.push(Branch { + code, + delta, + prefetch: pointer_of_centroids[i - 1][child as usize].0.clone(), + head: pointer_of_centroids[i - 1][child as usize].1, + extra: pointer_of_firsts[i - 1][child as usize], + norm: norm::(vector), + }); + } + let (mut tape, chunk) = tape.into_inner(); + H1TapeWriter::flush(&mut tape, O::Vector::count(dim) as _, chunk); + level.push(tape.first()); + } + } + pointer_of_firsts.push(level); + } + meta.push(MetaTuple { + dim, + height_of_root: structures.len() as u32, + is_residual, + rerank_in_heap: vchordrq_options.rerank_in_table, + indexed_vectors: Some(0), + centroids_first: centroids.first(), + vectors_first: vectors, + centroid_prefetch: pointer_of_centroids + .last() + .expect("internal error: empty structure")[0] + .0 + .clone(), + centroid_head: pointer_of_centroids + .last() + .expect("internal error: empty structure")[0] + .1, + centroid_norm: norm::( + structures + .last() + .expect("internal error: empty structure") + .centroids[0] + .as_borrowed(), + ), + first: pointer_of_firsts + .last() + .expect("internal error: empty structure")[0], + freepages_first: freepages.first(), + cells: structures.iter().map(|s| s.len() as _).collect(), + }); +} + +fn norm(vector: V::Borrowed<'_>) -> f32 { + V::squared_norm(vector).sqrt() +} diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs new file mode 100644 index 00000000..0aae88fe --- /dev/null +++ b/crates/vchordrq/src/bulkdelete.rs @@ -0,0 +1,198 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::{id_0, id_1}; +use crate::operator::Operator; +use crate::tape::by_next; +use crate::tuples::*; +use crate::{Opaque, tape}; +use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::FunctionalAccessor; +use std::num::NonZero; + +pub fn bulkdelete( + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) -> u64 +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let height_of_root = meta_tuple.height_of_root(); + + type State = Vec; + let mut state: State = vec![meta_tuple.first()]; + + drop(meta_guard); + + let mut live = 0_u64; + + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + by_next(index, first).inspect(|_| check()), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, _, first, _| results.push(first), + ); + } + results + }; + for _ in (1..height_of_root).rev() { + state = step(state); + } + for first in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut directory = tape::read_directory_tape::( + by_next(index, jump_tuple.directory_first()).inspect(|_| check()), + ); + { + let mut current = directory.next().unwrap_or(u32::MAX); + while current != u32::MAX { + check(); + let read = index.read(current); + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_ref(bytes); + if let FrozenTupleReader::_0(tuple) = tuple { + for p in tuple.payload().iter() { + if Some(true) == p.map(&callback) { + break 'scan (true, 0); + } + page_live += u64::from(p.is_some()); + } + } + } + (false, page_live) + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_mut(bytes); + if let FrozenTupleWriter::_0(mut tuple) = tuple { + for p in tuple.payload().iter_mut() { + if Some(true) == p.map(&callback) { + *p = None; + } + live += u64::from(p.is_some()); + } + } + } + } else { + live += page_live; + } + current = directory.next().unwrap_or(u32::MAX); + } + } + { + let mut current = jump_tuple.appendable_first(); + while current != u32::MAX { + check(); + let read = index.read(current); + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; + for i in 1..=read.len() { + let bytes = read.get(i).expect("data corruption"); + let tuple = AppendableTuple::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'scan (true, 0); + } + page_live += u64::from(p.is_some()); + } + (false, page_live) + }; + if flag { + drop(read); + let mut write = index.write(current, false); + for i in 1..=write.len() { + let bytes = write.get_mut(i).expect("data corruption"); + let mut tuple = AppendableTuple::deserialize_mut(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + *p = None; + } + live += u64::from(p.is_some()); + } + current = write.get_opaque().next; + } else { + live += page_live; + current = read.get_opaque().next; + } + } + } + } + + live +} + +pub fn bulkdelete_vectors( + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let vectors_first = meta_tuple.vectors_first().to_vec(); + + drop(meta_guard); + + for vectors_first in vectors_first { + let mut current = vectors_first; + while current != u32::MAX { + check(); + let read = index.read(current); + let flag = 'flag: { + for i in 1..=read.len() { + if let Some(bytes) = read.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + break 'flag true; + } + } + } + false + }; + if flag { + drop(read); + let mut write = index.write(current, true); + for i in 1..=write.len() { + if let Some(bytes) = write.get(i) { + let tuple = VectorTuple::::deserialize_ref(bytes); + let p = tuple.payload(); + if Some(true) == p.map(&callback) { + write.free(i); + } + }; + } + current = write.get_opaque().next; + } else { + current = read.get_opaque().next; + } + } + } +} diff --git a/crates/vchordrq/src/cache.rs b/crates/vchordrq/src/cache.rs new file mode 100644 index 00000000..d81ca98f --- /dev/null +++ b/crates/vchordrq/src/cache.rs @@ -0,0 +1,77 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::{id_0, id_1}; +use crate::tuples::{MetaTuple, WithReader}; +use crate::{Opaque, tape}; +use index::relation::{Page, PageGuard, RelationRead}; +use index_accessor::FunctionalAccessor; + +pub fn cache(index: &R, level: i32) -> Vec +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let height_of_root = meta_tuple.height_of_root(); + let centroids_first = meta_tuple.centroids_first(); + + type State = Vec; + let mut state: State = vec![meta_tuple.first()]; + + drop(meta_guard); + + let mut trace = Vec::new(); + + if level >= 0 { + // meta tuple + trace.push(0); + } + + if level >= 1 { + let mut step = |state: State| { + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + tape::by_next(index, first).inspect(|guard| trace.push(guard.id())), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, _, first, _| { + results.push(first); + }, + ); + } + results + }; + + // h1 tuples + for _ in (1..height_of_root).rev() { + state = step(state); + } + + // jump tuples + for first in state { + trace.push(first); + } + } + + if level >= 2 { + // centroid tuples + for guard in tape::by_next(index, centroids_first) { + trace.push(guard.id()); + } + } + + trace +} diff --git a/crates/vchordrq/src/centroids.rs b/crates/vchordrq/src/centroids.rs new file mode 100644 index 00000000..d2024c9e --- /dev/null +++ b/crates/vchordrq/src/centroids.rs @@ -0,0 +1,43 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::*; +use crate::tuples::*; +use index::relation::{Page, RelationRead}; +use index_accessor::Accessor1; + +pub fn read< + 'a, + R: RelationRead + 'a, + O: Operator, + A: Accessor1<::Element, ::Metadata>, +>( + mut prefetch: impl Iterator>, + head: u16, + accessor: A, +) -> A::Output { + let mut cursor = Err(head); + let mut result = accessor; + while let Err(head) = cursor { + let guard = prefetch.next().expect("data corruption"); + let bytes = guard.get(head).expect("data corruption"); + let tuple = CentroidTuple::::deserialize_ref(bytes); + result.push(tuple.elements()); + cursor = tuple.metadata_or_head(); + } + if prefetch.next().is_some() { + panic!("data corruption"); + } + result.finish(cursor.expect("data corruption")) +} diff --git a/crates/vchordrq/src/closure_lifetime_binder.rs b/crates/vchordrq/src/closure_lifetime_binder.rs new file mode 100644 index 00000000..86237624 --- /dev/null +++ b/crates/vchordrq/src/closure_lifetime_binder.rs @@ -0,0 +1,59 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +// Use stable language features as an alternative to `closure_lifetime_binder`. +// See https://github.com/rust-lang/rust/issues/97362. + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&'a mut A, &'a B) -> R, +{ + f +} + +#[inline(always)] +pub fn id_1(f: F) -> F +where + F: for<'a> FnMut(A, (&'a B, &'a C)) -> R, +{ + f +} + +#[inline(always)] +pub fn id_2(f: F) -> F +where + F: for<'a> FnMut(A, B, C, &'a D) -> R, +{ + f +} + +#[inline(always)] +pub fn id_3(f: F) -> F +where + T: index::relation::RelationWrite, + F: for<'a> Fn(&'a T, A, B) -> T::WriteGuard<'a>, +{ + f +} + +#[inline(always)] +pub fn id_4<'b, F, P, A, B, R: ?Sized>(f: F) -> F +where + P: index::prefetcher::Prefetcher<'b>, +

::Item: index::fetch::Fetch<'b>, + F: FnMut(A, P::Guards, B) -> R, +{ + f +} diff --git a/crates/vchordrq/src/consume.rs b/crates/vchordrq/src/consume.rs new file mode 100644 index 00000000..e68d9c0d --- /dev/null +++ b/crates/vchordrq/src/consume.rs @@ -0,0 +1,41 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::*; +use crate::{Opaque, freepages}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; + +pub fn consume(index: &R) +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let freepages_first = meta_tuple.freepages_first(); + + drop(meta_guard); + + let id = index + .extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + false, + ) + .id(); + + freepages::free(index, freepages_first, id); +} diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs new file mode 100644 index 00000000..c8e40252 --- /dev/null +++ b/crates/vchordrq/src/cost.rs @@ -0,0 +1,40 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::{MetaTuple, WithReader}; +use index::relation::{Page, RelationRead}; + +pub struct Cost { + pub dim: u32, + pub cells: Vec, + pub indexed_vectors: Option, +} + +#[must_use] +pub fn cost(index: &R) -> Cost { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let cells = meta_tuple.cells().to_vec(); + let indexed_vectors = meta_tuple.indexed_vectors(); + + drop(meta_guard); + + Cost { + dim, + cells, + indexed_vectors, + } +} diff --git a/crates/vchordrq/src/fast_heap.rs b/crates/vchordrq/src/fast_heap.rs new file mode 100644 index 00000000..c36f91a9 --- /dev/null +++ b/crates/vchordrq/src/fast_heap.rs @@ -0,0 +1,115 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use index::prefetcher::Sequence; +use std::collections::BinaryHeap; +use std::num::NonZero; + +pub struct SortHeap { + inner: Vec, + t: NonZero, +} + +impl SortHeap { + pub fn peek(&self) -> Option<&T> { + self.inner.last() + } +} + +pub enum FastHeap { + Sorted(SortHeap), + Binary(BinaryHeap), +} + +impl FastHeap { + pub fn from_vec(vec: Vec) -> Self { + let n = vec.len(); + if let Some(t) = NonZero::new(n / 384) { + let mut inner = vec; + let index = n - t.get(); + inner.select_nth_unstable(index); + inner[index..].sort_unstable(); + Self::Sorted(SortHeap { inner, t }) + } else { + Self::Binary(BinaryHeap::from(vec)) + } + } + pub fn pop(&mut self) -> Option { + match self { + FastHeap::Sorted(SortHeap { inner, t }) => { + let Some(k) = inner.pop() else { unreachable!() }; + if let Some(value) = NonZero::new(t.get() - 1) { + *t = value; + } else { + *self = FastHeap::Binary(std::mem::take(inner).into()); + } + Some(k) + } + FastHeap::Binary(x) => x.pop(), + } + } + pub fn peek(&self) -> Option<&T> { + match self { + FastHeap::Sorted(x) => x.peek(), + FastHeap::Binary(x) => x.peek(), + } + } +} + +impl From> for FastHeap { + fn from(value: Vec) -> Self { + Self::from_vec(value) + } +} + +impl Sequence for FastHeap { + type Item = T; + type Inner = std::vec::IntoIter; + fn next(&mut self) -> Option { + self.pop() + } + fn peek(&mut self) -> Option<&Self::Item> { + (self as &Self).peek() + } + fn into_inner(self) -> Self::Inner { + match self { + FastHeap::Sorted(sort_heap) => sort_heap.inner.into_iter(), + FastHeap::Binary(binary_heap) => binary_heap.into_vec().into_iter(), + } + } +} + +#[test] +fn test_fast_heap() { + for _ in 0..if cfg!(not(miri)) { 1000 } else { 1 } { + let sequence = (0..1000).map(|_| rand::random::()).collect::>(); + let answer = { + let mut x = sequence.clone(); + x.sort_by_key(|x| std::cmp::Reverse(*x)); + x + }; + let result = { + let mut x = FastHeap::from_vec(sequence.clone()); + std::iter::from_fn(|| x.pop()).collect::>() + }; + assert_eq!(answer, result); + } +} + +#[test] +fn test_issue_209() { + let mut heap = FastHeap::from_vec(vec![0]); + assert_eq!(heap.pop(), Some(0)); + assert_eq!(heap.pop(), None); +} diff --git a/crates/vchordrq/src/freepages.rs b/crates/vchordrq/src/freepages.rs new file mode 100644 index 00000000..9850f452 --- /dev/null +++ b/crates/vchordrq/src/freepages.rs @@ -0,0 +1,50 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::Opaque; +use crate::tuples::*; +use index::relation::{Page, RelationWrite}; + +pub fn alloc(index: &R, freepages_first: u32) -> Option> +where + R::Page: Page, +{ + let mut freepages_guard = index.write(freepages_first, false); + let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); + let mut freepages_tuple = FreepagesTuple::deserialize_mut(freepages_bytes); + let id = *freepages_tuple.first(); + if id != u32::MAX { + let mut guard = index.write(id, false); + *freepages_tuple.first() = guard.get_opaque_mut().next; + drop(freepages_guard); // write log of freespaces_guard + Some(guard) + } else { + None + } +} + +// the page must be inaccessible in the graph +pub fn free(index: &R, freepages_first: u32, id: u32) +where + R::Page: Page, +{ + let mut guard = index.write(id, false); + let mut freepages_guard = index.write(freepages_first, false); + let freepages_bytes = freepages_guard.get_mut(1).expect("data corruption"); + let mut freepages_tuple = FreepagesTuple::deserialize_mut(freepages_bytes); + guard.get_opaque_mut().next = *freepages_tuple.first(); + drop(guard); // write log of guard + *freepages_tuple.first() = id; + drop(freepages_guard); // write log of freepages_guard +} diff --git a/crates/vchordrq/src/insert.rs b/crates/vchordrq/src/insert.rs new file mode 100644 index 00000000..3d5308e8 --- /dev/null +++ b/crates/vchordrq/src/insert.rs @@ -0,0 +1,212 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::id_0; +use crate::linked_vec::LinkedVec; +use crate::operator::*; +use crate::tuples::*; +use crate::{Opaque, centroids, tape, vectors}; +use always_equal::AlwaysEqual; +use distance::Distance; +use index::bump::Bump; +use index::fetch::BorrowedIter; +use index::prefetcher::{Prefetcher, PrefetcherHeapFamily}; +use index::relation::{Page, RelationRead, RelationWrite}; +use index_accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; +use vector::{VectorBorrowed, VectorOwned}; + +pub trait InsertChooser { + fn choose(&mut self, n: NonZero) -> usize; +} + +type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); + +pub fn insert_vector( + index: &R, + payload: NonZero, + vector: ::Borrowed<'_>, + chooser: &mut impl InsertChooser, + skip_search: bool, +) -> (Vec, u16) +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let rerank_in_heap = meta_tuple.rerank_in_heap(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + let vectors_first = { + let l = meta_tuple.vectors_first(); + let n = NonZero::new(l.len()).expect("data corruption"); + let i = chooser.choose(n); + l[i] + }; + + drop(meta_guard); + + if !rerank_in_heap { + vectors::append::(index, vectors_first, vector, payload, skip_search) + } else { + (Vec::new(), 0) + } +} + +pub fn insert<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + payload: NonZero, + vector: ::Borrowed<'_>, + key: (Vec, u16), + bump: &'b impl Bump, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, + skip_freespaces: bool, +) where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + let freepages_first = meta_tuple.freepages_first(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + let epsilon = 1.9; + + type State = (Reverse, AlwaysEqual, AlwaysEqual); + let mut state: State = if is_residual { + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); + let head = meta_tuple.centroid_head(); + let distance = centroids::read::( + prefetch.map(|id| index.read(id)), + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + (Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first)) + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + (Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first)) + }; + + drop(meta_guard); + let lut = (O::Vector::block_preprocess(vector),); + + let mut step = |state: State| { + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + { + let (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) = state; + tape::read_h1_tape::( + tape::by_next(index, first), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + |(rough, err), head, norm, first, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + ))), + )); + }, + ); + } + let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); + let mut cache = BinaryHeap::<(_, _, _)>::new(); + { + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + { + let distance = centroids::read::( + prefetch, + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); + } + cache.pop() + } + .expect("invariant is violated: tree is not height-balanced") + }; + + for _ in (1..height_of_root).rev() { + state = step(state); + } + + let (_, _, AlwaysEqual(first)) = state; + + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + + let (code, delta) = O::build( + vector, + is_residual.then(|| { + centroids::read::( + jump_tuple + .centroid_prefetch() + .iter() + .map(|&id| index.read(id)), + jump_tuple.centroid_head(), + FunctionalAccessor::new( + (dim, Vec::new()), + id_0(|(_, elements): &mut (_, Vec<_>), slice| { + elements.extend_from_slice(slice) + }), + |(dim, elements), metadata| O::Vector::pack(dim, elements, metadata), + ), + ) + }), + ); + + let (prefetch, head) = key; + let serialized = AppendableTuple::serialize(&AppendableTuple { + metadata: [ + code.0.dis_u_2, + code.0.factor_cnt, + code.0.factor_ip, + code.0.factor_err, + ], + delta, + payload: Some(payload), + prefetch, + head, + elements: rabitq::bit::binary::pack_code(&code.1), + }); + + tape::append( + index, + jump_tuple.appendable_first(), + &serialized, + false, + (!skip_freespaces).then_some(freepages_first), + ); +} diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs new file mode 100644 index 00000000..b60c2e3b --- /dev/null +++ b/crates/vchordrq/src/lib.rs @@ -0,0 +1,81 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod build; +mod bulkdelete; +mod cache; +mod centroids; +mod closure_lifetime_binder; +mod consume; +mod cost; +mod fast_heap; +mod freepages; +mod insert; +mod linked_vec; +mod maintain; +mod maxsim_cost; +mod prewarm; +mod rerank; +mod search; +mod statistics; +mod tape; +mod tape_writer; +mod tuples; +mod vectors; + +pub mod operator; +pub mod types; + +pub use build::build; +pub use bulkdelete::{bulkdelete, bulkdelete_vectors}; +pub use cache::cache; +pub use consume::consume; +pub use cost::cost; +pub use fast_heap::FastHeap; +pub use insert::{InsertChooser, insert, insert_vector}; +pub use maintain::{MaintainChooser, maintain}; +pub use maxsim_cost::{ + MaxsimCostBackend, MaxsimCostEstimate, MaxsimCostInput, estimate_maxsim_cost, +}; +pub use prewarm::prewarm; +pub use rerank::{how, rerank_heap, rerank_index}; +pub use search::{default_search, maxsim_search}; +pub use statistics::set_indexed_vectors; + +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, FromBytes, IntoBytes, Immutable, KnownLayout)] +pub struct Opaque { + pub next: u32, + pub skip: u32, +} + +#[allow(unsafe_code)] +unsafe impl index::relation::Opaque for Opaque {} + +pub(crate) struct Branch { + pub code: rabitq::bit::Code, + pub delta: f32, + pub prefetch: Vec, + pub head: u16, + pub norm: f32, + pub extra: T, +} + +#[derive(Debug, Clone, Copy)] +pub enum RerankMethod { + Index, + Heap, +} diff --git a/crates/vchordrq/src/linked_vec.rs b/crates/vchordrq/src/linked_vec.rs new file mode 100644 index 00000000..a42f23fc --- /dev/null +++ b/crates/vchordrq/src/linked_vec.rs @@ -0,0 +1,50 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub struct LinkedVec { + inner: Vec>, + last: Vec, +} + +impl LinkedVec { + pub fn new() -> Self { + Self { + inner: Vec::new(), + last: Vec::with_capacity(4096), + } + } + pub fn push(&mut self, value: T) { + if self.last.len() == self.last.capacity() { + self.reserve(); + } + #[allow(unsafe_code)] + unsafe { + std::hint::assert_unchecked(self.last.len() != self.last.capacity()); + } + self.last.push(value); + } + #[cold] + fn reserve(&mut self) { + let fresh = Vec::with_capacity(self.last.capacity() * 4); + self.inner.push(core::mem::replace(&mut self.last, fresh)); + } + pub fn into_vec(self) -> Vec { + let mut last = self.last; + last.reserve(self.inner.iter().map(Vec::len).sum::()); + for x in self.inner { + last.extend(x); + } + last + } +} diff --git a/crates/vchordrq/src/maintain.rs b/crates/vchordrq/src/maintain.rs new file mode 100644 index 00000000..3d79ad7c --- /dev/null +++ b/crates/vchordrq/src/maintain.rs @@ -0,0 +1,316 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::{id_0, id_1, id_2, id_3}; +use crate::operator::{Operator, Vector}; +use crate::tape_writer::{DirectoryTapeWriter, FrozenTapeWriter}; +use crate::tuples::*; +use crate::{Branch, Opaque, freepages, tape}; +use index::prefetcher::PrefetcherSequenceFamily; +use index::relation::{ + Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; +use index_accessor::FunctionalAccessor; +use rabitq::packing::unpack; +use std::cell::RefCell; + +pub trait MaintainChooser { + fn choose(&mut self, i: usize) -> bool; +} + +pub struct Maintain { + pub number_of_formerly_allocated_pages: usize, + pub number_of_freshly_allocated_pages: usize, + pub number_of_freed_pages: usize, +} + +pub fn maintain<'b, R: RelationRead + RelationWrite, O: Operator>( + index: &'b R, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, + chooser: &mut impl MaintainChooser, + check: impl Fn(), +) -> Maintain +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let height_of_root = meta_tuple.height_of_root(); + let freepages_first = meta_tuple.freepages_first(); + + type State = Vec; + let mut state: State = vec![meta_tuple.first()]; + + drop(meta_guard); + + let step = |state: State| { + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + tape::by_next(index, first).inspect(|_| check()), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, _, first, _| results.push(first), + ); + } + results + }; + + for _ in (1..height_of_root).rev() { + state = step(state); + } + + struct Buffers { + pages: Vec, + number_of_formerly_allocated_pages: usize, + number_of_freshly_allocated_pages: usize, + } + + let buffers = RefCell::new(Buffers { + pages: Vec::new(), + number_of_formerly_allocated_pages: 0, + number_of_freshly_allocated_pages: 0, + }); + + for (idx, first) in state.into_iter().enumerate() { + if !chooser.choose(idx) { + continue; + } + + let mut jump_guard = index.write(first, false); + let jump_bytes = jump_guard.get_mut(1).expect("data corruption"); + let mut jump_tuple = JumpTuple::deserialize_mut(jump_bytes); + + let hooked_index = RelationHooked(index, { + id_3(|index: &R, opaque: Opaque, tracking_freespace: bool| { + if !tracking_freespace { + let mut buffers = buffers.borrow_mut(); + if let Some(id) = buffers.pages.pop() { + drop(buffers); + let mut guard = index.write(id, false); + guard.clear(opaque); + guard + } else if let Some(mut guard) = freepages::alloc(index, freepages_first) { + buffers.number_of_formerly_allocated_pages += 1; + drop(buffers); + guard.clear(opaque); + guard + } else { + buffers.number_of_freshly_allocated_pages += 1; + drop(buffers); + index.extend(opaque, false) + } + } else { + index.extend(opaque, true) + } + }) + }); + + let mut tape = FrozenTapeWriter::create(&hooked_index, O::Vector::count(dim) as _, false); + + let mut trace_directory = Vec::new(); + let mut trace_forzen = Vec::new(); + let mut trace_appendable = Vec::new(); + + let mut tuples = 0_u64; + let mut callback = id_2(|(code, delta): (_, _), head, payload, prefetch: &[_]| { + tape.push(Branch { + code, + delta, + prefetch: prefetch.to_vec(), + head, + norm: 0.0, + extra: payload, + }); + tuples += 1; + }); + let directory = tape::read_directory_tape::( + tape::by_next(index, *jump_tuple.directory_first()) + .inspect(|_| check()) + .inspect(|guard| trace_directory.push(guard.id())), + ); + tape::read_frozen_tape::( + tape::by_directory(&mut prefetch_h0_tuples, directory) + .inspect(|_| check()) + .inspect(|guard| trace_forzen.push(guard.id())), + || { + FunctionalAccessor::new( + Vec::<[u8; 16]>::new(), + Vec::<[u8; 16]>::extend_from_slice, + id_1( + |elements: Vec<_>, (metadata, delta): (&[[f32; 32]; 4], &[f32; 32])| { + let unpacked = unpack(&elements); + std::array::from_fn(|i| { + let f = |&x| [x & 1 != 0, x & 2 != 0, x & 4 != 0, x & 8 != 0]; + let signs = unpacked[i].iter().flat_map(f).collect::>(); + ( + ( + rabitq::bit::CodeMetadata { + dis_u_2: metadata[0][i], + factor_cnt: metadata[1][i], + factor_ip: metadata[2][i], + factor_err: metadata[3][i], + }, + signs, + ), + delta[i], + ) + }) + }, + ), + ) + }, + &mut callback, + ); + tape::read_appendable_tape::( + tape::by_next(index, *jump_tuple.appendable_first()) + .inspect(|_| check()) + .inspect(|guard| trace_appendable.push(guard.id())), + |metadata, elements, delta| { + let signs = elements + .iter() + .flat_map(|x| std::array::from_fn::<_, 64, _>(|i| *x & (1 << i) != 0)) + .take(dim as _) + .collect::>(); + ( + ( + rabitq::bit::CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }, + signs, + ), + delta, + ) + }, + &mut callback, + ); + + let (frozen_tape, branches) = tape.into_inner(); + + let mut appendable_tape = tape::TapeWriter::create(&hooked_index, false); + + for branch in branches { + appendable_tape.push(AppendableTuple { + metadata: [ + branch.code.0.dis_u_2, + branch.code.0.factor_cnt, + branch.code.0.factor_ip, + branch.code.0.factor_err, + ], + elements: rabitq::bit::binary::pack_code(&branch.code.1), + delta: branch.delta, + prefetch: branch.prefetch, + head: branch.head, + payload: Some(branch.extra), + }); + } + + let frozen_first = { frozen_tape }.first(); + + let directory = tape::by_next(index, frozen_first) + .inspect(|_| check()) + .map(|guard| guard.id()) + .collect::>(); + + let mut directory_tape = DirectoryTapeWriter::create(&hooked_index, false); + directory_tape.push(directory.as_slice()); + let directory_tape = directory_tape.into_inner(); + + *jump_tuple.directory_first() = { directory_tape }.first(); + *jump_tuple.frozen_first() = frozen_first; + *jump_tuple.appendable_first() = { appendable_tape }.first(); + *jump_tuple.tuples() = tuples; + + drop(jump_guard); + + let mut buffers = buffers.borrow_mut(); + buffers.pages.extend_from_slice(&trace_directory); + buffers.pages.extend_from_slice(&trace_forzen); + buffers.pages.extend_from_slice(&trace_appendable); + } + + let buffers = RefCell::into_inner(buffers); + for id in buffers.pages.iter().copied() { + freepages::free(index, freepages_first, id); + } + + Maintain { + number_of_formerly_allocated_pages: buffers.number_of_formerly_allocated_pages, + number_of_freshly_allocated_pages: buffers.number_of_freshly_allocated_pages, + number_of_freed_pages: buffers.pages.len(), + } +} + +#[derive(Clone)] +struct RelationHooked<'b, R, E>(&'b R, E); + +impl<'b, R, E> Relation for RelationHooked<'b, R, E> +where + R: Relation, + E: Clone, +{ + type Page = R::Page; +} + +impl<'b, R, E> RelationReadTypes for RelationHooked<'b, R, E> +where + R: RelationRead, + E: Clone, +{ + type ReadGuard<'a> = R::ReadGuard<'a>; +} + +impl<'b, R, E> RelationRead for RelationHooked<'b, R, E> +where + R: RelationRead, + E: Clone, +{ + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + self.0.read(id) + } +} + +impl<'b, R, E> RelationWriteTypes for RelationHooked<'b, R, E> +where + R: RelationWrite, + E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, +{ + type WriteGuard<'a> = R::WriteGuard<'a>; +} + +impl<'b, R, E> RelationWrite for RelationHooked<'b, R, E> +where + R: RelationWrite, + E: Clone + for<'a> Fn(&'a R, ::Opaque, bool) -> R::WriteGuard<'a>, +{ + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.0.write(id, tracking_freespace) + } + + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_> { + (self.1)(self.0, opaque, tracking_freespace) + } + + fn search(&self, freespace: usize) -> Option> { + self.0.search(freespace) + } +} diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs new file mode 100644 index 00000000..c63966d7 --- /dev/null +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -0,0 +1,206 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[derive(Clone, Copy, Debug)] +pub enum MaxsimCostBackend { + CoarseOnly, + CpuExact, + Gpu, + Auto, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostInput { + pub heap_rows: f64, + pub index_tokens: f64, + pub token_nodes_per_query: f64, + pub base_index_pages: f64, + pub dimension: u32, + pub element_bits: u32, + pub query_tokens: u32, + pub limit_tuples: Option, + pub filter_selectivity: f64, + pub candidate_limit: Option, + pub backend: MaxsimCostBackend, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostEstimate { + pub startup_cost: f64, + pub total_cost: f64, + pub selectivity: f64, + pub index_pages: f64, +} + +pub fn estimate_maxsim_cost(input: MaxsimCostInput) -> MaxsimCostEstimate { + let heap_rows = input.heap_rows.max(1.0); + let index_tokens = input.index_tokens.max(heap_rows); + let query_tokens = f64::from(input.query_tokens.max(1)); + let average_document_tokens = (index_tokens / heap_rows).clamp(1.0, 65_536.0); + let token_visits = input.token_nodes_per_query.max(1.0) * query_tokens; + + // We do not yet persist page-level candidate statistics. Until then, use a + // conservative occupancy estimate: token visits are sampled from the token + // index, and a page is a candidate when at least one of its average tokens + // is visited. This is intentionally bounded by the heap row count. + let token_visit_fraction = (token_visits / index_tokens).clamp(0.0, 1.0); + let candidate_probability = 1.0 - (1.0 - token_visit_fraction).powf(average_document_tokens); + let generated_pages = (heap_rows * candidate_probability).clamp(1.0, heap_rows); + + let filtered_limit = input + .limit_tuples + .map(|limit| limit.max(1.0) / input.filter_selectivity.clamp(1e-9, 1.0)); + let exact_candidate_count = input + .candidate_limit + .map_or(generated_pages, |limit| { + f64::from(limit).min(generated_pages) + }) + .clamp(1.0, heap_rows); + let returned_pages = match input.backend { + MaxsimCostBackend::CoarseOnly => filtered_limit + .unwrap_or(generated_pages) + .min(generated_pages), + MaxsimCostBackend::CpuExact | MaxsimCostBackend::Gpu | MaxsimCostBackend::Auto => { + exact_candidate_count + } + }; + + // Candidate generation and aggregation are eager in the current scanner, + // so their full work belongs to startup cost even when SQL has a small + // LIMIT. The constants are deliberately conservative placeholders until + // committed corpus benchmarks replace them with fitted values. + let search_cost = 0.001 * token_visits; + let aggregation_cost = 0.01 * token_visits + 0.05 * generated_pages; + let exact_components = exact_candidate_count + * average_document_tokens + * query_tokens + * f64::from(input.dimension.max(1)); + let tensor_bytes = exact_candidate_count + * average_document_tokens + * f64::from(input.dimension.max(1)) + * f64::from(input.element_bits.max(1)) + / 8.0; + let cpu_exact_cost = exact_candidate_count + exact_components * 1e-6; + let gpu_exact_cost = 5.0 + tensor_bytes * 1e-7 + exact_components * 5e-8; + let backend_cost = match input.backend { + MaxsimCostBackend::CoarseOnly => 0.0, + MaxsimCostBackend::CpuExact => cpu_exact_cost, + MaxsimCostBackend::Gpu => gpu_exact_cost, + // Price a small but nonzero fallback risk. Runtime still performs a + // complete CPU rerank on every GPU failure. + MaxsimCostBackend::Auto => gpu_exact_cost + 0.05 * cpu_exact_cost, + }; + let startup_cost = search_cost + aggregation_cost + backend_cost; + let total_cost = startup_cost + returned_pages; + let selectivity = (returned_pages / heap_rows).clamp(1e-9, 1.0); + let index_pages = + input.base_index_pages.max(1.0) * (1.0 + 0.25 * (query_tokens - 1.0).max(0.0)); + + MaxsimCostEstimate { + startup_cost, + total_cost, + selectivity, + index_pages, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input(backend: MaxsimCostBackend) -> MaxsimCostInput { + MaxsimCostInput { + heap_rows: 34_054.0, + index_tokens: 25_438_338.0, + token_nodes_per_query: 10_000.0, + base_index_pages: 20_000.0, + dimension: 320, + element_bits: 16, + query_tokens: 32, + limit_tuples: Some(20.0), + filter_selectivity: 1.0, + candidate_limit: Some(256), + backend, + } + } + + #[test] + fn maxsim_is_never_zero_cost() { + for backend in [ + MaxsimCostBackend::CoarseOnly, + MaxsimCostBackend::CpuExact, + MaxsimCostBackend::Gpu, + MaxsimCostBackend::Auto, + ] { + let estimate = estimate_maxsim_cost(input(backend)); + assert!(estimate.startup_cost > 0.0); + assert!(estimate.total_cost >= estimate.startup_cost); + assert!((1e-9..=1.0).contains(&estimate.selectivity)); + assert!(estimate.index_pages >= 1.0); + } + } + + #[test] + fn query_token_count_increases_eager_work() { + let one = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 1, + ..input(MaxsimCostBackend::CpuExact) + }); + let many = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 64, + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(many.startup_cost > one.startup_cost); + assert!(many.index_pages > one.index_pages); + } + + #[test] + fn exact_candidate_limit_bounds_rows_and_cost() { + let small = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(128), + ..input(MaxsimCostBackend::CpuExact) + }); + let large = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(2048), + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(small.selectivity < large.selectivity); + assert!(small.startup_cost < large.startup_cost); + } + + #[test] + fn auto_prices_more_than_gpu_for_fallback_risk() { + let gpu = estimate_maxsim_cost(input(MaxsimCostBackend::Gpu)); + let auto = estimate_maxsim_cost(input(MaxsimCostBackend::Auto)); + assert!(auto.startup_cost > gpu.startup_cost); + } + + #[test] + fn missing_stats_remain_finite() { + let estimate = estimate_maxsim_cost(MaxsimCostInput { + heap_rows: -1.0, + index_tokens: 0.0, + token_nodes_per_query: 0.0, + base_index_pages: 0.0, + filter_selectivity: 0.0, + limit_tuples: None, + candidate_limit: None, + ..input(MaxsimCostBackend::CoarseOnly) + }); + assert!(estimate.startup_cost.is_finite()); + assert!(estimate.total_cost.is_finite()); + assert!(estimate.selectivity.is_finite()); + assert!(estimate.index_pages.is_finite()); + } +} diff --git a/crates/vchordrq/src/operator.rs b/crates/vchordrq/src/operator.rs new file mode 100644 index 00000000..d3a893d2 --- /dev/null +++ b/crates/vchordrq/src/operator.rs @@ -0,0 +1,987 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use distance::Distance; +use index_accessor::{ + Accessor1, Accessor2, ByteDistanceAccessor, DefaultWithDimension, DistanceAccessor, Dot, + HalfbyteDistanceAccessor, L2S, RAccess, +}; +use rabitq::bit::CodeMetadata; +use rabitq::bit::binary::BinaryLut; +use rabitq::bit::block::{BlockLut, STEP}; +use simd::{Floating, f16}; +use std::fmt::Debug; +use std::marker::PhantomData; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; +use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +#[derive(Debug)] +pub struct BlockAccessor([u32; 32], F); + +impl> + Accessor2<[u8; 16], [u8; 16], (&[[f32; 32]; 4], &[f32; 32]), ()> for BlockAccessor +{ + type Output = [F::Output; 32]; + + #[inline(always)] + fn push(&mut self, input: &[[u8; 16]], target: &[[u8; 16]]) { + use std::iter::zip; + + for (input, target) in zip(input.chunks(STEP), target.chunks(STEP)) { + let delta = simd::fast_scan::scan(input, target); + simd::fast_scan::accu(&mut self.0, &delta); + } + } + + #[inline(always)] + fn finish(mut self, (metadata, delta): (&[[f32; 32]; 4], &[f32; 32]), (): ()) -> Self::Output { + std::array::from_fn(|i| { + (self.1).call( + self.0[i], + CodeMetadata { + dis_u_2: metadata[0][i], + factor_cnt: metadata[1][i], + factor_ip: metadata[2][i], + factor_err: metadata[3][i], + }, + delta[i], + ) + }) + } +} + +#[derive(Debug, Clone)] +pub struct CloneAccessor(Vec); + +impl Default for CloneAccessor { + #[inline(always)] + fn default() -> Self { + Self(Vec::new()) + } +} + +impl Accessor1 for CloneAccessor { + type Output = V; + + #[inline(always)] + fn push(&mut self, input: &[V::Element]) { + self.0.extend(input); + } + + #[inline(always)] + fn finish(self, (metadata, dim): (V::Metadata, u32)) -> Self::Output { + V::pack(dim, self.0, metadata) + } +} + +pub trait Vector: VectorOwned { + type Element: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + + type Metadata: Debug + Copy + FromBytes + IntoBytes + Immutable + KnownLayout; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[Self::Element]>, Self::Metadata); + + fn count(dim: u32) -> u32; + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata); + + fn pack(dim: u32, elements: Vec, metadata: Self::Metadata) -> Self; + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut; + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut); + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code; + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32; +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f32; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[f32]>, ()) { + let vector = vector.slice(); + ( + match vector.len() { + 0 => unreachable!(), + 1..=960 => vec![vector], + 961..=1280 => vec![&vector[..640], &vector[640..]], + 1281.. => vector.chunks(1920).collect(), + }, + (), + ) + } + + fn count(dim: u32) -> u32 { + match dim { + 0 => unreachable!(), + 1..=960 => 1, + 961..=1280 => 2, + 1281.. => dim.div_ceil(1920), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + rabitq::bit::block::preprocess(vector.slice()) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + rabitq::bit::preprocess(vector.slice()) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + rabitq::bit::code(vector.slice()) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + f32::reduce_sum_of_x2(vector.slice()) + } +} + +impl Vector for VectOwned { + type Metadata = (); + + type Element = f16; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[f16]>, ()) { + let vector = vector.slice(); + ( + match vector.len() { + 0 => unreachable!(), + 1..=1920 => vec![vector], + 1921..=2560 => vec![&vector[..1280], &vector[1280..]], + 2561.. => vector.chunks(3840).collect(), + }, + (), + ) + } + + fn count(dim: u32) -> u32 { + match dim { + 0 => unreachable!(), + 1..=1920 => 1, + 1921..=2560 => 2, + 2561.. => dim.div_ceil(3840), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + (vector.slice(), ()) + } + + fn pack(_: u32, elements: Vec, (): Self::Metadata) -> Self { + VectOwned::new(elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + rabitq::bit::block::preprocess(&f16::vector_to_f32(vector.slice())) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + rabitq::bit::preprocess(&f16::vector_to_f32(vector.slice())) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + rabitq::bit::code(&f16::vector_to_f32(vector.slice())) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + f16::reduce_sum_of_x2(vector.slice()) + } +} + +impl Vector for Rabitq8Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[u8]>, [f32; 4]) { + ( + match vector.packed_code().len() { + 0 => unreachable!(), + 1..=3840 => vec![vector.packed_code()], + 3841..=5120 => vec![&vector.packed_code()[..1280], &vector.packed_code()[1280..]], + 5121.. => vector.packed_code().chunks(7680).collect(), + }, + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn count(dim: u32) -> u32 { + match dim { + 0 => unreachable!(), + 1..=3840 => 1, + 3841..=5120 => 2, + 5121.. => dim.div_ceil(7680), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn pack(dim: u32, elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq8Owned::new(dim, _0, _1, _2, _3, elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::block::preprocess(&result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::preprocess(&result) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + let n = vector.dim(); + let sum_of_abs_x = vector.sum_of_abs_x(); + let sum_of_x_2 = vector.sum_of_x2(); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.unpacked_code().filter(|&x| x >= 128).count(); + let cnt_neg = vector.unpacked_code().filter(|&x| x <= 127).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let vector = vector.unpacked_code(); + let mut signs = Vec::new(); + for x in vector { + signs.push(x >= 128); + } + signs + }, + ) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + vector.sum_of_x2() + } +} + +impl Vector for Rabitq4Owned { + type Metadata = [f32; 4]; + + type Element = u8; + + fn split(vector: Self::Borrowed<'_>) -> (Vec<&[u8]>, [f32; 4]) { + ( + match vector.packed_code().len() { + 0 => unreachable!(), + 1..=3840 => vec![vector.packed_code()], + 3841..=5120 => vec![&vector.packed_code()[..1280], &vector.packed_code()[1280..]], + 5121.. => vector.packed_code().chunks(7680).collect(), + }, + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn count(dim: u32) -> u32 { + match dim { + 0 => unreachable!(), + 1..=7680 => 1, + 7681..=10240 => 2, + 10241.. => dim.div_ceil(15360), + } + } + + fn unpack(vector: Self::Borrowed<'_>) -> (&[Self::Element], Self::Metadata) { + ( + vector.packed_code(), + [ + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + ], + ) + } + + fn pack(dim: u32, elements: Vec, [_0, _1, _2, _3]: Self::Metadata) -> Self { + Rabitq4Owned::new(dim, _0, _1, _2, _3, elements) + } + + fn block_preprocess(vector: Self::Borrowed<'_>) -> BlockLut { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::block::preprocess(&result) + } + + fn preprocess(vector: Self::Borrowed<'_>) -> (BlockLut, BinaryLut) { + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::bit::preprocess(&result) + } + + fn code(vector: Self::Borrowed<'_>) -> rabitq::bit::Code { + let n = vector.dim(); + let sum_of_abs_x = vector.sum_of_abs_x(); + let sum_of_x_2 = vector.sum_of_x2(); + ( + CodeMetadata { + dis_u_2: sum_of_x_2, + factor_cnt: { + let cnt_pos = vector.unpacked_code().filter(|&x| x >= 8).count(); + let cnt_neg = vector.unpacked_code().filter(|&x| x <= 7).count(); + cnt_pos as f32 - cnt_neg as f32 + }, + factor_ip: sum_of_x_2 / sum_of_abs_x, + factor_err: { + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (n as f32).sqrt(); + dis_u * (1.0 / (x_0 * x_0) - 1.0).sqrt() / (n as f32 - 1.0).sqrt() + }, + }, + { + let vector = vector.unpacked_code(); + let mut signs = Vec::new(); + for x in vector { + signs.push(x >= 8); + } + signs + }, + ) + } + + fn squared_norm(vector: Self::Borrowed<'_>) -> f32 { + vector.sum_of_x2() + } +} + +pub trait Operator: 'static + Debug + Copy { + type Vector: Vector; + + type DistanceAccessor: DefaultWithDimension + + Accessor2< + ::Element, + ::Element, + ::Metadata, + ::Metadata, + Output = Distance, + >; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]>; + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32); + + fn build( + vector: ::Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32); +} + +#[derive(Debug)] +pub struct Op(PhantomData V>, PhantomData D>); + +impl Clone for Op { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for Op {} + +impl Operator for Op, L2S> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_l2s(value, code, lut.0) + } else { + rabitq::bit::block::half_process_l2s_residual(value, code, lut.0, dis_f, delta) + } + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_l2s(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_l2s_residual(value, code, lut.0, dis_f, delta) + } + } + } + + fn build( + vector: VectBorrowed<'_, f32>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dim = vector.dim(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .sum::() + / (dim as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); + 2.0 * dis_u * t / x_0 + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } + } +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + } else { + rabitq::bit::block::half_process_dot_residual( + sum, code, lut.0, dis_f, delta, norm, + ) + } + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } else { + rabitq::bit::binary::half_process_dot_residual(sum, code, lut.0, dis_f, delta, norm) + } + } + } + + fn build( + vector: VectBorrowed<'_, f32>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f32::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dim = vector.dim(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .sum::() + / (dim as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); + dis_u * t / x_0 - f32::reduce_sum_of_xy(residual.slice(), centroid.slice()) + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } + } +} + +impl Operator for Op, L2S> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, L2S>; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_l2s(value, code, lut.0) + } else { + rabitq::bit::block::half_process_l2s_residual(value, code, lut.0, dis_f, delta) + } + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_l2s(value, code, lut.0) + } else { + rabitq::bit::binary::half_process_l2s_residual(value, code, lut.0, dis_f, delta) + } + } + } + + fn build( + vector: VectBorrowed<'_, f16>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dim = vector.dim(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .map(simd::F16::_to_f32) + .sum::() + / (dim as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); + 2.0 * dis_u * t / x_0 + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } + } +} + +impl Operator for Op, Dot> { + type Vector = VectOwned; + + type DistanceAccessor = DistanceAccessor, Dot>; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, delta| { + if !is_residual { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + } else { + rabitq::bit::block::half_process_dot_residual( + sum, code, lut.0, dis_f, delta, norm, + ) + } + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + dis_f: f32, + norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + move |metadata: [f32; 4], elements: &[u64], delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + if !is_residual { + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } else { + rabitq::bit::binary::half_process_dot_residual(sum, code, lut.0, dis_f, delta, norm) + } + } + } + + fn build( + vector: VectBorrowed<'_, f16>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if let Some(centroid) = centroid { + let residual = VectOwned::new(f16::vector_sub(vector.slice(), centroid.slice())); + let code = Self::Vector::code(residual.as_borrowed()); + let delta = { + use std::iter::zip; + let dim = vector.dim(); + let t = zip(&code.1, centroid.slice()) + .map(|(&sign, &num)| std::hint::select_unpredictable(sign, num, -num)) + .map(simd::F16::_to_f32) + .sum::() + / (dim as f32).sqrt(); + let sum_of_x_2 = code.0.dis_u_2; + let sum_of_abs_x = sum_of_x_2 / code.0.factor_ip; + let dis_u = sum_of_x_2.sqrt(); + let x_0 = sum_of_abs_x / dis_u / (dim as f32).sqrt(); + dis_u * t / x_0 - f16::reduce_sum_of_xy(residual.slice(), centroid.slice()) + }; + (code, delta) + } else { + let code = Self::Vector::code(vector); + let delta = 0.0; + (code, delta) + } + } +} + +impl Operator for Op { + type Vector = Rabitq8Owned; + + type DistanceAccessor = ByteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, _delta| { + rabitq::bit::block::half_process_l2s(value, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_l2s(value, code, lut.0) + } + } + + fn build( + vector: Rabitq8Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq8Owned; + + type DistanceAccessor = ByteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, _delta| { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } + } + + fn build( + vector: Rabitq8Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |value, code, _delta| { + rabitq::bit::block::half_process_l2s(value, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let value = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_l2s(value, code, lut.0) + } + } + + fn build( + vector: Rabitq4Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +impl Operator for Op { + type Vector = Rabitq4Owned; + + type DistanceAccessor = HalfbyteDistanceAccessor; + + fn block_access( + lut: &BlockLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl for<'x> Accessor1<[u8; 16], (&'x [[f32; 32]; 4], &'x [f32; 32]), Output = [(f32, f32); 32]> + { + assert!(!is_residual); + RAccess::new( + (&lut.1, ()), + BlockAccessor([0_u32; 32], move |sum, code, _delta| { + rabitq::bit::block::half_process_dot(sum, code, lut.0) + }), + ) + } + + fn binary_access( + lut: &BinaryLut, + is_residual: bool, + _dis_f: f32, + _norm: f32, + ) -> impl FnMut([f32; 4], &[u64], f32) -> (f32, f32) { + assert!(!is_residual); + move |metadata: [f32; 4], elements: &[u64], _delta: f32| { + let sum = rabitq::bit::binary::accumulate(elements, &lut.1); + let code = CodeMetadata { + dis_u_2: metadata[0], + factor_cnt: metadata[1], + factor_ip: metadata[2], + factor_err: metadata[3], + }; + rabitq::bit::binary::half_process_dot(sum, code, lut.0) + } + } + + fn build( + vector: Rabitq4Borrowed<'_>, + centroid: Option, + ) -> (rabitq::bit::Code, f32) { + if centroid.is_some() { + unimplemented!(); + } + (Self::Vector::code(vector), 0.0) + } +} + +pub trait Call { + type Output; + + fn call(&mut self, a: A, b: B, c: C) -> Self::Output; +} + +impl R, R> Call for F { + type Output = R; + + #[inline(always)] + fn call(&mut self, a: A, b: B, c: C) -> R { + (self)(a, b, c) + } +} diff --git a/crates/vchordrq/src/prewarm.rs b/crates/vchordrq/src/prewarm.rs new file mode 100644 index 00000000..02cb9b03 --- /dev/null +++ b/crates/vchordrq/src/prewarm.rs @@ -0,0 +1,124 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::{id_0, id_1, id_2}; +use crate::operator::Operator; +use crate::tape::{by_directory, by_next}; +use crate::tuples::*; +use crate::{Opaque, centroids, tape}; +use index::prefetcher::PrefetcherSequenceFamily; +use index::relation::{Page, RelationRead}; +use index_accessor::FunctionalAccessor; +use std::fmt::Write; + +pub fn prewarm<'b, R: RelationRead, O: Operator>( + index: &'b R, + height: i32, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, +) -> String +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let height_of_root = meta_tuple.height_of_root(); + + let mut message = String::new(); + writeln!(message, "height of root: {height_of_root}").unwrap(); + let prewarm_max_height = if height < 0 { 0 } else { height as u32 }; + if prewarm_max_height > height_of_root { + return message; + } + + type State = Vec; + let mut state: State = { + let mut results = Vec::new(); + let prefetch = meta_tuple.centroid_prefetch().to_vec().into_iter(); + let head = meta_tuple.centroid_head(); + let first = meta_tuple.first(); + centroids::read::(prefetch.map(|id| index.read(id)), head, ()); + results.push(first); + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {}", 1).unwrap(); + results + }; + + drop(meta_guard); + + let mut step = |state: State| { + let mut counter = 0_usize; + let mut results = Vec::new(); + for first in state { + tape::read_h1_tape::( + by_next(index, first).inspect(|_| counter += 1), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), head, _, first, prefetch| { + centroids::read::(prefetch.iter().map(|&id| index.read(id)), head, ()); + results.push(first); + }, + ); + } + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {counter}").unwrap(); + results + }; + + for _ in (std::cmp::max(1, prewarm_max_height)..height_of_root).rev() { + state = step(state); + } + + if prewarm_max_height == 0 { + let mut counter = 0_usize; + let mut results = Vec::new(); + for first in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory).inspect(|_| counter += 1), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + id_2(|_, _, _, _| { + results.push(()); + }), + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()).inspect(|_| counter += 1), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + id_2(|_, _, _, _| { + results.push(()); + }), + ); + } + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()).inspect(|_| counter += 1), + |_, _, _| (), + id_2(|_, _, _, _| { + results.push(()); + }), + ); + } + writeln!(message, "------------------------").unwrap(); + writeln!(message, "number of nodes: {}", results.len()).unwrap(); + writeln!(message, "number of pages: {counter}").unwrap(); + } + + message +} diff --git a/crates/vchordrq/src/rerank.rs b/crates/vchordrq/src/rerank.rs new file mode 100644 index 00000000..08f089ff --- /dev/null +++ b/crates/vchordrq/src/rerank.rs @@ -0,0 +1,137 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::id_4; +use crate::operator::*; +use crate::tuples::{MetaTuple, WithReader}; +use crate::{RerankMethod, vectors}; +use always_equal::AlwaysEqual; +use distance::Distance; +use index::fetch::BorrowedIter; +use index::packed::PackedRefMut; +use index::prefetcher::Prefetcher; +use index::relation::{Page, RelationRead}; +use index_accessor::{Accessor2, DefaultWithDimension, LTryAccess}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::marker::PhantomData; +use std::num::NonZero; +use vector::{VectorBorrowed, VectorOwned}; + +type Result = (Reverse, AlwaysEqual>); + +pub fn how(index: &impl RelationRead) -> RerankMethod { + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let rerank_in_heap = meta_tuple.rerank_in_heap(); + if rerank_in_heap { + RerankMethod::Heap + } else { + RerankMethod::Index + } +} + +pub struct Reranker { + prefetcher: P, + cache: BinaryHeap, + f: F, + _phantom: PhantomData (T, W)>, +} + +impl<'b, T, F, P, W> Iterator for Reranker +where + F: FnMut(NonZero, P::Guards, u16) -> Option, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, +{ + type Item = (Distance, NonZero); + + fn next(&mut self) -> Option { + while let Some(((_, AlwaysEqual(mut w)), prefetch)) = self + .prefetcher + .next_if(|((d, _), ..)| Some(*d) > self.cache.peek().map(|(d, ..)| *d)) + { + let &mut (payload, head, ..) = w.get_mut(); + if let Some(distance) = (self.f)(payload, prefetch, head) { + self.cache.push((Reverse(distance), AlwaysEqual(payload))); + }; + } + let (Reverse(distance), AlwaysEqual(payload)) = self.cache.pop()?; + Some((distance, payload)) + } +} + +impl Reranker { + pub fn finish(self) -> (P, impl Iterator) { + (self.prefetcher, self.cache.into_iter()) + } +} + +pub fn rerank_index< + 'b, + O: Operator, + T, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, +>( + vector: O::Vector, + prefetcher: P, +) -> Reranker, P::Guards, u16) -> Option, P, W> { + let dim = vector.as_borrowed().dim(); + Reranker { + prefetcher, + cache: BinaryHeap::new(), + f: id_4::<_, P, _, _, _>(move |payload, prefetch, head| { + vectors::read::( + prefetch, + head, + payload, + LTryAccess::new( + O::Vector::unpack(vector.as_borrowed()), + O::DistanceAccessor::default_with_dimension(dim), + ), + ) + }), + _phantom: PhantomData, + } +} + +pub fn rerank_heap< + 'b, + O: Operator, + T, + P: Prefetcher<'b, Item = ((Reverse, AlwaysEqual), AlwaysEqual)>, + W: 'b + PackedRefMut, u16, BorrowedIter<'b>)>, +>( + vector: O::Vector, + prefetcher: P, + mut fetch: impl FnMut(NonZero) -> Option + 'b, +) -> Reranker, P::Guards, u16) -> Option, P, W> { + let dim = vector.as_borrowed().dim(); + Reranker { + prefetcher, + cache: BinaryHeap::new(), + f: id_4::<_, P, _, _, _>(move |payload, _, _| { + let unpack = O::Vector::unpack(vector.as_borrowed()); + let vector = fetch(payload)?; + let vector = O::Vector::unpack(vector.as_borrowed()); + let mut accessor = O::DistanceAccessor::default_with_dimension(dim); + accessor.push(unpack.0, vector.0); + let distance = accessor.finish(unpack.1, vector.1); + Some(distance) + }), + _phantom: PhantomData, + } +} diff --git a/crates/vchordrq/src/search.rs b/crates/vchordrq/src/search.rs new file mode 100644 index 00000000..5e5a9e02 --- /dev/null +++ b/crates/vchordrq/src/search.rs @@ -0,0 +1,380 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::closure_lifetime_binder::{id_0, id_1, id_2}; +use crate::linked_vec::LinkedVec; +use crate::operator::*; +use crate::tape::{by_directory, by_next}; +use crate::tuples::*; +use crate::{Opaque, centroids, tape}; +use always_equal::AlwaysEqual; +use distance::Distance; +use index::bump::Bump; +use index::fetch::BorrowedIter; +use index::packed::{PackedRefMut4, PackedRefMut8}; +use index::prefetcher::{Prefetcher, PrefetcherHeapFamily, PrefetcherSequenceFamily}; +use index::relation::{Page, RelationRead}; +use index_accessor::{DefaultWithDimension, FunctionalAccessor, LAccess}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; +use vector::{VectorBorrowed, VectorOwned}; + +type Extra1<'b> = &'b mut (u32, f32, u16, BorrowedIter<'b>); + +pub fn default_search<'b, R: RelationRead, O: Operator>( + index: &'b R, + vector: ::Borrowed<'_>, + probes: Vec, + epsilon: f32, + bump: &'b impl Bump, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, +) -> Vec<( + (Reverse, AlwaysEqual<()>), + AlwaysEqual, u16, BorrowedIter<'b>)>>, +)> +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + let cells = meta_tuple.cells().to_vec(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + if height_of_root as usize != 1 + probes.len() { + panic!( + "usage: need {} probes, but {} probes provided", + height_of_root - 1, + probes.len() + ); + } + debug_assert_eq!(cells[(height_of_root - 1) as usize], 1); + + type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; + let mut state: State = if is_residual { + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); + let head = meta_tuple.centroid_head(); + let distance = centroids::read::( + prefetch.map(|id| index.read(id)), + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + }; + + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + + let mut step = |state: State| { + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + |(rough, err), head, norm, first, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + ))), + )); + }, + ); + } + let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); + let mut cache = BinaryHeap::<(_, _, _)>::new(); + std::iter::from_fn(move || { + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + { + let distance = centroids::read::( + prefetch, + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); + } + cache.pop() + }) + }; + + for i in 1..height_of_root { + let partial_scan = probes[i as usize - 1] < cells[(height_of_root - 1 - i) as usize]; + if partial_scan || is_residual { + state = step(state).take(probes[i as usize - 1] as _).collect(); + } else { + // fast path + let mut results = LinkedVec::new(); + for (Reverse(_), AlwaysEqual(_), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, norm, first, _| { + results.push(( + Reverse(Distance::ZERO), + AlwaysEqual(norm), + AlwaysEqual(first), + )); + }, + ); + } + state = results.into_vec(); + } + } + + let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut callback = id_2(|(rough, err), head, payload, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + results.push(( + (Reverse(lowerbound), AlwaysEqual(())), + AlwaysEqual(PackedRefMut4(bump.alloc(( + payload, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + )))), + )); + }); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + } + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()), + O::binary_access(&lut.1, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + } + results.into_vec() +} + +pub fn maxsim_search<'b, R: RelationRead, O: Operator>( + index: &'b R, + vector: ::Borrowed<'_>, + probes: Vec, + epsilon: f32, + mut threshold: u32, + bump: &'b impl Bump, + mut prefetch_h1_vectors: impl PrefetcherHeapFamily<'b, R>, + mut prefetch_h0_tuples: impl PrefetcherSequenceFamily<'b, R>, +) -> ( + Vec<( + (Reverse, AlwaysEqual), + AlwaysEqual, u16, BorrowedIter<'b>)>>, + )>, + Distance, +) +where + R::Page: Page, +{ + let meta_guard = index.read(0); + let meta_bytes = meta_guard.get(1).expect("data corruption"); + let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); + let dim = meta_tuple.dim(); + let is_residual = meta_tuple.is_residual(); + let height_of_root = meta_tuple.height_of_root(); + let cells = meta_tuple.cells().to_vec(); + assert_eq!(dim, vector.dim(), "unmatched dimensions"); + if height_of_root as usize != 1 + probes.len() { + panic!( + "usage: need {} probes, but {} probes provided", + height_of_root - 1, + probes.len() + ); + } + debug_assert_eq!(cells[(height_of_root - 1) as usize], 1); + + type State = Vec<(Reverse, AlwaysEqual, AlwaysEqual)>; + let mut state: State = if is_residual { + let prefetch = + BorrowedIter::from_slice(meta_tuple.centroid_prefetch(), |x| bump.alloc_slice(x)); + let head = meta_tuple.centroid_head(); + let distance = centroids::read::( + prefetch.map(|id| index.read(id)), + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + } else { + // fast path + let distance = Distance::ZERO; + let norm = meta_tuple.centroid_norm(); + let first = meta_tuple.first(); + vec![(Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))] + }; + + drop(meta_guard); + let lut = O::Vector::preprocess(vector); + + let mut step = |state: State| { + let mut results = LinkedVec::<(_, AlwaysEqual>)>::new(); + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + |(rough, err), head, norm, first, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + results.push(( + Reverse(lowerbound), + AlwaysEqual(bump.alloc(( + first, + norm, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + ))), + )); + }, + ); + } + let mut heap = prefetch_h1_vectors.prefetch(results.into_vec()); + let mut cache = BinaryHeap::<(_, _, _)>::new(); + std::iter::from_fn(move || { + while let Some(((Reverse(_), AlwaysEqual(&mut (first, norm, head, ..))), prefetch)) = + heap.next_if(|(d, _)| Some(*d) > cache.peek().map(|(d, ..)| *d)) + { + let distance = centroids::read::( + prefetch, + head, + LAccess::new( + O::Vector::unpack(vector), + O::DistanceAccessor::default_with_dimension(dim), + ), + ); + cache.push((Reverse(distance), AlwaysEqual(norm), AlwaysEqual(first))); + } + cache.pop() + }) + }; + + let mut it = None; + for i in 1..height_of_root { + let partial_scan = probes[i as usize - 1] < cells[(height_of_root - 1 - i) as usize]; + let needs_sort = i + 1 == height_of_root && threshold != 0; + if partial_scan || is_residual || needs_sort { + let it = it.insert(step(state)); + state = it.take(probes[i as usize - 1] as _).collect(); + } else { + // fast path + let mut results = LinkedVec::new(); + for (Reverse(_), AlwaysEqual(_), AlwaysEqual(first)) in state { + tape::read_h1_tape::( + by_next(index, first), + || FunctionalAccessor::new((), id_0(|_, _| ()), id_1(|_, _| [(); _])), + |(), _, norm, first, _| { + results.push(( + Reverse(Distance::ZERO), + AlwaysEqual(norm), + AlwaysEqual(first), + )); + }, + ); + } + state = results.into_vec(); + } + } + + let mut results = LinkedVec::<(_, AlwaysEqual<_>)>::new(); + for (Reverse(dis_f), AlwaysEqual(norm), AlwaysEqual(first)) in state { + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + let mut callback = id_2(|(rough, err), head, payload, prefetch| { + let lowerbound = Distance::from_f32(rough - err * epsilon); + let rough = Distance::from_f32(rough); + results.push(( + (Reverse(lowerbound), AlwaysEqual(rough)), + AlwaysEqual(PackedRefMut8(bump.alloc(( + payload, + head, + BorrowedIter::from_slice(prefetch, |x| bump.alloc_slice(x)), + )))), + )); + }); + if prefetch_h0_tuples.is_not_plain() { + let directory = + tape::read_directory_tape::(by_next(index, jump_tuple.directory_first())); + tape::read_frozen_tape::( + by_directory(&mut prefetch_h0_tuples, directory), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + } else { + tape::read_frozen_tape::( + by_next(index, jump_tuple.frozen_first()), + || O::block_access(&lut.0, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + } + tape::read_appendable_tape::( + by_next(index, jump_tuple.appendable_first()), + O::binary_access(&lut.1, is_residual, dis_f.to_f32(), norm), + &mut callback, + ); + threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); + } + let mut estimation_by_threshold = Distance::NEG_INFINITY; + for (Reverse(distance), AlwaysEqual(_), AlwaysEqual(first)) in it.into_iter().flatten() { + if threshold == 0 { + break; + } + let jump_guard = index.read(first); + let jump_bytes = jump_guard.get(1).expect("data corruption"); + let jump_tuple = JumpTuple::deserialize_ref(jump_bytes); + threshold = threshold.saturating_sub(jump_tuple.tuples().min(u32::MAX as _) as _); + estimation_by_threshold = distance; + } + (results.into_vec(), estimation_by_threshold) +} diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs new file mode 100644 index 00000000..80d850e2 --- /dev/null +++ b/crates/vchordrq/src/statistics.rs @@ -0,0 +1,30 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::{MetaTuple, WithWriter}; +use index::relation::{Page, RelationWrite}; + +/// Store the number of live vector nodes observed by the latest complete +/// build or vacuum pass. +/// +/// This is deliberately refreshed in bulk instead of on every insert. A +/// per-insert update would serialize all writers on the metapage, which is a +/// poor tradeoff for a planner statistic. Like PostgreSQL's relation +/// statistics, the value may be stale between maintenance passes. +pub fn set_indexed_vectors(index: &R, indexed_vectors: u64) { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + meta_tuple.set_indexed_vectors(indexed_vectors); +} diff --git a/crates/vchordrq/src/tape.rs b/crates/vchordrq/src/tape.rs new file mode 100644 index 00000000..fb0d8878 --- /dev/null +++ b/crates/vchordrq/src/tape.rs @@ -0,0 +1,404 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::*; +use crate::{Opaque, freepages}; +use index::prefetcher::{Prefetcher, PrefetcherSequenceFamily}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::Accessor1; +use std::marker::PhantomData; +use std::num::NonZero; + +pub struct TapeWriter<'a, R, T> +where + R: RelationWrite + 'a, +{ + head: R::WriteGuard<'a>, + first: u32, + index: &'a R, + tracking_freespace: bool, + _phantom: PhantomData T>, +} + +impl<'a, R, T> TapeWriter<'a, R, T> +where + R: RelationWrite + 'a, + R::Page: Page, +{ + pub fn create(index: &'a R, tracking_freespace: bool) -> Self { + let mut head = index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ); + head.get_opaque_mut().skip = head.id(); + let first = head.id(); + Self { + head, + first, + index, + tracking_freespace, + _phantom: PhantomData, + } + } + pub fn first(&self) -> u32 { + self.first + } + pub fn freespace(&self) -> u16 { + self.head.freespace() + } + pub fn tape_move(&mut self) { + if self.head.len() == 0 { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + let next = self.index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + self.tracking_freespace, + ); + self.head.get_opaque_mut().next = next.id(); + self.head = next; + } +} + +impl<'a, R, T> TapeWriter<'a, R, T> +where + R: RelationWrite + 'a, + R::Page: Page, + T: Tuple, +{ + pub fn push(&mut self, x: T) -> (u32, u16) { + let bytes = T::serialize(&x); + if let Some(i) = self.head.alloc(&bytes) { + (self.head.id(), i) + } else { + let next = self.index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + self.tracking_freespace, + ); + self.head.get_opaque_mut().next = next.id(); + self.head = next; + if let Some(i) = self.head.alloc(&bytes) { + (self.head.id(), i) + } else { + panic!("implementation: a free page cannot accommodate a single tuple") + } + } + } + pub fn tape_put(&mut self, x: T) -> (u32, u16) { + let bytes = T::serialize(&x); + if let Some(i) = self.head.alloc(&bytes) { + (self.head.id(), i) + } else { + panic!("implementation: a free page cannot accommodate a single tuple") + } + } +} + +pub fn read_directory_tape<'b, R>( + iter: impl Iterator>, +) -> impl Iterator +where + R: RelationRead + 'b, +{ + use std::pin::Pin; + use std::ptr::NonNull; + + #[pin_project::pin_project] + struct State<'b, R: RelationRead + 'b, I> { + slice: NonNull<[u32]>, + #[pin] + now: Option<(R::ReadGuard<'b>, u16)>, + iter: I, + } + + impl<'b, R: RelationRead + 'b, I: Iterator>> State<'b, R, I> { + fn init(self: Pin<&mut Self>) { + let mut this = self.project(); + let now = this.iter.next().map(|guard| (guard, 0)); + this.now.set(now); + } + + fn next(mut self: Pin<&mut Self>) -> Option { + loop { + let mut this = self.as_mut().project(); + // Safety: If the slice is not empty, the function will return immediately, + // so the guard will not be moved or dropped and the slice remains valid. If + // the slice is empty, a pointer is trivially never dangling, so it's safe + // to use. + #[allow(unsafe_code)] + if let Some((first, more)) = unsafe { this.slice.as_ref() }.split_first() { + *this.slice = more.into(); + return Some(*first); + } + // Safety: `guard` is never moved in this block + #[allow(unsafe_code)] + if let Some((guard, i)) = unsafe { this.now.as_mut().get_unchecked_mut() } { + if *i < guard.len() { + *i += 1; + let bytes = guard.get(*i).expect("data corruption"); + let tuple = DirectoryTuple::deserialize_ref(bytes); + *this.slice = match tuple { + DirectoryTupleReader::_0(tuple) => tuple.elements(), + DirectoryTupleReader::_1(tuple) => tuple.elements(), + } + .into(); + continue; + } + } else { + return None; + } + let now = this.iter.next().map(|guard| (guard, 0)); + this.now.set(now); + } + } + } + + let mut state = Box::pin(State::<'b, R, _> { + slice: NonNull::from(&mut []), + now: None, + iter, + }); + + impl<'b, R: RelationRead + 'b, I: Iterator>> Iterator + for Pin>> + { + type Item = u32; + + fn next(&mut self) -> Option { + self.as_mut().next() + } + } + + state.as_mut().init(); + + state +} + +pub fn by_directory<'b, R>( + p: &mut impl PrefetcherSequenceFamily<'b, R>, + iter: impl Iterator, +) -> impl Iterator> +where + R: RelationRead + 'b, +{ + let mut t = p.prefetch(iter.peekable()); + std::iter::from_fn(move || { + let (_, mut x) = t.next()?; + let ret = x.next().expect("should be at least one element"); + assert!(x.next().is_none(), "should be at most one element"); + Some(ret) + }) +} + +pub fn by_next<'b, R>(index: &'b R, first: u32) -> impl Iterator> +where + R: RelationRead + 'b, + R::Page: Page, +{ + let mut current = first; + std::iter::from_fn(move || { + if current != u32::MAX { + let guard = index.read(current); + current = guard.get_opaque().next; + Some(guard) + } else { + None + } + }) +} + +pub fn read_h1_tape<'b, R, A, T>( + iter: impl Iterator>, + accessor: impl Fn() -> A, + mut callback: impl for<'a> FnMut(T, u16, f32, u32, &'a [u32]), +) where + R: RelationRead + 'b, + A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, +{ + let mut x = None; + for guard in iter { + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = H1Tuple::deserialize_ref(bytes); + match tuple { + H1TupleReader::_0(tuple) => { + let mut x = x.take().unwrap_or_else(&accessor); + x.push(tuple.elements()); + let values = x.finish((tuple.metadata(), tuple.delta())); + let prefetch = tuple.prefetch(); + let flattened = prefetch.as_flattened(); + let step = prefetch.len(); + for (j, value) in values.into_iter().enumerate() { + if j < tuple.len() as usize { + callback( + value, + tuple.head()[j], + tuple.norm()[j], + tuple.first()[j], + &flattened[j * step..][..step], + ); + } + } + } + H1TupleReader::_1(tuple) => { + x.get_or_insert_with(&accessor).push(tuple.elements()); + } + } + } + } +} + +pub fn read_frozen_tape<'b, R, A, T>( + iter: impl Iterator>, + accessor: impl Fn() -> A, + mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), +) where + R: RelationRead + 'b, + A: for<'a> Accessor1<[u8; 16], (&'a [[f32; 32]; 4], &'a [f32; 32]), Output = [T; 32]>, +{ + let mut x = None; + for guard in iter { + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = FrozenTuple::deserialize_ref(bytes); + match tuple { + FrozenTupleReader::_0(tuple) => { + let mut x = x.take().unwrap_or_else(&accessor); + x.push(tuple.elements()); + let values = x.finish((tuple.metadata(), tuple.delta())); + let prefetch = tuple.prefetch(); + let flattened = prefetch.as_flattened(); + let step = prefetch.len(); + for (j, value) in values.into_iter().enumerate() { + if let Some(payload) = tuple.payload()[j] { + callback( + value, + tuple.head()[j], + payload, + &flattened[j * step..][..step], + ); + } + } + } + FrozenTupleReader::_1(tuple) => { + x.get_or_insert_with(&accessor).push(tuple.elements()); + } + } + } + } +} + +pub fn read_appendable_tape<'b, R, T>( + iter: impl Iterator>, + mut access: impl for<'a> FnMut([f32; 4], &'a [u64], f32) -> T, + mut callback: impl for<'a> FnMut(T, u16, NonZero, &'a [u32]), +) where + R: RelationRead + 'b, +{ + for guard in iter { + for i in 1..=guard.len() { + let bytes = guard.get(i).expect("data corruption"); + let tuple = AppendableTuple::deserialize_ref(bytes); + if let Some(payload) = tuple.payload() { + let value = access(tuple.metadata(), tuple.elements(), tuple.delta()); + callback(value, tuple.head(), payload, tuple.prefetch()); + } + } + } +} + +#[allow(clippy::collapsible_else_if)] +pub fn append( + index: &R, + first: u32, + bytes: &[u8], + tracking_freespace: bool, + freepages_first: Option, +) -> (u32, u16) +where + R::Page: Page, +{ + assert!(!tracking_freespace || freepages_first.is_none()); + assert!(first != u32::MAX); + let mut current = first; + loop { + let read = index.read(current); + if read.freespace() as usize >= bytes.len() || read.get_opaque().next == u32::MAX { + drop(read); + let mut write = index.write(current, tracking_freespace); + if write.get_opaque().next == u32::MAX { + if let Some(i) = write.alloc(bytes) { + return (current, i); + } + let mut extend = { + if let Some(freepages_first) = freepages_first { + if let Some(mut guard) = freepages::alloc(index, freepages_first) { + guard.clear(Opaque { + next: u32::MAX, + skip: u32::MAX, + }); + guard + } else { + index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ) + } + } else { + index.extend( + Opaque { + next: u32::MAX, + skip: u32::MAX, + }, + tracking_freespace, + ) + } + }; + write.get_opaque_mut().next = extend.id(); + drop(write); + let fresh = extend.id(); + if let Some(i) = extend.alloc(bytes) { + drop(extend); + let mut past = index.write(first, tracking_freespace); + past.get_opaque_mut().skip = fresh.max(past.get_opaque().skip); + return (fresh, i); + } else { + panic!("implementation: a clear page cannot accommodate a single tuple"); + } + } + if current == first && write.get_opaque().skip != first { + current = write.get_opaque().skip; + } else { + current = write.get_opaque().next; + } + } else { + if current == first && read.get_opaque().skip != first { + current = read.get_opaque().skip; + } else { + current = read.get_opaque().next; + } + } + } +} diff --git a/crates/vchordrq/src/tape_writer.rs b/crates/vchordrq/src/tape_writer.rs new file mode 100644 index 00000000..3bb71a20 --- /dev/null +++ b/crates/vchordrq/src/tape_writer.rs @@ -0,0 +1,261 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tape::TapeWriter; +use crate::tuples::*; +use crate::{Branch, Opaque}; +use index::relation::{Page, RelationWrite}; +use rabitq::packing::{any_pack, padding_pack}; +use std::num::NonZero; + +pub struct DirectoryTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, DirectoryTuple>, +} + +impl<'a, R> DirectoryTapeWriter<'a, R> +where + R: RelationWrite + 'a, + R::Page: Page, +{ + pub fn create(index: &'a R, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + } + } + pub fn push(&mut self, branch: &[u32]) { + let mut remain = branch; + loop { + let freespace = self.tape.freespace(); + if DirectoryTuple::estimate_size_0(remain.len()) <= freespace as usize { + self.tape.tape_put(DirectoryTuple::_0 { + elements: remain.to_vec(), + }); + break; + } + if let Some(w) = DirectoryTuple::fit_1(freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(DirectoryTuple::_1 { + elements: left.to_vec(), + }); + remain = right; + } else { + self.tape.tape_move(); + } + } + } + pub fn into_inner(self) -> TapeWriter<'a, R, DirectoryTuple> { + self.tape + } +} + +pub struct H1TapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, H1Tuple>, + branches: Vec>, + prefetch: usize, +} + +impl<'a, R> H1TapeWriter<'a, R> +where + R: RelationWrite + 'a, + R::Page: Page, +{ + pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + branches: Vec::new(), + prefetch, + } + } + pub fn push(&mut self, branch: Branch) { + self.branches.push(branch); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { + let elements = + padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); + let mut remain = elements.as_slice(); + loop { + let freespace = self.tape.freespace(); + if H1Tuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { + self.tape.tape_put(H1Tuple::_0 { + metadata: [ + chunk.each_ref().map(|x| x.code.0.dis_u_2), + chunk.each_ref().map(|x| x.code.0.factor_cnt), + chunk.each_ref().map(|x| x.code.0.factor_ip), + chunk.each_ref().map(|x| x.code.0.factor_err), + ], + delta: chunk.each_ref().map(|x| x.delta), + prefetch: fix_good(chunk.each_ref().map(|x| x.prefetch.as_slice())), + norm: chunk.each_ref().map(|x| x.norm), + head: chunk.each_ref().map(|x| x.head), + first: chunk.each_ref().map(|x| x.extra), + len: chunk.len() as _, + elements: remain.to_vec(), + }); + break; + } + if let Some(w) = H1Tuple::fit_1(self.prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right; + } else { + self.tape.tape_move(); + } + } + self.branches.clear(); + } + } + pub fn into_inner(self) -> (TapeWriter<'a, R, H1Tuple>, Vec>) { + (self.tape, self.branches) + } + pub fn flush(tape: &mut TapeWriter<'_, R, H1Tuple>, prefetch: usize, chunk: Vec>) { + if chunk.is_empty() { + return; + } + let mut remain = padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); + loop { + let freespace = tape.freespace(); + if H1Tuple::estimate_size_0(prefetch, remain.len()) <= freespace as usize { + tape.tape_put(H1Tuple::_0 { + metadata: [ + any_pack(chunk.iter().map(|x| x.code.0.dis_u_2)), + any_pack(chunk.iter().map(|x| x.code.0.factor_cnt)), + any_pack(chunk.iter().map(|x| x.code.0.factor_ip)), + any_pack(chunk.iter().map(|x| x.code.0.factor_err)), + ], + delta: any_pack(chunk.iter().map(|x| x.delta)), + prefetch: fix_bad(chunk.iter().map(|x| x.prefetch.as_slice())), + head: any_pack(chunk.iter().map(|x| x.head)), + norm: any_pack(chunk.iter().map(|x| x.norm)), + first: any_pack(chunk.iter().map(|x| x.extra)), + len: chunk.len() as _, + elements: remain, + }); + break; + } + if let Some(w) = H1Tuple::fit_1(prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + tape.tape_put(H1Tuple::_1 { + elements: left.to_vec(), + }); + remain = right.to_vec(); + } else { + tape.tape_move(); + } + } + } +} + +pub struct FrozenTapeWriter<'a, R> +where + R: RelationWrite + 'a, +{ + tape: TapeWriter<'a, R, FrozenTuple>, + branches: Vec>>, + prefetch: usize, +} + +impl<'a, R> FrozenTapeWriter<'a, R> +where + R: RelationWrite + 'a, + R::Page: Page, +{ + pub fn create(index: &'a R, prefetch: usize, tracking_freespace: bool) -> Self { + Self { + tape: TapeWriter::create(index, tracking_freespace), + branches: Vec::new(), + prefetch, + } + } + pub fn push(&mut self, branch: Branch>) { + self.branches.push(branch); + if let Ok(chunk) = <&[_; 32]>::try_from(self.branches.as_slice()) { + let elements = + padding_pack(chunk.iter().map(|x| rabitq::packing::pack_to_u4(&x.code.1))); + let mut remain = elements.as_slice(); + loop { + let freespace = self.tape.freespace(); + if FrozenTuple::estimate_size_0(self.prefetch, remain.len()) <= freespace as usize { + self.tape.tape_put(FrozenTuple::_0 { + metadata: [ + chunk.each_ref().map(|x| x.code.0.dis_u_2), + chunk.each_ref().map(|x| x.code.0.factor_cnt), + chunk.each_ref().map(|x| x.code.0.factor_ip), + chunk.each_ref().map(|x| x.code.0.factor_err), + ], + delta: chunk.each_ref().map(|x| x.delta), + prefetch: fix_good(chunk.each_ref().map(|x| x.prefetch.as_slice())), + head: chunk.each_ref().map(|x| x.head), + payload: chunk.each_ref().map(|x| Some(x.extra)), + elements: remain.to_vec(), + }); + break; + } + if let Some(w) = FrozenTuple::fit_1(self.prefetch, freespace) { + let (left, right) = remain.split_at(std::cmp::min(w, remain.len())); + self.tape.tape_put(FrozenTuple::_1 { + elements: left.to_vec(), + }); + remain = right; + } else { + self.tape.tape_move(); + } + } + self.branches.clear(); + } + } + pub fn into_inner(self) -> (TapeWriter<'a, R, FrozenTuple>, Vec>>) { + (self.tape, self.branches) + } +} + +fn fix_good<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + let slices = into_iter.into_iter().collect::>(); + if slices.len() != 32 { + panic!("too many or too few slices"); + } + let min = slices.iter().map(|x| x.len()).min().unwrap_or_default(); + let max = slices.iter().map(|x| x.len()).max().unwrap_or_default(); + if min != max { + panic!("the number of pages for prefetching is not a constant"); + } + let flattened = slices.into_iter().flatten().copied().collect::>(); + let (arrays, remainder) = flattened.as_chunks::<32>(); + assert!(remainder.is_empty()); + arrays.to_vec() +} + +fn fix_bad<'a>(into_iter: impl IntoIterator) -> Vec<[u32; 32]> { + let mut slices = into_iter.into_iter().collect::>(); + if slices.len() > 32 { + panic!("too many slices"); + } + let min = slices.iter().map(|x| x.len()).min().unwrap_or_default(); + let max = slices.iter().map(|x| x.len()).max().unwrap_or_default(); + if min != max { + panic!("the number of pages for prefetching is not a constant"); + } + let temp = vec![u32::MAX; max]; + slices.resize(32, temp.as_slice()); + let flattened = slices.into_iter().flatten().copied().collect::>(); + let (arrays, remainder) = flattened.as_chunks::<32>(); + assert!(remainder.is_empty()); + arrays.to_vec() +} diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs new file mode 100644 index 00000000..85dee786 --- /dev/null +++ b/crates/vchordrq/src/tuples.rs @@ -0,0 +1,1630 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::Vector; +use index::tuples::{Bool, MutChecker, Padding, RefChecker}; +use std::num::NonZero; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +pub const ALIGN: usize = 8; +pub type Tag = u64; +const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); +const VERSION: u64 = 1001; +const STATISTICS_VERSION: u16 = 1; +const MAX_INDEXED_VECTORS: u64 = (1_u64 << 48) - 1; + +#[inline(always)] +fn tag(source: &[u8]) -> Tag { + assert!(source.len() >= size_of::()); + #[allow(unsafe_code)] + unsafe { + source.as_ptr().cast::().read_unaligned() + } +} + +pub trait Tuple: 'static { + fn serialize(&self) -> Vec; +} + +pub trait WithReader: Tuple { + type Reader<'a>; + fn deserialize_ref(source: &[u8]) -> Self::Reader<'_>; +} + +pub trait WithWriter: Tuple { + type Writer<'a>; + fn deserialize_mut(source: &mut [u8]) -> Self::Writer<'_>; +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct MetaTupleHeader { + version: u64, + dim: u32, + height_of_root: u32, + is_residual: Bool, + rerank_in_heap: Bool, + cells_s: u16, + cells_e: u16, + statistics_version: u16, + centroids_first: u32, + vectors_first_s: u16, + vectors_first_e: u16, + freepages_first: u32, + indexed_vectors_low: u32, + indexed_vectors_high: u16, + // tree + centroid_prefetch_s: u16, + centroid_prefetch_e: u16, + centroid_head: u16, + centroid_norm: f32, + first: u32, +} + +// Statistics deliberately replace the old 2-byte and 6-byte padding regions. +// Keep these assertions in non-test builds: changing any offset would require +// an index format version bump and REINDEX instead of the compatibility path. +const _: () = { + assert!(size_of::() == 56); + assert!(std::mem::offset_of!(MetaTupleHeader, statistics_version) == 22); + assert!(std::mem::offset_of!(MetaTupleHeader, centroids_first) == 24); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low) == 36); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high) == 40); + assert!(std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s) == 42); +}; + +pub struct MetaTuple { + pub dim: u32, + pub height_of_root: u32, + pub is_residual: bool, + pub rerank_in_heap: bool, + pub indexed_vectors: Option, + pub cells: Vec, + pub centroids_first: u32, + pub vectors_first: Vec, + pub freepages_first: u32, + pub centroid_prefetch: Vec, + pub centroid_head: u16, + pub centroid_norm: f32, + pub first: u32, +} + +impl Tuple for MetaTuple { + #[allow(clippy::match_single_binding)] + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + MetaTuple { + dim, + height_of_root, + is_residual, + rerank_in_heap, + indexed_vectors, + cells, + centroids_first, + vectors_first, + freepages_first, + centroid_prefetch, + centroid_head, + centroid_norm, + first, + } => { + if let Some(indexed_vectors) = indexed_vectors { + assert!( + *indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + } + buffer.extend((MAGIC as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // cells + let cells_s = buffer.len() as u16; + buffer.extend(cells.as_bytes()); + let cells_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // vectors_first + let vectors_first_s = buffer.len() as u16; + buffer.extend(vectors_first.as_bytes()); + let vectors_first_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // centroid_prefetch + let centroid_prefetch_s = buffer.len() as u16; + buffer.extend(centroid_prefetch.as_bytes()); + let centroid_prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + MetaTupleHeader { + version: VERSION, + dim: *dim, + height_of_root: *height_of_root, + is_residual: (*is_residual).into(), + rerank_in_heap: (*rerank_in_heap).into(), + cells_s, + cells_e, + statistics_version: indexed_vectors + .map(|_| STATISTICS_VERSION) + .unwrap_or(0), + centroids_first: *centroids_first, + vectors_first_s, + vectors_first_e, + freepages_first: *freepages_first, + indexed_vectors_low: indexed_vectors.unwrap_or(0) as u32, + indexed_vectors_high: (indexed_vectors.unwrap_or(0) >> 32) as u16, + centroid_prefetch_s, + centroid_prefetch_e, + centroid_head: *centroid_head, + centroid_norm: *centroid_norm, + first: *first, + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for MetaTuple { + type Reader<'a> = MetaTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> MetaTupleReader<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let checker = RefChecker::new(source); + if VERSION != *checker.prefix::(size_of::()) { + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); + } + let header: &MetaTupleHeader = checker.prefix(size_of::()); + let cells = checker.bytes(header.cells_s, header.cells_e); + let vectors_first = checker.bytes(header.vectors_first_s, header.vectors_first_e); + let centroid_prefetch = + checker.bytes(header.centroid_prefetch_s, header.centroid_prefetch_e); + MetaTupleReader { + header, + cells, + vectors_first, + centroid_prefetch, + } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + if VERSION != header.version { + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); + } + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct MetaTupleReader<'a> { + header: &'a MetaTupleHeader, + cells: &'a [u32], + vectors_first: &'a [u32], + centroid_prefetch: &'a [u32], +} + +impl<'a> MetaTupleReader<'a> { + pub fn dim(self) -> u32 { + self.header.dim + } + pub fn height_of_root(self) -> u32 { + self.header.height_of_root + } + pub fn is_residual(self) -> bool { + self.header.is_residual.into() + } + pub fn rerank_in_heap(self) -> bool { + self.header.rerank_in_heap.into() + } + pub fn indexed_vectors(self) -> Option { + match self.header.statistics_version { + 0 => None, + STATISTICS_VERSION => Some( + u64::from(self.header.indexed_vectors_low) + | (u64::from(self.header.indexed_vectors_high) << 32), + ), + _ => panic!("deserialization: unsupported statistics version"), + } + } + pub fn cells(self) -> &'a [u32] { + self.cells + } + pub fn centroids_first(self) -> u32 { + self.header.centroids_first + } + pub fn vectors_first(self) -> &'a [u32] { + self.vectors_first + } + pub fn freepages_first(self) -> u32 { + self.header.freepages_first + } + pub fn centroid_prefetch(self) -> &'a [u32] { + self.centroid_prefetch + } + pub fn centroid_head(self) -> u16 { + self.header.centroid_head + } + pub fn centroid_norm(self) -> f32 { + self.header.centroid_norm + } + pub fn first(self) -> u32 { + self.header.first + } +} + +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl MetaTupleWriter<'_> { + pub fn set_indexed_vectors(&mut self, indexed_vectors: u64) { + assert!( + indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + self.header.statistics_version = STATISTICS_VERSION; + self.header.indexed_vectors_low = indexed_vectors as u32; + self.header.indexed_vectors_high = (indexed_vectors >> 32) as u16; + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct FreepagesTupleHeader { + first: u32, + _padding_0: [Padding; 4], +} + +#[derive(Debug, Clone)] +pub struct FreepagesTuple {} + +impl Tuple for FreepagesTuple { + fn serialize(&self) -> Vec { + FreepagesTupleHeader { + first: u32::MAX, + _padding_0: Default::default(), + } + .as_bytes() + .to_vec() + } +} + +impl WithWriter for FreepagesTuple { + type Writer<'a> = FreepagesTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> FreepagesTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header = checker.prefix(0_u16); + FreepagesTupleWriter { header } + } +} + +pub struct FreepagesTupleWriter<'a> { + header: &'a mut FreepagesTupleHeader, +} + +impl FreepagesTupleWriter<'_> { + pub fn first(&mut self) -> &mut u32 { + &mut self.header.first + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct CentroidTupleHeader0 { + metadata_s: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct CentroidTupleHeader1 { + head: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[derive(Debug, Clone)] +pub enum CentroidTuple { + _0 { + metadata: V::Metadata, + elements: Vec, + }, + _1 { + head: u16, + elements: Vec, + }, +} + +impl Tuple for CentroidTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + CentroidTuple::_0 { metadata, elements } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // metadata + let metadata_s = buffer.len() as u16; + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + CentroidTupleHeader0 { + metadata_s, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + CentroidTuple::_1 { head, elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + CentroidTupleHeader1 { + head: *head, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for CentroidTuple { + type Reader<'a> = CentroidTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> CentroidTupleReader<'_, V> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &CentroidTupleHeader0 = checker.prefix(size_of::()); + let metadata = checker.prefix(header.metadata_s); + let elements = checker.bytes(header.elements_s, header.elements_e); + CentroidTupleReader::_0(CentroidTupleReader0 { + header, + elements, + metadata, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &CentroidTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + CentroidTupleReader::_1(CentroidTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Clone)] +pub struct CentroidTupleReader0<'a, V: Vector> { + #[allow(dead_code)] + header: &'a CentroidTupleHeader0, + metadata: &'a V::Metadata, + elements: &'a [V::Element], +} + +impl Copy for CentroidTupleReader0<'_, V> {} + +#[derive(Clone)] +pub struct CentroidTupleReader1<'a, V: Vector> { + header: &'a CentroidTupleHeader1, + elements: &'a [V::Element], +} + +impl Copy for CentroidTupleReader1<'_, V> {} + +#[derive(Clone)] +pub enum CentroidTupleReader<'a, V: Vector> { + _0(CentroidTupleReader0<'a, V>), + _1(CentroidTupleReader1<'a, V>), +} + +impl Copy for CentroidTupleReader<'_, V> {} + +impl<'a, V: Vector> CentroidTupleReader<'a, V> { + pub fn elements(self) -> &'a [::Element] { + match self { + CentroidTupleReader::_0(this) => this.elements, + CentroidTupleReader::_1(this) => this.elements, + } + } + pub fn metadata_or_head(self) -> Result { + match self { + CentroidTupleReader::_0(this) => Ok(*this.metadata), + CentroidTupleReader::_1(this) => Err(this.header.head), + } + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader0 { + payload: Option>, + metadata_s: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct VectorTupleHeader1 { + payload: Option>, + head: u16, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 2], +} + +#[derive(Debug, Clone)] +pub enum VectorTuple { + _0 { + payload: Option>, + metadata: V::Metadata, + elements: Vec, + }, + _1 { + payload: Option>, + head: u16, + elements: Vec, + }, +} + +impl Tuple for VectorTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + VectorTuple::_0 { + payload, + metadata, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // metadata + let metadata_s = buffer.len() as u16; + buffer.extend(metadata.as_bytes()); + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader0 { + payload: *payload, + metadata_s, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + VectorTuple::_1 { + payload, + head, + elements, + } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + VectorTupleHeader1 { + payload: *payload, + head: *head, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for VectorTuple { + type Reader<'a> = VectorTupleReader<'a, V>; + + fn deserialize_ref(source: &[u8]) -> VectorTupleReader<'_, V> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader0 = checker.prefix(size_of::()); + let metadata = checker.prefix(header.metadata_s); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_0(VectorTupleReader0 { + header, + elements, + metadata, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VectorTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + VectorTupleReader::_1(VectorTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Clone)] +pub struct VectorTupleReader0<'a, V: Vector> { + header: &'a VectorTupleHeader0, + metadata: &'a V::Metadata, + elements: &'a [V::Element], +} + +impl Copy for VectorTupleReader0<'_, V> {} + +#[derive(Clone)] +pub struct VectorTupleReader1<'a, V: Vector> { + header: &'a VectorTupleHeader1, + elements: &'a [V::Element], +} + +impl Copy for VectorTupleReader1<'_, V> {} + +#[derive(Clone)] +pub enum VectorTupleReader<'a, V: Vector> { + _0(VectorTupleReader0<'a, V>), + _1(VectorTupleReader1<'a, V>), +} + +impl Copy for VectorTupleReader<'_, V> {} + +impl<'a, V: Vector> VectorTupleReader<'a, V> { + pub fn payload(self) -> Option> { + match self { + VectorTupleReader::_0(this) => this.header.payload, + VectorTupleReader::_1(this) => this.header.payload, + } + } + pub fn elements(self) -> &'a [::Element] { + match self { + VectorTupleReader::_0(this) => this.elements, + VectorTupleReader::_1(this) => this.elements, + } + } + pub fn metadata_or_head(self) -> Result { + match self { + VectorTupleReader::_0(this) => Ok(*this.metadata), + VectorTupleReader::_1(this) => Err(this.header.head), + } + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct DirectoryTupleHeader0 { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct DirectoryTupleHeader1 { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], +} + +#[derive(Debug, Clone)] +pub enum DirectoryTuple { + _0 { elements: Vec }, + _1 { elements: Vec }, +} + +impl DirectoryTuple { + pub fn estimate_size_0(elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += (elements * size_of::()).next_multiple_of(ALIGN); + size + } + pub fn fit_1(freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::()) + } else { + None + } + } +} + +impl Tuple for DirectoryTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + Self::_0 { elements } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + DirectoryTupleHeader0 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + Self::_1 { elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + DirectoryTupleHeader1 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for DirectoryTuple { + type Reader<'a> = DirectoryTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> DirectoryTupleReader<'_> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &DirectoryTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + DirectoryTupleReader::_0(DirectoryTupleReader0 { header, elements }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &DirectoryTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + DirectoryTupleReader::_1(DirectoryTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub enum DirectoryTupleReader<'a> { + _0(DirectoryTupleReader0<'a>), + _1(DirectoryTupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy)] +pub struct DirectoryTupleReader0<'a> { + #[allow(dead_code)] + header: &'a DirectoryTupleHeader0, + elements: &'a [u32], +} + +#[derive(Debug, Clone, Copy)] +pub struct DirectoryTupleReader1<'a> { + #[allow(dead_code)] + header: &'a DirectoryTupleHeader1, + elements: &'a [u32], +} + +impl<'a> DirectoryTupleReader0<'a> { + pub fn elements(&self) -> &'a [u32] { + self.elements + } +} + +impl<'a> DirectoryTupleReader1<'a> { + pub fn elements(&self) -> &'a [u32] { + self.elements + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H1TupleHeader0 { + metadata: [[f32; 32]; 4], + delta: [f32; 32], + prefetch_s: u16, + prefetch_e: u16, + head: [u16; 32], + norm: [f32; 32], + first: [u32; 32], + len: u32, + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct H1TupleHeader1 { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone)] +pub enum H1Tuple { + _0 { + metadata: [[f32; 32]; 4], + delta: [f32; 32], + prefetch: Vec<[u32; 32]>, + head: [u16; 32], + norm: [f32; 32], + first: [u32; 32], + len: u32, + elements: Vec<[u8; 16]>, + }, + _1 { + elements: Vec<[u8; 16]>, + }, +} + +impl H1Tuple { + pub fn estimate_size_0(prefetch: usize, elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN); + size += (elements * size_of::<[u8; 16]>()).next_multiple_of(ALIGN); + size + } + pub fn fit_1(prefetch: usize, freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; + freespace &= !(ALIGN - 1) as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u8; 16]>()) + } else { + None + } + } +} + +impl Tuple for H1Tuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + Self::_0 { + head, + norm, + metadata, + delta, + first, + prefetch, + len, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch + let prefetch_s = buffer.len() as u16; + buffer.extend(prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + H1TupleHeader0 { + head: *head, + norm: *norm, + metadata: *metadata, + delta: *delta, + first: *first, + len: *len, + prefetch_s, + prefetch_e, + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + Self::_1 { elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + H1TupleHeader1 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for H1Tuple { + type Reader<'a> = H1TupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> H1TupleReader<'_> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &H1TupleHeader0 = checker.prefix(size_of::()); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); + let elements = checker.bytes(header.elements_s, header.elements_e); + H1TupleReader::_0(H1TupleReader0 { + header, + prefetch, + elements, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &H1TupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + H1TupleReader::_1(H1TupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub enum H1TupleReader<'a> { + _0(H1TupleReader0<'a>), + _1(H1TupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy)] +pub struct H1TupleReader0<'a> { + header: &'a H1TupleHeader0, + prefetch: &'a [[u32; 32]], + elements: &'a [[u8; 16]], +} + +#[derive(Debug, Clone, Copy)] +pub struct H1TupleReader1<'a> { + #[allow(dead_code)] + header: &'a H1TupleHeader1, + elements: &'a [[u8; 16]], +} + +impl<'a> H1TupleReader0<'a> { + pub fn metadata(self) -> &'a [[f32; 32]; 4] { + &self.header.metadata + } + pub fn delta(self) -> &'a [f32; 32] { + &self.header.delta + } + pub fn prefetch(self) -> &'a [[u32; 32]] { + self.prefetch + } + pub fn head(self) -> &'a [u16; 32] { + &self.header.head + } + pub fn norm(self) -> &'a [f32; 32] { + &self.header.norm + } + pub fn first(self) -> &'a [u32; 32] { + &self.header.first + } + pub fn len(self) -> u32 { + self.header.len + } + pub fn elements(&self) -> &'a [[u8; 16]] { + self.elements + } +} + +impl<'a> H1TupleReader1<'a> { + pub fn elements(&self) -> &'a [[u8; 16]] { + self.elements + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct JumpTupleHeader { + centroid_prefetch_s: u16, + centroid_prefetch_e: u16, + centroid_head: u16, + _padding_0: [Padding; 6], + directory_first: u32, + frozen_first: u32, + appendable_first: u32, + tuples: u64, +} + +#[derive(Debug, Clone)] +pub struct JumpTuple { + pub centroid_prefetch: Vec, + pub centroid_head: u16, + pub directory_first: u32, + pub frozen_first: u32, + pub appendable_first: u32, + pub tuples: u64, +} + +impl Tuple for JumpTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // centroid_prefetch + let centroid_prefetch_s = buffer.len() as u16; + buffer.extend(self.centroid_prefetch.as_bytes()); + let centroid_prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[..size_of::()].copy_from_slice( + JumpTupleHeader { + centroid_prefetch_s, + centroid_prefetch_e, + centroid_head: self.centroid_head, + directory_first: self.directory_first, + frozen_first: self.frozen_first, + appendable_first: self.appendable_first, + tuples: self.tuples, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + } +} + +impl WithReader for JumpTuple { + type Reader<'a> = JumpTupleReader<'a>; + fn deserialize_ref(source: &[u8]) -> JumpTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &JumpTupleHeader = checker.prefix(0_u16); + let centroid_prefetch = + checker.bytes(header.centroid_prefetch_s, header.centroid_prefetch_e); + JumpTupleReader { + header, + centroid_prefetch, + } + } +} + +impl WithWriter for JumpTuple { + type Writer<'a> = JumpTupleWriter<'a>; + fn deserialize_mut(source: &mut [u8]) -> JumpTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut JumpTupleHeader = checker.prefix(0_u16); + JumpTupleWriter { header } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct JumpTupleReader<'a> { + header: &'a JumpTupleHeader, + centroid_prefetch: &'a [u32], +} + +impl<'a> JumpTupleReader<'a> { + pub fn centroid_prefetch(self) -> &'a [u32] { + self.centroid_prefetch + } + pub fn centroid_head(self) -> u16 { + self.header.centroid_head + } + pub fn directory_first(self) -> u32 { + self.header.directory_first + } + pub fn frozen_first(self) -> u32 { + self.header.frozen_first + } + pub fn appendable_first(self) -> u32 { + self.header.appendable_first + } + pub fn tuples(self) -> u64 { + self.header.tuples + } +} + +#[derive(Debug)] +pub struct JumpTupleWriter<'a> { + header: &'a mut JumpTupleHeader, +} + +impl JumpTupleWriter<'_> { + pub fn directory_first(&mut self) -> &mut u32 { + &mut self.header.directory_first + } + pub fn frozen_first(&mut self) -> &mut u32 { + &mut self.header.frozen_first + } + pub fn appendable_first(&mut self) -> &mut u32 { + &mut self.header.appendable_first + } + pub fn tuples(&mut self) -> &mut u64 { + &mut self.header.tuples + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct FrozenTupleHeader0 { + metadata: [[f32; 32]; 4], + delta: [f32; 32], + // it's not last field for reducing padding bytes + payload: [Option>; 32], + prefetch_s: u16, + prefetch_e: u16, + head: [u16; 32], + elements_s: u16, + elements_e: u16, +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct FrozenTupleHeader1 { + elements_s: u16, + elements_e: u16, + _padding_0: [Padding; 4], +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone)] +pub enum FrozenTuple { + _0 { + metadata: [[f32; 32]; 4], + delta: [f32; 32], + payload: [Option>; 32], + prefetch: Vec<[u32; 32]>, + head: [u16; 32], + elements: Vec<[u8; 16]>, + }, + _1 { + elements: Vec<[u8; 16]>, + }, +} + +impl FrozenTuple { + pub fn estimate_size_0(prefetch: usize, elements: usize) -> usize { + let mut size = 0_usize; + size += size_of::(); + size += size_of::(); + size += (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN); + size += (elements * size_of::<[u8; 16]>()).next_multiple_of(ALIGN); + size + } + pub fn fit_1(prefetch: usize, freespace: u16) -> Option { + let mut freespace = freespace as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= size_of::() as isize; + freespace &= !(ALIGN - 1) as isize; + freespace -= (prefetch * size_of::<[u32; 32]>()).next_multiple_of(ALIGN) as isize; + freespace &= !(ALIGN - 1) as isize; + if freespace >= 0 { + Some(freespace as usize / size_of::<[u8; 16]>()) + } else { + None + } + } +} + +impl Tuple for FrozenTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + match self { + FrozenTuple::_0 { + metadata, + delta, + payload, + prefetch, + head, + elements, + } => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch + let prefetch_s = buffer.len() as u16; + buffer.extend(prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + FrozenTupleHeader0 { + head: *head, + metadata: *metadata, + delta: *delta, + payload: *payload, + elements_s, + elements_e, + prefetch_s, + prefetch_e, + } + .as_bytes(), + ); + } + Self::_1 { elements } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // elements + let elements_s = buffer.len() as u16; + buffer.extend(elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[size_of::()..][..size_of::()].copy_from_slice( + FrozenTupleHeader1 { + elements_s, + elements_e, + _padding_0: Default::default(), + } + .as_bytes(), + ); + } + } + buffer + } +} + +impl WithReader for FrozenTuple { + type Reader<'a> = FrozenTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> FrozenTupleReader<'_> { + let tag = tag(source); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &FrozenTupleHeader0 = checker.prefix(size_of::()); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleReader::_0(FrozenTupleReader0 { + header, + prefetch, + elements, + }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &FrozenTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleReader::_1(FrozenTupleReader1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +impl WithWriter for FrozenTuple { + type Writer<'a> = FrozenTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> FrozenTupleWriter<'_> { + let tag = tag(source); + match tag { + 0 => { + let mut checker = MutChecker::new(source); + let header: &mut FrozenTupleHeader0 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleWriter::_0(FrozenTupleWriter0 { header, elements }) + } + 1 => { + let mut checker = MutChecker::new(source); + let header: &mut FrozenTupleHeader1 = checker.prefix(size_of::()); + let elements = checker.bytes(header.elements_s, header.elements_e); + FrozenTupleWriter::_1(FrozenTupleWriter1 { header, elements }) + } + _ => panic!("deserialization: bad bytes"), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub enum FrozenTupleReader<'a> { + _0(FrozenTupleReader0<'a>), + _1(FrozenTupleReader1<'a>), +} + +#[derive(Debug, Clone, Copy)] +pub struct FrozenTupleReader0<'a> { + header: &'a FrozenTupleHeader0, + prefetch: &'a [[u32; 32]], + elements: &'a [[u8; 16]], +} + +impl<'a> FrozenTupleReader0<'a> { + pub fn metadata(self) -> &'a [[f32; 32]; 4] { + &self.header.metadata + } + pub fn delta(self) -> &'a [f32; 32] { + &self.header.delta + } + pub fn payload(self) -> &'a [Option>; 32] { + &self.header.payload + } + pub fn prefetch(self) -> &'a [[u32; 32]] { + self.prefetch + } + pub fn head(self) -> &'a [u16; 32] { + &self.header.head + } + pub fn elements(self) -> &'a [[u8; 16]] { + self.elements + } +} + +#[derive(Debug, Clone, Copy)] +pub struct FrozenTupleReader1<'a> { + #[allow(dead_code)] + header: &'a FrozenTupleHeader1, + elements: &'a [[u8; 16]], +} + +impl<'a> FrozenTupleReader1<'a> { + pub fn elements(self) -> &'a [[u8; 16]] { + self.elements + } +} + +#[derive(Debug)] +pub enum FrozenTupleWriter<'a> { + _0(FrozenTupleWriter0<'a>), + #[allow(dead_code)] + _1(FrozenTupleWriter1<'a>), +} + +#[derive(Debug)] +pub struct FrozenTupleWriter0<'a> { + header: &'a mut FrozenTupleHeader0, + #[allow(dead_code)] + elements: &'a mut [[u8; 16]], +} + +#[derive(Debug)] +pub struct FrozenTupleWriter1<'a> { + #[allow(dead_code)] + header: &'a mut FrozenTupleHeader1, + #[allow(dead_code)] + elements: &'a mut [[u8; 16]], +} + +impl FrozenTupleWriter0<'_> { + pub fn payload(&mut self) -> &mut [Option>; 32] { + &mut self.header.payload + } +} + +#[repr(C, align(8))] +#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] +struct AppendableTupleHeader { + metadata: [f32; 4], + delta: f32, + prefetch_s: u16, + prefetch_e: u16, + head: u16, + _padding_0: [Padding; 2], + elements_s: u16, + elements_e: u16, + // it's the last field for reducing padding bytes + payload: Option>, +} + +#[derive(Debug, Clone)] +pub struct AppendableTuple { + pub metadata: [f32; 4], + pub delta: f32, + pub prefetch: Vec, + pub head: u16, + pub elements: Vec, + pub payload: Option>, +} + +impl Tuple for AppendableTuple { + fn serialize(&self) -> Vec { + let mut buffer = Vec::::new(); + buffer.extend(std::iter::repeat_n(0, size_of::())); + // prefetch + let prefetch_s = buffer.len() as u16; + buffer.extend(self.prefetch.as_bytes()); + let prefetch_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // elements + let elements_s = buffer.len() as u16; + buffer.extend(self.elements.as_bytes()); + let elements_e = buffer.len() as u16; + while buffer.len() % ALIGN != 0 { + buffer.push(0); + } + // header + buffer[..size_of::()].copy_from_slice( + AppendableTupleHeader { + metadata: self.metadata, + delta: self.delta, + prefetch_s, + prefetch_e, + head: self.head, + elements_s, + elements_e, + payload: self.payload, + _padding_0: Default::default(), + } + .as_bytes(), + ); + buffer + } +} + +impl WithReader for AppendableTuple { + type Reader<'a> = AppendableTupleReader<'a>; + + fn deserialize_ref(source: &[u8]) -> AppendableTupleReader<'_> { + let checker = RefChecker::new(source); + let header: &AppendableTupleHeader = checker.prefix(0_u16); + let prefetch = checker.bytes(header.prefetch_s, header.prefetch_e); + let elements = checker.bytes(header.elements_s, header.elements_e); + AppendableTupleReader { + header, + prefetch, + elements, + } + } +} + +impl WithWriter for AppendableTuple { + type Writer<'a> = AppendableTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> AppendableTupleWriter<'_> { + let mut checker = MutChecker::new(source); + let header: &mut AppendableTupleHeader = checker.prefix(0_u16); + let elements = checker.bytes(header.elements_s, header.elements_e); + AppendableTupleWriter { header, elements } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct AppendableTupleReader<'a> { + header: &'a AppendableTupleHeader, + prefetch: &'a [u32], + elements: &'a [u64], +} + +impl<'a> AppendableTupleReader<'a> { + pub fn metadata(self) -> [f32; 4] { + self.header.metadata + } + pub fn delta(self) -> f32 { + self.header.delta + } + pub fn prefetch(self) -> &'a [u32] { + self.prefetch + } + pub fn head(self) -> u16 { + self.header.head + } + pub fn payload(self) -> Option> { + self.header.payload + } + pub fn elements(self) -> &'a [u64] { + self.elements + } +} + +#[derive(Debug)] +pub struct AppendableTupleWriter<'a> { + header: &'a mut AppendableTupleHeader, + #[allow(dead_code)] + elements: &'a mut [u64], +} + +impl AppendableTupleWriter<'_> { + pub fn payload(&mut self) -> &mut Option> { + &mut self.header.payload + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn meta(indexed_vectors: Option) -> MetaTuple { + MetaTuple { + dim: 3, + height_of_root: 1, + is_residual: false, + rerank_in_heap: false, + indexed_vectors, + cells: vec![4], + centroids_first: 1, + vectors_first: vec![2], + freepages_first: 3, + centroid_prefetch: vec![4], + centroid_head: 0, + centroid_norm: 1.0, + first: 5, + } + } + + #[test] + fn meta_statistics_reuse_padding_without_moving_fields() { + assert_eq!(size_of::(), 56); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, statistics_version), + 22 + ); + assert_eq!(std::mem::offset_of!(MetaTupleHeader, centroids_first), 24); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low), + 36 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high), + 40 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s), + 42 + ); + } + + #[test] + fn meta_statistics_roundtrip_full_48_bit_range() { + for expected in [0, 1, u32::MAX as u64 + 1, MAX_INDEXED_VECTORS] { + let bytes = meta(Some(expected)).serialize(); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(expected) + ); + } + } + + #[test] + fn old_meta_padding_reads_as_missing_statistics_and_can_be_upgraded() { + let mut bytes = meta(None).serialize(); + assert_eq!(MetaTuple::deserialize_ref(&bytes).indexed_vectors(), None); + + MetaTuple::deserialize_mut(&mut bytes).set_indexed_vectors(42); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(42) + ); + } +} diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs new file mode 100644 index 00000000..d9b85cbc --- /dev/null +++ b/crates/vchordrq/src/types.rs @@ -0,0 +1,163 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use distance::Distance; +use serde::{Deserialize, Serialize}; +use simd::f16; +use validator::{Validate, ValidationError}; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; +use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqIndexOptions { + #[serde(default = "VchordrqIndexOptions::default_residual_quantization")] + pub residual_quantization: bool, + #[serde(default = "VchordrqIndexOptions::default_rerank_in_table")] + pub rerank_in_table: bool, + #[serde(default = "VchordrqIndexOptions::default_degree_of_parallelism")] + #[validate(range(min = 1, max = 256))] + pub degree_of_parallelism: u32, +} + +impl VchordrqIndexOptions { + fn default_residual_quantization() -> bool { + false + } + fn default_rerank_in_table() -> bool { + false + } + fn default_degree_of_parallelism() -> u32 { + 32 + } +} + +impl Default for VchordrqIndexOptions { + fn default() -> Self { + Self { + residual_quantization: Self::default_residual_quantization(), + rerank_in_table: Self::default_rerank_in_table(), + degree_of_parallelism: Self::default_degree_of_parallelism(), + } + } +} + +#[derive(Debug, Clone)] +pub enum OwnedVector { + Vecf32(VectOwned), + Vecf16(VectOwned), + Rabitq8(Rabitq8Owned), + Rabitq4(Rabitq4Owned), +} + +impl OwnedVector { + pub fn dim(&self) -> u32 { + match self { + Self::Vecf32(vector) => vector.as_borrowed().dim(), + Self::Vecf16(vector) => vector.as_borrowed().dim(), + Self::Rabitq8(vector) => vector.as_borrowed().dim(), + Self::Rabitq4(vector) => vector.as_borrowed().dim(), + } + } + + pub fn operator_dot(&self, rhs: &Self) -> Option { + if self.dim() != rhs.dim() { + return None; + } + match (self, rhs) { + (Self::Vecf32(lhs), Self::Vecf32(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Vecf16(lhs), Self::Vecf16(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq8(lhs), Self::Rabitq8(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq4(lhs), Self::Rabitq4(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + _ => None, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub enum BorrowedVector<'a> { + Vecf32(VectBorrowed<'a, f32>), + Vecf16(VectBorrowed<'a, f16>), + Rabitq8(Rabitq8Borrowed<'a>), + Rabitq4(Rabitq4Borrowed<'a>), +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum DistanceKind { + L2S, + Dot, +} + +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum VectorKind { + Vecf32, + Vecf16, + Rabitq8, + Rabitq4, +} + +impl VectorKind { + pub fn number_of_bits_of_an_elements(self) -> u32 { + match self { + VectorKind::Vecf32 => 32, + VectorKind::Vecf16 => 16, + VectorKind::Rabitq8 => 8, + VectorKind::Rabitq4 => 8, + } + } +} + +#[derive(Debug, Clone, Copy, Validate)] +#[validate(schema(function = "Self::validate_self"))] +pub struct VectorOptions { + #[validate(range(min = 1))] + pub dim: u32, + pub v: VectorKind, + pub d: DistanceKind, +} + +impl VectorOptions { + pub fn validate_self(&self) -> Result<(), ValidationError> { + match (self.v, self.d, self.dim) { + (_, _, 1..=60000) => Ok(()), + _ => Err(ValidationError::new("invalid vector options")), + } + } +} + +pub struct Structure { + pub centroids: Vec, + pub children: Vec>, +} + +impl Structure { + pub fn len(&self) -> usize { + self.children.len() + } + pub fn is_empty(&self) -> bool { + self.children.is_empty() + } +} diff --git a/crates/vchordrq/src/vectors.rs b/crates/vchordrq/src/vectors.rs new file mode 100644 index 00000000..3531efa2 --- /dev/null +++ b/crates/vchordrq/src/vectors.rs @@ -0,0 +1,107 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::operator::*; +use crate::tuples::*; +use crate::{Opaque, tape}; +use index::relation::{Page, PageGuard, RelationRead, RelationWrite}; +use index_accessor::TryAccessor1; +use std::num::NonZero; +use vector::VectorOwned; + +pub fn read< + 'a, + R: RelationRead + 'a, + O: Operator, + A: TryAccessor1<::Element, ::Metadata>, +>( + mut prefetch: impl Iterator>, + head: u16, + payload: NonZero, + accessor: A, +) -> Option { + let mut cursor = Err(head); + let mut result = accessor; + while let Err(head) = cursor { + let guard = prefetch.next()?; + let bytes = guard.get(head)?; + let tuple = VectorTuple::::deserialize_ref(bytes); + if tuple.payload().is_none() { + panic!("data corruption"); + } + if tuple.payload() != Some(payload) { + return None; + } + result.push(tuple.elements())?; + cursor = tuple.metadata_or_head(); + } + if prefetch.next().is_some() { + return None; + } + result.finish(cursor.ok()?) +} + +pub fn append( + index: &R, + vectors_first: u32, + vector: ::Borrowed<'_>, + payload: NonZero, + skip_search: bool, +) -> (Vec, u16) +where + R::Page: Page, +{ + fn append( + index: &R, + first: u32, + bytes: &[u8], + skip_search: bool, + ) -> (u32, u16) + where + R::Page: Page, + { + if !skip_search && let Some(mut write) = index.search(bytes.len()) { + let i = write + .alloc(bytes) + .expect("implementation: a free page cannot accommodate a single tuple"); + return (write.id(), i); + } + tape::append(index, first, bytes, true, None) + } + let (slices, metadata) = O::Vector::split(vector); + let mut chain = Ok(metadata); + let mut prefetch = Vec::new(); + for i in (0..slices.len()).rev() { + let bytes = VectorTuple::::serialize(&match chain { + Ok(metadata) => VectorTuple::_0 { + elements: slices[i].to_vec(), + payload: Some(payload), + metadata, + }, + Err(head) => VectorTuple::_1 { + elements: slices[i].to_vec(), + payload: Some(payload), + head, + }, + }); + let (id, head) = append(index, vectors_first, &bytes, skip_search); + chain = Err(head); + prefetch.push(id); + } + prefetch.reverse(); + ( + prefetch, + chain.expect_err("internal error: 0-dimensional vector"), + ) +} diff --git a/crates/vector/Cargo.toml b/crates/vector/Cargo.toml new file mode 100644 index 00000000..d68f1f92 --- /dev/null +++ b/crates/vector/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "vector" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +distance = { path = "../distance" } +rabitq = { path = "../rabitq" } +simd = { path = "../simd" } + +[lints] +workspace = true diff --git a/crates/vector/src/bvect.rs b/crates/vector/src/bvect.rs new file mode 100644 index 00000000..716d26ea --- /dev/null +++ b/crates/vector/src/bvect.rs @@ -0,0 +1,256 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; + +pub const BVECTOR_WIDTH: u32 = u64::BITS; + +// When using binary vector, please ensure that the padding bits are always zero. +#[derive(Debug, Clone)] +pub struct BVectOwned { + dim: u32, + data: Vec, +} + +impl BVectOwned { + #[inline(always)] + pub fn new(dim: u32, data: Vec) -> Self { + Self::new_checked(dim, data).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dim: u32, data: Vec) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if data.len() != dim.div_ceil(BVECTOR_WIDTH) as usize { + return None; + } + if dim % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dim % BVECTOR_WIDTH) != 0 { + return None; + } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dim, data)) + } + } + + /// # Safety + /// + /// * `dim` must be in `1..=65535`. + /// * `data` must be of the correct length. + /// * The padding bits must be zero. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked(dim: u32, data: Vec) -> Self { + Self { dim, data } + } +} + +impl VectorOwned for BVectOwned { + type Borrowed<'a> = BVectBorrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> BVectBorrowed<'_> { + BVectBorrowed { + dim: self.dim, + data: &self.data, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct BVectBorrowed<'a> { + dim: u32, + data: &'a [u64], +} + +impl<'a> BVectBorrowed<'a> { + #[inline(always)] + pub fn new(dim: u32, data: &'a [u64]) -> Self { + Self::new_checked(dim, data).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dim: u32, data: &'a [u64]) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if data.len() != dim.div_ceil(BVECTOR_WIDTH) as usize { + return None; + } + if dim % BVECTOR_WIDTH != 0 && data[data.len() - 1] >> (dim % BVECTOR_WIDTH) != 0 { + return None; + } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dim, data)) + } + } + + /// # Safety + /// + /// * `dim` must be in `1..=65535`. + /// * `data` must be of the correct length. + /// * The padding bits must be zero. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked(dim: u32, data: &'a [u64]) -> Self { + Self { dim, data } + } + + #[inline(always)] + pub fn data(&self) -> &'a [u64] { + self.data + } + + #[inline(always)] + pub fn get(&self, index: u32) -> bool { + assert!(index < self.dim); + self.data[(index / BVECTOR_WIDTH) as usize] & (1 << (index % BVECTOR_WIDTH)) != 0 + } + + #[inline(always)] + pub fn iter(self) -> impl Iterator + 'a { + let mut index = 0_u32; + std::iter::from_fn(move || { + if index < self.dim { + let result = self.data[(index / BVECTOR_WIDTH) as usize] + & (1 << (index % BVECTOR_WIDTH)) + != 0; + index += 1; + Some(result) + } else { + None + } + }) + } +} + +impl VectorBorrowed for BVectBorrowed<'_> { + type Owned = BVectOwned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.dim + } + + fn own(&self) -> BVectOwned { + BVectOwned { + dim: self.dim, + data: self.data.to_vec(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + (simd::bit::reduce_sum_of_x(self.data) as f32).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + Distance::from(-(simd::bit::reduce_sum_of_and(self.data, rhs.data) as f32)) + } + + #[inline(always)] + fn operator_l2s(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_cos(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_hamming(self, rhs: Self) -> Distance { + Distance::from(simd::bit::reduce_sum_of_xor(self.data, rhs.data) as f32) + } + + #[inline(always)] + fn operator_jaccard(self, rhs: Self) -> Distance { + let (and, or) = simd::bit::reduce_sum_of_and_or(self.data, rhs.data); + Distance::from(1.0 - (and as f32 / or as f32)) + } + + #[inline(always)] + fn function_normalize(&self) -> BVectOwned { + unimplemented!() + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let data = simd::bit::vector_and(self.data, rhs.data); + BVectOwned::new(self.dim, data) + } + + fn operator_or(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let data = simd::bit::vector_or(self.data, rhs.data); + BVectOwned::new(self.dim, data) + } + + fn operator_xor(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let data = simd::bit::vector_xor(self.data, rhs.data); + BVectOwned::new(self.dim, data) + } +} + +impl PartialEq for BVectBorrowed<'_> { + fn eq(&self, other: &Self) -> bool { + if self.dim != other.dim { + return false; + } + for (&l, &r) in self.data.iter().zip(other.data.iter()) { + let l = l.reverse_bits(); + let r = r.reverse_bits(); + if l != r { + return false; + } + } + true + } +} + +impl PartialOrd for BVectBorrowed<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + use std::cmp::Ordering; + if self.dim != other.dim { + return None; + } + for (&l, &r) in self.data.iter().zip(other.data.iter()) { + let l = l.reverse_bits(); + let r = r.reverse_bits(); + match l.cmp(&r) { + Ordering::Equal => (), + x => return Some(x), + } + } + Some(Ordering::Equal) + } +} diff --git a/crates/vector/src/lib.rs b/crates/vector/src/lib.rs new file mode 100644 index 00000000..07301250 --- /dev/null +++ b/crates/vector/src/lib.rs @@ -0,0 +1,59 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod bvect; +pub mod rabitq4; +pub mod rabitq8; +pub mod svect; +pub mod vect; + +pub trait VectorOwned: Clone + 'static { + type Borrowed<'a>: VectorBorrowed; + + fn as_borrowed(&self) -> Self::Borrowed<'_>; +} + +pub trait VectorBorrowed: Copy { + type Owned: VectorOwned; + + fn own(&self) -> Self::Owned; + + fn dim(&self) -> u32; + + fn norm(&self) -> f32; + + fn operator_dot(self, rhs: Self) -> distance::Distance; + + fn operator_l2s(self, rhs: Self) -> distance::Distance; + + fn operator_cos(self, rhs: Self) -> distance::Distance; + + fn operator_hamming(self, rhs: Self) -> distance::Distance; + + fn operator_jaccard(self, rhs: Self) -> distance::Distance; + + fn function_normalize(&self) -> Self::Owned; + + fn operator_add(&self, rhs: Self) -> Self::Owned; + + fn operator_sub(&self, rhs: Self) -> Self::Owned; + + fn operator_mul(&self, rhs: Self) -> Self::Owned; + + fn operator_and(&self, rhs: Self) -> Self::Owned; + + fn operator_or(&self, rhs: Self) -> Self::Owned; + + fn operator_xor(&self, rhs: Self) -> Self::Owned; +} diff --git a/crates/vector/src/rabitq4.rs b/crates/vector/src/rabitq4.rs new file mode 100644 index 00000000..4010c786 --- /dev/null +++ b/crates/vector/src/rabitq4.rs @@ -0,0 +1,384 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; + +#[derive(Debug, Clone)] +pub struct Rabitq4Owned { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, +} + +impl Rabitq4Owned { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(2) as usize != packed_code.len() { + return None; + } + if dim % 2 == 1 && packed_code.last().copied().unwrap_or_default() >> 4 != 0 { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to 1/2 of `code.len()`, rounding to infinity. + /// * `packed_code` is filled with zero bits correctly. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } +} + +impl VectorOwned for Rabitq4Owned { + type Borrowed<'a> = Rabitq4Borrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + Rabitq4Borrowed { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.as_slice(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Rabitq4Borrowed<'a> { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], +} + +impl<'a> Rabitq4Borrowed<'a> { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Option { + if !(1..=65535).contains(&packed_code.len()) { + return None; + } + if dim.div_ceil(2) as usize != packed_code.len() { + return None; + } + if dim % 2 == 1 && packed_code.last().copied().unwrap_or_default() >> 4 != 0 { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to 1/2 of `code.len()`, rounding to infinity. + /// * `packed_code` is filled with zero bits correctly. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } + + #[inline(always)] + pub fn sum_of_x2(&self) -> f32 { + self.sum_of_x2 + } + + #[inline(always)] + pub fn norm_of_lattice(&self) -> f32 { + self.norm_of_lattice + } + + #[inline(always)] + pub fn sum_of_code(&self) -> f32 { + self.sum_of_code + } + + #[inline(always)] + pub fn sum_of_abs_x(&self) -> f32 { + self.sum_of_abs_x + } + + #[inline(always)] + pub fn packed_code(&self) -> &'a [u8] { + self.packed_code + } + + #[inline(always)] + pub fn unpacked_code(&self) -> impl Iterator { + self.packed_code + .iter() + .flat_map(|x| [x & 0xf, x >> 4]) + .take(self.dim as _) + } +} + +impl VectorBorrowed for Rabitq4Borrowed<'_> { + type Owned = Rabitq4Owned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.dim + } + + #[inline(always)] + fn own(&self) -> Rabitq4Owned { + Rabitq4Owned { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.to_owned(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + self.sum_of_x2.sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_dot( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_l2s( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::halfbyte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::halfbyte::binary::half_process_cos( + dim, + sum, + rabitq::halfbyte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::halfbyte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> Rabitq4Owned { + Rabitq4Owned { + dim: self.dim, + sum_of_x2: 1.0, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x / self.sum_of_x2.sqrt(), + packed_code: self.packed_code.to_owned(), + } + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} diff --git a/crates/vector/src/rabitq8.rs b/crates/vector/src/rabitq8.rs new file mode 100644 index 00000000..d02a40c5 --- /dev/null +++ b/crates/vector/src/rabitq8.rs @@ -0,0 +1,373 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; + +#[derive(Debug, Clone)] +pub struct Rabitq8Owned { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, +} + +impl Rabitq8Owned { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(1) as usize != packed_code.len() { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to `code.len()`. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: Vec, + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } +} + +impl VectorOwned for Rabitq8Owned { + type Borrowed<'a> = Rabitq8Borrowed<'a>; + + #[inline(always)] + fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + Rabitq8Borrowed { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.as_slice(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct Rabitq8Borrowed<'a> { + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], +} + +impl<'a> Rabitq8Borrowed<'a> { + #[inline(always)] + pub fn new( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + .expect("invalid data") + } + + #[inline(always)] + pub fn new_checked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Option { + if !(1..=65535).contains(&dim) { + return None; + } + if dim.div_ceil(1) as usize != packed_code.len() { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { + Self::new_unchecked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + ) + }) + } + + /// # Safety + /// + /// * `dim` must not be zero. + /// * `dim` must be less than 65536. + /// * `dim` must be equal to `code.len()`. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked( + dim: u32, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + packed_code: &'a [u8], + ) -> Self { + Self { + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + packed_code, + } + } + + #[inline(always)] + pub fn sum_of_x2(&self) -> f32 { + self.sum_of_x2 + } + + #[inline(always)] + pub fn norm_of_lattice(&self) -> f32 { + self.norm_of_lattice + } + + #[inline(always)] + pub fn sum_of_code(&self) -> f32 { + self.sum_of_code + } + + #[inline(always)] + pub fn sum_of_abs_x(&self) -> f32 { + self.sum_of_abs_x + } + + #[inline(always)] + pub fn packed_code(&self) -> &'a [u8] { + self.packed_code + } + + #[inline(always)] + pub fn unpacked_code(&self) -> impl Iterator { + self.packed_code.iter().copied() + } +} + +impl VectorBorrowed for Rabitq8Borrowed<'_> { + type Owned = Rabitq8Owned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.dim + } + + #[inline(always)] + fn own(&self) -> Rabitq8Owned { + Rabitq8Owned { + dim: self.dim, + sum_of_x2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x, + packed_code: self.packed_code.to_owned(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + self.sum_of_x2.sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::byte::binary::half_process_dot( + dim, + sum, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::byte::binary::half_process_l2s( + dim, + sum, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let dim = self.dim(); + let sum = rabitq::byte::binary::accumulate(self.packed_code, rhs.packed_code); + Distance::from_f32( + rabitq::byte::binary::half_process_cos( + dim, + sum, + rabitq::byte::CodeMetadata { + dis_u_2: self.sum_of_x2, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + }, + rabitq::byte::CodeMetadata { + dis_u_2: rhs.sum_of_x2, + norm_of_lattice: rhs.norm_of_lattice, + sum_of_code: rhs.sum_of_code, + }, + ) + .0, + ) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> Rabitq8Owned { + Rabitq8Owned { + dim: self.dim, + sum_of_x2: 1.0, + norm_of_lattice: self.norm_of_lattice, + sum_of_code: self.sum_of_code, + sum_of_abs_x: self.sum_of_abs_x / self.sum_of_x2.sqrt(), + packed_code: self.packed_code.to_owned(), + } + } + + fn operator_add(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_sub(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_mul(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} diff --git a/crates/vector/src/svect.rs b/crates/vector/src/svect.rs new file mode 100644 index 00000000..f04afe86 --- /dev/null +++ b/crates/vector/src/svect.rs @@ -0,0 +1,436 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; +use simd::Floating; + +#[derive(Debug, Clone)] +pub struct SVectOwned { + dim: u32, + indexes: Vec, + values: Vec, +} + +impl SVectOwned { + #[inline(always)] + pub fn new(dim: u32, indexes: Vec, values: Vec) -> Self { + Self::new_checked(dim, indexes, values).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dim: u32, indexes: Vec, values: Vec) -> Option { + if !(1..=1_048_575).contains(&dim) { + return None; + } + if indexes.len() != values.len() { + return None; + } + let len = indexes.len(); + for i in 1..len { + if !(indexes[i - 1] < indexes[i]) { + return None; + } + } + if len != 0 && !(indexes[len - 1] < dim) { + return None; + } + if S::reduce_or_of_is_zero_x(&values) { + return None; + } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dim, indexes, values)) + } + } + + /// # Safety + /// + /// * `dim` must be in `1..=1_048_575`. + /// * `indexes.len()` must be equal to `values.len()`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dim`. + /// * A floating number in `values` must not be positive zero or negative zero. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked(dim: u32, indexes: Vec, values: Vec) -> Self { + Self { + dim, + indexes, + values, + } + } + + #[inline(always)] + pub fn indexes(&self) -> &[u32] { + &self.indexes + } + + #[inline(always)] + pub fn values(&self) -> &[S] { + &self.values + } +} + +impl VectorOwned for SVectOwned { + type Borrowed<'a> = SVectBorrowed<'a, S>; + + #[inline(always)] + fn as_borrowed(&self) -> SVectBorrowed<'_, S> { + SVectBorrowed { + dim: self.dim, + indexes: &self.indexes, + values: &self.values, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct SVectBorrowed<'a, S> { + dim: u32, + indexes: &'a [u32], + values: &'a [S], +} + +impl<'a, S: Floating> SVectBorrowed<'a, S> { + #[inline(always)] + pub fn new(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + Self::new_checked(dim, indexes, values).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Option { + if !(1..=1_048_575).contains(&dim) { + return None; + } + if indexes.len() != values.len() { + return None; + } + let len = indexes.len(); + for i in 1..len { + if !(indexes[i - 1] < indexes[i]) { + return None; + } + } + if len != 0 && !(indexes[len - 1] < dim) { + return None; + } + for i in 0..len { + if values[i] == S::zero() { + return None; + } + } + #[allow(unsafe_code)] + unsafe { + Some(Self::new_unchecked(dim, indexes, values)) + } + } + + /// # Safety + /// + /// * `dim` must be in `1..=1_048_575`. + /// * `indexes.len()` must be equal to `values.len()`. + /// * `indexes` must be a strictly increasing sequence and the last in the sequence must be less than `dim`. + /// * A floating number in `values` must not be positive zero or negative zero. + #[inline(always)] + #[allow(unsafe_code)] + pub unsafe fn new_unchecked(dim: u32, indexes: &'a [u32], values: &'a [S]) -> Self { + Self { + dim, + indexes, + values, + } + } + + #[inline(always)] + pub fn indexes(&self) -> &'a [u32] { + self.indexes + } + + #[inline(always)] + pub fn values(&self) -> &'a [S] { + self.values + } + + #[inline(always)] + pub fn len(&self) -> u32 { + self.indexes.len() as u32 + } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.indexes.is_empty() + } +} + +impl VectorBorrowed for SVectBorrowed<'_, S> { + type Owned = SVectOwned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.dim + } + + #[inline(always)] + fn own(&self) -> SVectOwned { + SVectOwned { + dim: self.dim, + indexes: self.indexes.to_vec(), + values: self.values.to_vec(), + } + } + + #[inline(always)] + fn norm(&self) -> f32 { + S::reduce_sum_of_x2(self.values).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + Distance::from(-xy) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + let d2 = S::reduce_sum_of_d2_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + Distance::from(d2) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy_sparse(self.indexes, self.values, rhs.indexes, rhs.values); + let x2 = S::reduce_sum_of_x2(self.values); + let y2 = S::reduce_sum_of_x2(rhs.values); + Distance::from(1.0 - xy / (x2 * y2).sqrt()) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> SVectOwned { + let l = S::reduce_sum_of_x2(self.values).sqrt(); + let mut indexes = self.indexes.to_vec(); + let mut values = self.values.to_vec(); + let n = indexes.len(); + S::vector_mul_scalar_inplace(&mut values, 1.0 / l); + let mut j = 0_usize; + for i in 0..n { + if values[i] != S::zero() { + indexes[j] = indexes[i]; + values[j] = values[i]; + j += 1; + } + } + indexes.truncate(j); + values.truncate(j); + SVectOwned::new(self.dim, indexes, values) + } + + fn operator_add(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; (size1 + size2) as _]; + let mut values = vec![S::zero(); (size1 + size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index.min(rhs_index); + values[pos] = S::scalar_add( + lhs_value.mask(lhs_index <= rhs_index), + rhs_value.mask(lhs_index >= rhs_index), + ); + pos1 += (lhs_index <= rhs_index) as u32; + pos2 += (lhs_index >= rhs_index) as u32; + pos += (values[pos] != S::zero()) as usize; + } + for i in pos1..size1 { + indexes[pos] = self.indexes[i as usize]; + values[pos] = self.values[i as usize]; + pos += 1; + } + for i in pos2..size2 { + indexes[pos] = rhs.indexes[i as usize]; + values[pos] = rhs.values[i as usize]; + pos += 1; + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dim, indexes, values) + } + + fn operator_sub(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; (size1 + size2) as _]; + let mut values = vec![S::zero(); (size1 + size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index.min(rhs_index); + values[pos] = S::scalar_sub( + lhs_value.mask(lhs_index <= rhs_index), + rhs_value.mask(lhs_index >= rhs_index), + ); + pos1 += (lhs_index <= rhs_index) as u32; + pos2 += (lhs_index >= rhs_index) as u32; + pos += (values[pos] != S::zero()) as usize; + } + for i in pos1..size1 { + indexes[pos] = self.indexes[i as usize]; + values[pos] = self.values[i as usize]; + pos += 1; + } + for i in pos2..size2 { + indexes[pos] = rhs.indexes[i as usize]; + values[pos] = S::scalar_neg(rhs.values[i as usize]); + pos += 1; + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dim, indexes, values) + } + + fn operator_mul(&self, rhs: Self) -> Self::Owned { + assert_eq!(self.dim, rhs.dim); + let size1 = self.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; std::cmp::min(size1, size2) as _]; + let mut values = vec![S::zero(); std::cmp::min(size1, size2) as _]; + while pos1 < size1 && pos2 < size2 { + let lhs_index = self.indexes[pos1 as usize]; + let rhs_index = rhs.indexes[pos2 as usize]; + match lhs_index.cmp(&rhs_index) { + std::cmp::Ordering::Less => { + pos1 += 1; + } + std::cmp::Ordering::Equal => { + // only both indexes are not zero, values are multiplied + let lhs_value = self.values[pos1 as usize]; + let rhs_value = rhs.values[pos2 as usize]; + indexes[pos] = lhs_index; + values[pos] = S::scalar_mul(lhs_value, rhs_value); + pos1 += 1; + pos2 += 1; + // only increment pos if the value is not zero + pos += (values[pos] != S::zero()) as usize; + } + std::cmp::Ordering::Greater => { + pos2 += 1; + } + } + } + indexes.truncate(pos); + values.truncate(pos); + SVectOwned::new(self.dim, indexes, values) + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} + +impl PartialEq for SVectBorrowed<'_, S> { + fn eq(&self, other: &Self) -> bool { + if self.dim != other.dim { + return false; + } + if self.indexes.len() != other.indexes.len() { + return false; + } + for (&l, &r) in self.indexes.iter().zip(other.indexes.iter()) { + if l != r { + return false; + } + } + for (&l, &r) in self.values.iter().zip(other.values.iter()) { + if l != r { + return false; + } + } + true + } +} + +impl PartialOrd for SVectBorrowed<'_, S> { + fn partial_cmp(&self, other: &Self) -> Option { + use std::cmp::Ordering; + if self.dim != other.dim { + return None; + } + let mut lhs = self + .indexes + .iter() + .copied() + .zip(self.values.iter().copied()); + let mut rhs = other + .indexes + .iter() + .copied() + .zip(other.values.iter().copied()); + loop { + return match (lhs.next(), rhs.next()) { + (Some(lh), Some(rh)) => match lh.0.cmp(&rh.0) { + Ordering::Equal => match lh.1.partial_cmp(&rh.1)? { + Ordering::Equal => continue, + x => Some(x), + }, + Ordering::Less => Some(if lh.1 < S::zero() { + Ordering::Less + } else { + Ordering::Greater + }), + Ordering::Greater => Some(if S::zero() < rh.1 { + Ordering::Less + } else { + Ordering::Greater + }), + }, + (Some((_, x)), None) => Some(PartialOrd::partial_cmp(&x, &S::zero())?), + (None, Some((_, y))) => Some(PartialOrd::partial_cmp(&S::zero(), &y)?), + (None, None) => Some(Ordering::Equal), + }; + } + } +} diff --git a/crates/vector/src/vect.rs b/crates/vector/src/vect.rs new file mode 100644 index 00000000..84a658f3 --- /dev/null +++ b/crates/vector/src/vect.rs @@ -0,0 +1,218 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::{VectorBorrowed, VectorOwned}; +use distance::Distance; +use simd::Floating; +use std::cmp::Ordering; + +#[derive(Debug, Clone)] +#[repr(transparent)] +pub struct VectOwned(Vec); + +impl VectOwned { + #[inline(always)] + pub fn new(slice: Vec) -> Self { + Self::new_checked(slice).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(slice: Vec) -> Option { + if !(1..=65535).contains(&slice.len()) { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { Self::new_unchecked(slice) }) + } + + /// # Safety + /// + /// * `slice.len()` must not be zero. + /// * `slice.len()` must be less than 65536. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked(slice: Vec) -> Self { + Self(slice) + } + + #[inline(always)] + pub fn slice(&self) -> &[S] { + self.0.as_slice() + } + + #[inline(always)] + pub fn slice_mut(&mut self) -> &mut [S] { + self.0.as_mut_slice() + } + + #[inline(always)] + pub fn into_vec(self) -> Vec { + self.0 + } +} + +impl VectorOwned for VectOwned { + type Borrowed<'a> = VectBorrowed<'a, S>; + + #[inline(always)] + fn as_borrowed(&self) -> VectBorrowed<'_, S> { + VectBorrowed(self.0.as_slice()) + } +} + +#[derive(Debug, Clone, Copy)] +#[repr(transparent)] +pub struct VectBorrowed<'a, S>(&'a [S]); + +impl<'a, S: Floating> VectBorrowed<'a, S> { + #[inline(always)] + pub fn new(slice: &'a [S]) -> Self { + Self::new_checked(slice).expect("invalid data") + } + + #[inline(always)] + pub fn new_checked(slice: &'a [S]) -> Option { + if !(1..=65535).contains(&slice.len()) { + return None; + } + #[allow(unsafe_code)] + Some(unsafe { Self::new_unchecked(slice) }) + } + + /// # Safety + /// + /// * `slice.len()` must not be zero. + /// * `slice.len()` must be less than 65536. + #[allow(unsafe_code)] + #[inline(always)] + pub unsafe fn new_unchecked(slice: &'a [S]) -> Self { + Self(slice) + } + + #[inline(always)] + pub fn slice(&self) -> &'a [S] { + self.0 + } +} + +impl VectorBorrowed for VectBorrowed<'_, S> { + type Owned = VectOwned; + + #[inline(always)] + fn dim(&self) -> u32 { + self.0.len() as u32 + } + + #[inline(always)] + fn own(&self) -> VectOwned { + VectOwned(self.0.to_vec()) + } + + #[inline(always)] + fn norm(&self) -> f32 { + S::reduce_sum_of_x2(self.0).sqrt() + } + + #[inline(always)] + fn operator_dot(self, rhs: Self) -> Distance { + Distance::from(-S::reduce_sum_of_xy(self.slice(), rhs.slice())) + } + + #[inline(always)] + fn operator_l2s(self, rhs: Self) -> Distance { + Distance::from(S::reduce_sum_of_d2(self.slice(), rhs.slice())) + } + + #[inline(always)] + fn operator_cos(self, rhs: Self) -> Distance { + let xy = S::reduce_sum_of_xy(self.slice(), rhs.slice()); + let x2 = S::reduce_sum_of_x2(self.0); + let y2 = S::reduce_sum_of_x2(rhs.0); + Distance::from(1.0 - xy / (x2 * y2).sqrt()) + } + + #[inline(always)] + fn operator_hamming(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn operator_jaccard(self, _: Self) -> Distance { + unimplemented!() + } + + #[inline(always)] + fn function_normalize(&self) -> VectOwned { + let mut data = self.0.to_vec(); + let l = S::reduce_sum_of_x2(&data).sqrt(); + S::vector_mul_scalar_inplace(&mut data, 1.0 / l); + VectOwned(data) + } + + fn operator_add(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_add(self.slice(), rhs.slice())) + } + + fn operator_sub(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_sub(self.slice(), rhs.slice())) + } + + fn operator_mul(&self, rhs: Self) -> Self::Owned { + VectOwned::new(S::vector_mul(self.slice(), rhs.slice())) + } + + fn operator_and(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_or(&self, _: Self) -> Self::Owned { + unimplemented!() + } + + fn operator_xor(&self, _: Self) -> Self::Owned { + unimplemented!() + } +} + +impl PartialEq for VectBorrowed<'_, S> { + fn eq(&self, other: &Self) -> bool { + if self.0.len() != other.0.len() { + return false; + } + let n = self.0.len(); + for i in 0..n { + if self.0[i] != other.0[i] { + return false; + } + } + true + } +} + +impl PartialOrd for VectBorrowed<'_, S> { + fn partial_cmp(&self, other: &Self) -> Option { + if self.0.len() != other.0.len() { + return None; + } + let n = self.0.len(); + for i in 0..n { + match PartialOrd::partial_cmp(&self.0[i], &other.0[i])? { + Ordering::Less => return Some(Ordering::Less), + Ordering::Equal => continue, + Ordering::Greater => return Some(Ordering::Greater), + } + } + Some(Ordering::Equal) + } +} diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml new file mode 100644 index 00000000..b9d11560 --- /dev/null +++ b/crates/xtask/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "xtask" +version.workspace = true +edition.workspace = true +publish = false + +[dependencies] +clap = { version = "4.6.1", features = ["derive", "env"] } +object = { version = "0.39.0", features = ["read", "wasm"] } +serde.workspace = true +serde_json = "1.0.149" +shlex = "1.3.0" +target-triple = "1.0.0" +tempfile = "3.27.0" + +[lints] +workspace = true diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs new file mode 100644 index 00000000..66afc023 --- /dev/null +++ b/crates/xtask/src/main.rs @@ -0,0 +1,580 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use clap::{Args, Parser, Subcommand}; +use object::{Object, ObjectSymbol}; +use serde::Deserialize; +use std::collections::{HashMap, HashSet}; +use std::error::Error; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; + +#[derive(Parser)] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + Build(BuildArgs), + Clippy(ClippyArgs), +} + +#[derive(Args)] +struct BuildArgs { + #[arg(short, long, default_value = "./build")] + output: String, + #[arg(long, default_value = target_triple::TARGET, env = "TARGET")] + target: String, + #[arg(long, default_value = "release", env = "PROFILE")] + profile: String, + #[arg(long, env = "RUNNER")] + runner: Option, +} + +#[derive(Args)] +struct ClippyArgs { + #[arg(long, default_value = target_triple::TARGET, env = "TARGET")] + target: String, + #[arg(long, default_value = "release", env = "PROFILE")] + profile: String, +} + +struct RustcCfg { + is_macos: bool, + is_windows: bool, + is_emscripten: bool, + is_unix: bool, + is_powerpc64: bool, +} + +impl RustcCfg { + fn dll_prefix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok("lib") + } else if self.is_windows || self.is_emscripten { + Ok("") + } else if self.is_unix { + Ok("lib") + } else { + Err("unknown operating system".into()) + } + } + fn dll_suffix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok(".dylib") + } else if self.is_windows { + Ok(".dll") + } else if self.is_emscripten { + Ok(".wasm") + } else if self.is_unix { + Ok(".so") + } else { + Err("unknown operating system".into()) + } + } + fn exe_suffix(&self) -> Result<&'static str, Box> { + if self.is_macos { + Ok("") + } else if self.is_windows { + Ok(".exe") + } else if self.is_emscripten { + Ok(".js") + } else if self.is_unix { + Ok("") + } else { + Err("unknown operating system".into()) + } + } + fn ext_suffix(&self, fork: &str) -> Result<&'static str, Box> { + if self.is_macos { + Ok(if matches!(fork, "pg14" | "pg15") { + ".so" + } else { + ".dylib" + }) + } else if self.is_windows { + Ok(".dll") + } else if self.is_emscripten || self.is_unix { + Ok(".so") + } else { + Err("unknown operating system".into()) + } + } +} + +#[derive(Deserialize)] +struct CargoMetadata { + target_directory: String, +} + +fn pg_config(pg_config: impl AsRef) -> Result, Box> { + let mut command = Command::new(pg_config.as_ref()); + command.stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("PostgreSQL failed: {command_status}").into()); + } + let contents = String::from_utf8(command_output.stdout)?; + let mut result = HashMap::new(); + for line in contents.lines() { + if let Some((key, value)) = line.split_once(" = ") { + result.insert(key.to_string(), value.to_string()); + } + } + Ok(result) +} + +fn control_file(path: impl AsRef) -> Result, Box> { + let path = path.as_ref(); + eprintln!("Reading {path:?}"); + let contents = std::fs::read_to_string(path)?; + let mut result = HashMap::new(); + for line in contents.lines() { + if let Some((key, prefix_stripped)) = line.split_once(" = '") + && let Some(value) = prefix_stripped.strip_suffix("'") + { + result.insert(key.to_string(), value.to_string()); + } + } + Ok(result) +} + +fn rustc_cfg(target: &str) -> Result> { + let mut command = Command::new("rustc"); + command + .args(["--print", "cfg"]) + .args(["--target", target]) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("Rust failed: {command_status}").into()); + } + let contents = String::from_utf8(command_output.stdout)?; + let mut cfgs = HashSet::new(); + for line in contents.lines() { + cfgs.insert(line.to_string()); + } + Ok(RustcCfg { + is_macos: cfgs.contains("target_os=\"macos\""), + is_unix: cfgs.contains("target_family=\"unix\""), + is_emscripten: cfgs.contains("target_os=\"emscripten\""), + is_windows: cfgs.contains("target_os=\"windows\""), + is_powerpc64: cfgs.contains("target_arch=\"powerpc64\""), + }) +} + +fn cargo_metadata() -> Result> { + let mut command = Command::new("cargo"); + command + .args(["metadata", "--format-version", "1"]) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + let contents = String::from_utf8(command_output.stdout)?; + let cargo_metadata: CargoMetadata = serde_json::from_str(&contents)?; + Ok(cargo_metadata) +} + +fn build( + pg_config: impl AsRef, + pg_version: &str, + rustc_cfg: &RustcCfg, + cargo_metadata: &CargoMetadata, + profile: &str, + target: &str, +) -> Result> { + let mut command = Command::new("cargo"); + command.args(["rustc"]); + if !matches!(profile, "dev" | "test") { + command.args(["--crate-type", "cdylib"]); + } + command + .args(["-p", "vchord", "--lib"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + let mut result = PathBuf::from(&cargo_metadata.target_directory); + result.push(target); + result.push(match profile { + "dev" | "test" => "debug", + "release" | "bench" => "release", + profile => profile, + }); + result.push(format!( + "{}vchord{}", + rustc_cfg.dll_prefix()?, + rustc_cfg.dll_suffix()? + )); + Ok(result) +} + +fn parse(rustc_cfg: &RustcCfg, obj: impl AsRef) -> Result, Box> { + let obj = obj.as_ref(); + eprintln!("Reading {obj:?}"); + let contents = std::fs::read(obj)?; + let object = object::File::parse(contents.as_slice())?; + let exports; + if rustc_cfg.is_macos { + exports = object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .flat_map(|x| x.strip_prefix("_")) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } else if rustc_cfg.is_emscripten { + exports = object + .symbols() + .flat_map(|x| x.name().ok()) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } else { + exports = object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .filter(|x| x.starts_with("__pgrx_internals")) + .map(str::to_string) + .collect(); + } + Ok(exports) +} + +fn generate( + runner: &Option>, + pg_config: impl AsRef, + pg_version: &str, + rustc_cfg: &RustcCfg, + cargo_metadata: &CargoMetadata, + profile: &str, + target: &str, + exports: Vec, + postmaster: impl AsRef, +) -> Result> { + let imports = if rustc_cfg.is_powerpc64 { + let postmaster = postmaster.as_ref(); + eprintln!("Reading {postmaster:?}"); + let contents = std::fs::read(postmaster)?; + let object = object::File::parse(contents.as_slice())?; + object + .exports()? + .into_iter() + .flat_map(|x| std::str::from_utf8(x.name())) + .filter(|x| should_stub_postmaster_export(x)) + .map(str::to_string) + .collect::>() + } else { + Vec::new() + }; + let mut pgrx_embed = tempfile::NamedTempFile::new()?; + eprintln!("Writing {:?}", pgrx_embed.path()); + let contents = format!( + "crate::schema_generation!({}; {});", + exports.join(" "), + imports.join(" ") + ); + std::io::Write::write_all(pgrx_embed.as_file_mut(), contents.as_bytes())?; + let mut command = Command::new("cargo"); + command + .args(["rustc"]) + .args(["-p", "vchord", "--bin", "pgrx_embed_vchord"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .args(["--", "-C", "lto=off", "--cfg", "pgrx_embed"]) + .env("PGRX_EMBED", pgrx_embed.path()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + let mut result = PathBuf::from(&cargo_metadata.target_directory); + result.push(target); + result.push(match profile { + "dev" | "test" => "debug", + "release" | "bench" => "release", + profile => profile, + }); + result.push(format!("pgrx_embed_vchord{}", rustc_cfg.exe_suffix()?)); + let mut command; + if let Some(runner) = runner { + command = Command::new(&runner[0]); + for arg in runner[1..].iter() { + command.arg(arg); + } + command.arg(result); + } else { + command = Command::new(result); + } + command.stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_output = command.output()?; + let command_status = command_output.status; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + let command_stdout = String::from_utf8(command_output.stdout)?.replace("\t", " "); + Ok(command_stdout) +} + +fn should_stub_postmaster_export(symbol: &str) -> bool { + // These symbols are supplied by the C runtime or linker when the schema + // helper executable is linked. Defining a PostgreSQL stub for any of them + // creates duplicate symbols on some targets (notably powerpc64le). + !matches!( + symbol, + "_start" + | "_IO_stdin_used" + | "main" + | "__data_start" + | "data_start" + | "__bss_start" + | "_edata" + | "_end" + ) +} + +#[cfg(test)] +mod tests { + use super::should_stub_postmaster_export; + + #[test] + fn schema_helper_does_not_stub_crt_or_linker_symbols() { + for symbol in [ + "_start", + "_IO_stdin_used", + "main", + "__data_start", + "data_start", + "__bss_start", + "_edata", + "_end", + ] { + assert!(!should_stub_postmaster_export(symbol), "{symbol}"); + } + assert!(should_stub_postmaster_export("PostgresMain")); + } +} + +fn install_by_copying( + src: impl AsRef, + dst: impl AsRef, + #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, +) -> Result<(), Box> { + eprintln!("Copying {:?} to {:?}", src.as_ref(), dst.as_ref()); + std::fs::copy(src, &dst)?; + #[cfg(target_family = "unix")] + { + use std::fs::Permissions; + use std::os::unix::fs::PermissionsExt; + let perm = Permissions::from_mode(if !is_executable { 0o644 } else { 0o755 }); + std::fs::set_permissions(dst, perm)?; + } + Ok(()) +} + +fn install_by_writing( + contents: impl AsRef<[u8]>, + dst: impl AsRef, + #[cfg_attr(not(target_family = "unix"), expect(unused_variables))] is_executable: bool, +) -> Result<(), Box> { + eprintln!("Writing {:?}", dst.as_ref()); + std::fs::write(&dst, contents)?; + #[cfg(target_family = "unix")] + { + use std::fs::Permissions; + use std::os::unix::fs::PermissionsExt; + let perm = Permissions::from_mode(if !is_executable { 0o644 } else { 0o755 }); + std::fs::set_permissions(dst, perm)?; + } + Ok(()) +} + +fn clippy( + pg_config: impl AsRef, + pg_version: &str, + profile: &str, + target: &str, +) -> Result<(), Box> { + let mut command = Command::new("cargo"); + command + .args(["clippy"]) + .args(["--workspace"]) + .args(["--profile", profile]) + .args(["--target", target]) + .args(["--features", pg_version]) + .env("PGRX_PG_CONFIG_PATH", pg_config.as_ref()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + eprintln!("Running {command:?}"); + let command_status = command.spawn()?.wait()?; + if !command_status.success() { + return Err(format!("Cargo failed: {command_status}").into()); + } + Ok(()) +} + +fn main() -> Result<(), Box> { + #[allow(unsafe_code)] + unsafe { + std::env::remove_var("PGRX_PG_CONFIG_PATH"); + } + let cli = Cli::parse(); + match cli.command { + Commands::Build(BuildArgs { + output, + target, + profile, + runner, + }) => { + if !std::fs::exists("./vchord.control")? { + return Err("The script must be run from the VectorChord source directory.".into()); + } + let vchord_version = control_file("./vchord.control")?["default_version"].clone(); + let runner = runner.and_then(|runner| shlex::split(&runner)); + let path = if let Some(value) = std::env::var_os("PG_CONFIG") { + PathBuf::from(value) + } else { + return Err("Environment variable `PG_CONFIG` is not set.".into()); + }; + let pg_config = pg_config(&path)?; + let pg_version = { + let version = pg_config["VERSION"].clone(); + if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = + prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) + { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) + } + } else { + return Err("PostgreSQL version is invalid.".into()); + } + }; + let postmaster = format!("{}/postgres", pg_config["BINDIR"]); + let rustc_cfg = rustc_cfg(&target)?; + let cargo_metadata = cargo_metadata()?; + let obj = build( + &path, + &pg_version, + &rustc_cfg, + &cargo_metadata, + &profile, + &target, + )?; + let pkglibdir = format!("{output}/pkglibdir"); + let sharedir = format!("{output}/sharedir"); + let sharedir_extension = format!("{sharedir}/extension"); + if std::fs::exists(&output)? { + std::fs::remove_dir_all(&output)?; + } + std::fs::create_dir_all(&output)?; + std::fs::create_dir_all(&pkglibdir)?; + std::fs::create_dir_all(&sharedir)?; + std::fs::create_dir_all(&sharedir_extension)?; + install_by_copying( + &obj, + format!("{pkglibdir}/vchord{}", rustc_cfg.ext_suffix(&pg_version)?), + true, + )?; + install_by_copying( + "./vchord.control", + format!("{sharedir}/extension/vchord.control"), + false, + )?; + if vchord_version != "0.0.0" { + for e in std::fs::read_dir("./sql/upgrade")?.collect::, _>>()? { + install_by_copying( + e.path(), + format!("{sharedir}/extension/{}", e.file_name().display()), + false, + )?; + } + install_by_copying( + format!("./sql/install/vchord--{vchord_version}.sql"), + format!("{sharedir}/extension/vchord--{vchord_version}.sql"), + false, + )?; + } else { + let exports = parse(&rustc_cfg, obj)?; + install_by_writing( + generate( + &runner, + &path, + &pg_version, + &rustc_cfg, + &cargo_metadata, + &profile, + &target, + exports, + postmaster, + )?, + format!("{sharedir_extension}/vchord--0.0.0.sql"), + false, + )?; + } + } + Commands::Clippy(ClippyArgs { target, profile }) => { + if !std::fs::exists("./vchord.control")? { + return Err("The script must be run from the VectorChord source directory.".into()); + } + let path = if let Some(value) = std::env::var_os("PG_CONFIG") { + PathBuf::from(value) + } else { + return Err("Environment variable `PG_CONFIG` is not set.".into()); + }; + let pg_config = pg_config(&path)?; + let pg_version = { + let version = pg_config["VERSION"].clone(); + if let Some(prefix_stripped) = version.strip_prefix("PostgreSQL ") { + if let Some((stripped, _)) = + prefix_stripped.split_once(|c: char| !c.is_ascii_digit()) + { + format!("pg{stripped}",) + } else { + format!("pg{prefix_stripped}",) + } + } else { + return Err("PostgreSQL version is invalid.".into()); + } + }; + clippy(&path, &pg_version, &profile, &target)?; + } + } + Ok(()) +} diff --git a/deny.toml b/deny.toml new file mode 100644 index 00000000..ae84eb04 --- /dev/null +++ b/deny.toml @@ -0,0 +1,25 @@ +[graph] +all-features = true + +[advisories] +ignore = [ + "RUSTSEC-2021-0127", # serde_cbor is unmaintained + "RUSTSEC-2024-0436", # paste - no longer maintained + # Build-time-only transitive dependency of validator_derive 0.20.0. The + # advisory has no maintained compatible release yet; remove this exception + # as soon as validator publishes a replacement. + "RUSTSEC-2026-0173", +] + +[licenses] +allow = [ + "Apache-2.0", + "Apache-2.0 WITH LLVM-exception", + "MIT", + # "MPL-2.0", + "BSD-3-Clause", + "ISC", + "Unicode-3.0", +] +exceptions = [{ allow = ["Zlib"], crate = "foldhash" }] +private = { ignore = true } diff --git a/deploy/systemd/tilemaxsimd.env.example b/deploy/systemd/tilemaxsimd.env.example new file mode 100644 index 00000000..6f3b6267 --- /dev/null +++ b/deploy/systemd/tilemaxsimd.env.example @@ -0,0 +1,3 @@ +# This unit is opt-in. Do not enable it on nodes where TileMaxSim is disabled. +# Memory values are GiB; startup fails if the requested CUDA arena is unavailable. +TILEMAXSIMD_ARGS="--socket /run/vectorchord/tilemaxsim.sock --status-socket /run/vectorchord/tilemaxsim-status.sock --socket-mode 660 --status-socket-mode 660 --gpu-memory-gb 0=20 --gpu-workspace-gb 2 --host-cache-gb 8 --max-inflight-request-gb 1 --contract-root MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors --scheduler-policy fair-priority --scheduler-quantum-fmas 4000000000 --max-connections 256 --max-queued-requests 128 --max-tenant-queued-requests 16" diff --git a/deploy/systemd/tilemaxsimd.service b/deploy/systemd/tilemaxsimd.service new file mode 100644 index 00000000..e13b94ef --- /dev/null +++ b/deploy/systemd/tilemaxsimd.service @@ -0,0 +1,34 @@ +[Unit] +Description=VectorChord native CUDA TileMaxSim daemon +Documentation=https://github.com/HuXinjing/VectorChord +After=local-fs.target +StartLimitIntervalSec=60 +StartLimitBurst=3 + +[Service] +Type=exec +User=vectorchord +Group=vectorchord +RuntimeDirectory=vectorchord +RuntimeDirectoryMode=0750 +StateDirectory=vectorchord +StateDirectoryMode=0750 +EnvironmentFile=/etc/vectorchord/tilemaxsimd.env +ExecStart=/usr/local/bin/tilemaxsimd $TILEMAXSIMD_ARGS +ExecStartPost=/usr/local/bin/tilemaxsimctl --socket /run/vectorchord/tilemaxsim-status.sock --wait-timeout-ms 300000 +ExecReload=/bin/kill -HUP $MAINPID +Restart=on-failure +RestartSec=2 +TimeoutStopSec=120 +KillSignal=SIGTERM +UMask=0027 +LimitMEMLOCK=infinity +LimitNOFILE=1048576 +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=strict +ProtectHome=true +ReadWritePaths=/run/vectorchord /var/lib/vectorchord + +[Install] +WantedBy=multi-user.target diff --git a/devtools/dump-codegen.sh b/devtools/dump-codegen.sh new file mode 100755 index 00000000..7ea0f795 --- /dev/null +++ b/devtools/dump-codegen.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -e + +cat << EOF +#include "postgres.h" + +#include "access/amapi.h" +#include "fmgr.h" + +#include + +#define DECLARE(funcname) \ + extern Datum funcname() { exit(1); } \ + extern const Pg_finfo_record *pg_finfo_##funcname(void); \ + const Pg_finfo_record *pg_finfo_##funcname(void) { return &my_finfo; } \ + extern int no_such_variable + +#define DECLARE_AMHANDLER(funcname) \ + extern Datum funcname() { return amhandler(); } \ + extern const Pg_finfo_record *pg_finfo_##funcname(void); \ + const Pg_finfo_record *pg_finfo_##funcname(void) { return &my_finfo; } \ + extern int no_such_variable + +static const Pg_finfo_record my_finfo = {1}; + +PG_MODULE_MAGIC; + +Datum amhandler() { + IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); + amroutine->amsupport = 1; + amroutine->amcanorderbyop = true; + return (Datum) amroutine; +} +EOF + +printf "\n" + +while read -r line; do + if [[ $line == *"amhandler"* ]]; then + echo "DECLARE_AMHANDLER($line);" + else + echo "DECLARE($line);" + fi +done <<< $(grep -ohr "'\w\+_wrapper'" $(dirname "$0")/../sql | sort | uniq | sed "s/'//g") \ No newline at end of file diff --git a/devtools/dump.sh b/devtools/dump.sh new file mode 100755 index 00000000..63e21672 --- /dev/null +++ b/devtools/dump.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -e + +mkdir -p $(dirname "$0")/../target/devtools + +if [ $# -eq 0 ] ; then + echo "Usage: $0 INITIAL_VERSION VERSIONS.." + echo "Dump the extension members. Install INITIAL_VERSION first and upgrade to every version in VERSIONS." + echo "The extension members are installed in database \"vchord\", so use \"createdb vchord\" to setup." + echo "Examples:" + echo " ./devtools/dump.sh 0.1.11 0.2.0 > ./dump_upgrade.sql" + echo " ./devtools/dump.sh 0.2.0 > ./dump_install.sql" + echo " diff ./dump_upgrade.sql ./dump_install.sql" + exit 0 +fi + +f=() +prev_arg="" +for arg in "$@"; do + if [ "$prev_arg" = "" ]; then + x=$(realpath "$(dirname "$0")/../sql/install/vchord--${arg}.sql") + else + x=$(realpath "$(dirname "$0")/../sql/upgrade/vchord--${prev_arg}--${arg}.sql") + fi + prev_arg=$arg + f+=("$x") +done + +so=$(realpath $(dirname "$0")/../target/devtools/vchord.so) +$(dirname "$0")/dump-codegen.sh | gcc -I $(pg_config --includedir-server) -fPIC -shared -o $so -x c - + +sql=$(mktemp) +echo "BEGIN;" >> $sql +echo "CREATE SCHEMA vchord;" >> $sql +echo "SET LOCAL search_path TO vchord,public;" >> $sql +cat ${f[@]} \ + | grep -v '^\\' \ + | sed "s|@extschema@|vchord|g" \ + | sed "s|MODULE_PATHNAME|$so|g" \ + >> $sql +echo "END;" >> $sql + +psql -d vchord -f $sql 1>&2 +pg_dump -s -d vchord +psql -d vchord -c "DROP SCHEMA IF EXISTS vchord CASCADE;" 1>&2 diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..19f0f01a --- /dev/null +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -0,0 +1,441 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +try: + from . import tilemaxsim_reference_sidecar as sidecar +except ImportError: + import tilemaxsim_reference_sidecar as sidecar + + +def request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + candidates: list[tuple[int, list[list[float]]]], +) -> bytes: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + body = bytearray( + sidecar.REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + ) + ) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + for candidate_id, tensor in candidates: + body.extend(sidecar.CANDIDATE_FIXED.pack(candidate_id, len(tensor))) + body.extend(struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, []))) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body + ) + + +def external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], +) -> tuple[bytes, dict[str, bytes]]: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + contract = model_contract_id.encode() + body = bytearray( + sidecar.EXTERNAL_REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + len(contract), + ) + ) + body.extend(contract) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + objects = {} + for candidate_id, tensor_ref, tensor in candidates: + payload = struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, [])) + objects[tensor_ref] = payload + reference = tensor_ref.encode() + checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}".encode() + body.extend( + sidecar.EXTERNAL_CANDIDATE_FIXED.pack( + candidate_id, len(tensor), len(reference), len(checksum) + ) + ) + body.extend(reference) + body.extend(checksum) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body, + objects, + ) + + +def scheduled_external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], + tenant: str, + priority: int, + timeout_ms: int, +) -> tuple[bytes, dict[str, bytes]]: + legacy, objects = external_request_frame( + request_id, dtype, query, model_contract_id, candidates + ) + body = legacy[sidecar.HEADER.size :] + fixed = sidecar.EXTERNAL_REQUEST_FIXED.unpack_from(body) + tenant_bytes = tenant.encode() + scheduled = bytearray( + sidecar.SCHEDULED_EXTERNAL_REQUEST_FIXED.pack( + *fixed, priority, timeout_ms, len(tenant_bytes) + ) + ) + contract_length = fixed[-1] + scheduled.extend(body[sidecar.EXTERNAL_REQUEST_FIXED.size :][:contract_length]) + scheduled.extend(tenant_bytes) + scheduled.extend(body[sidecar.EXTERNAL_REQUEST_FIXED.size + contract_length :]) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.SCHEDULED_EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(scheduled), + ) + + scheduled, + objects, + ) +def decode_response(frame: bytes) -> tuple[int, int, list[tuple[int, float]] | str]: + magic, version, kind, request_id, body_len = sidecar.HEADER.unpack_from(frame) + assert magic == sidecar.MAGIC + assert version in sidecar.SUPPORTED_VERSIONS + assert kind == sidecar.RESPONSE_KIND + assert len(frame) == sidecar.HEADER.size + body_len + status, count_or_length = sidecar.RESPONSE_FIXED.unpack_from( + frame, sidecar.HEADER.size + ) + offset = sidecar.HEADER.size + sidecar.RESPONSE_FIXED.size + if status: + return request_id, status, frame[offset : offset + count_or_length].decode() + results = [] + for _ in range(count_or_length): + results.append(sidecar.RESULT.unpack_from(frame, offset)) + offset += sidecar.RESULT.size + assert offset == len(frame) + return request_id, status, results + + +class ReferenceSidecarTest(unittest.TestCase): + def test_f32_exact_scores_and_opaque_ids(self) -> None: + frame = request_frame( + 41, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5]]), + ], + ) + request_id, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(request_id, 41) + self.assertEqual(status, 0) + self.assertEqual(results, [(17, 2.0), (3, 1.0)]) + + def test_f16_exact_scores(self) -> None: + frame = request_frame( + 42, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + [(0, [[0.75, 0.0], [0.0, 0.5]])], + ) + _, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(status, 0) + self.assertEqual(results, [(0, 1.25)]) + + def test_shared_raw_decoder_preserves_inline_payloads(self) -> None: + frame = request_frame( + 420, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + [(11, [[0.5, 0.25]])], + ) + request = sidecar.parse_request_frame(frame) + self.assertIsInstance(request, sidecar.InlineTensorRequest) + assert isinstance(request, sidecar.InlineTensorRequest) + self.assertEqual(request.request_id, 420) + self.assertEqual(request.query_rows, 1) + self.assertEqual(request.dimension, 2) + self.assertEqual([item.candidate_id for item in request.candidates], [11]) + self.assertEqual(len(request.query_payload), 4) + self.assertEqual(len(request.candidates[0].payload), 4) + + def test_duplicate_id_and_non_finite_input_fail_closed(self) -> None: + duplicate = request_frame( + 43, + sidecar.DTYPE_F32, + [[1.0]], + [(7, [[1.0]]), (7, [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(duplicate)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("duplicate candidate ID", message) + + non_finite = request_frame( + 44, + sidecar.DTYPE_F32, + [[float("nan")]], + [(0, [[1.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(non_finite)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("non-finite", message) + + def test_truncated_and_oversized_frames_fail_closed(self) -> None: + valid = request_frame(45, sidecar.DTYPE_F32, [[1.0]], [(0, [[1.0]])]) + _, status, message = decode_response(sidecar.process_frame(valid[:-1])) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("length mismatch", message) + + _, status, message = decode_response( + sidecar.process_frame(valid, sidecar.Limits(max_request_bytes=32)) + ) + self.assertEqual(status, sidecar.STATUS_RESOURCE_LIMIT) + self.assertIn("byte limit", message) + + def test_header_reserved_trailing_and_token_limit_fail_closed(self) -> None: + valid = request_frame(47, sidecar.DTYPE_F32, [[1.0]], [(9, [[1.0]])]) + invalid_frames = [] + for offset, value in ((0, 0), (4, 4), (6, 2)): + invalid = bytearray(valid) + invalid[offset] = value + invalid_frames.append(bytes(invalid)) + + reserved = bytearray(valid) + reserved[sidecar.HEADER.size + 14] = 1 + invalid_frames.append(bytes(reserved)) + + trailing = bytearray(valid) + trailing.extend(b"x") + struct.pack_into(" None: + frame, objects = external_request_frame( + 48, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "colqwen@immutable-revision", + [ + (77, "object://fixture/page-1", [[1.0, 0.0], [0.0, 1.0]]), + (4, "object://fixture/page-2", [[0.5, 0.5]]), + ], + ) + seen = [] + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + seen.append(request) + return objects[request.tensor_ref] + + request_id, status, results = decode_response( + sidecar.process_frame(frame, resolver=resolver) + ) + + self.assertEqual(request_id, 48) + self.assertEqual(status, 0) + self.assertEqual(results, [(77, 2.0), (4, 1.0)]) + self.assertEqual( + {request.model_contract_id for request in seen}, + {"colqwen@immutable-revision"}, + ) + self.assertEqual({request.tensor_ref for request in seen}, set(objects)) + + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.model_contract_id, "colqwen@immutable-revision") + self.assertEqual( + [candidate.candidate_id for candidate in parsed.candidates], [77, 4] + ) + + def test_external_v3_preserves_scheduler_metadata_and_v2_scoring(self) -> None: + frame, objects = scheduled_external_request_frame( + 49, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + "model@1", + [(7, "sha256://opaque", [[1.0, 0.0]])], + "tenant-a", + 17, + 4_000, + ) + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.scheduler_tenant, "tenant-a") + self.assertEqual(parsed.scheduler_priority, 17) + self.assertEqual(parsed.timeout_ms, 4_000) + request_id, status, results = decode_response( + sidecar.process_frame( + frame, resolver=lambda request: objects[request.tensor_ref] + ) + ) + self.assertEqual((request_id, status), (49, 0)) + assert isinstance(results, list) + self.assertAlmostEqual(results[0][1], 1.0) + + def test_external_v2_fails_closed_without_resolver_or_on_checksum_mismatch( + self, + ) -> None: + frame, objects = external_request_frame( + 49, + sidecar.DTYPE_F32, + [[1.0]], + "contract@1", + [(0, "object://immutable/page", [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(frame)) + self.assertEqual(status, sidecar.STATUS_COMPUTE_ERROR) + self.assertIn("resolver is not configured", message) + + def corrupt_resolver(request: sidecar.ExternalTensorRequest) -> bytes: + return objects[request.tensor_ref] + b"x" + + _, status, message = decode_response( + sidecar.process_frame(frame, resolver=corrupt_resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("byte length", message) + + def test_external_v2_validates_complete_control_frame_before_resolution( + self, + ) -> None: + frame, objects = external_request_frame( + 50, + sidecar.DTYPE_F32, + [[1.0, 0.0]], + "contract@1", + [(0, "object://immutable/page", [[1.0, 0.0]])], + ) + invalid = bytearray(frame) + contract_length = len("contract@1") + candidate_offset = ( + sidecar.HEADER.size + + sidecar.EXTERNAL_REQUEST_FIXED.size + + contract_length + + 8 + ) + reference_offset = candidate_offset + sidecar.EXTERNAL_CANDIDATE_FIXED.size + invalid[reference_offset] = 0 + called = False + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + nonlocal called + called = True + return objects[request.tensor_ref] + + _, status, message = decode_response( + sidecar.process_frame(bytes(invalid), resolver=resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("control characters", message) + self.assertFalse(called) + + def test_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilemaxsim.sock" + thread = threading.Thread( + target=sidecar.serve, + args=(path, sidecar.Limits()), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if path.exists() and stat.S_IMODE(path.stat().st_mode) == 0o600: + break + time.sleep(0.01) + else: + self.fail("sidecar socket was not created with mode 0600") + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) + + frame = request_frame( + 46, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [(5, [[1.0, 0.0], [0.0, 1.0]])], + ) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(path)) + connection.sendall(frame) + header = connection.recv(sidecar.HEADER.size) + while len(header) < sidecar.HEADER.size: + header += connection.recv(sidecar.HEADER.size - len(header)) + body_len = sidecar.HEADER.unpack(header)[4] + body = b"" + while len(body) < body_len: + body += connection.recv(body_len - len(body)) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + + request_id, status, results = decode_response(header + body) + self.assertEqual(request_id, 46) + self.assertEqual(status, 0) + self.assertEqual(results, [(5, 2.0)]) + self.assertFalse(path.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..267a77fc --- /dev/null +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -0,0 +1,803 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. + +This executable is intentionally simple and single-threaded. It is a protocol +oracle and end-to-end development aid, not a production or GPU implementation. +The CLI serves inline v1 requests. External-descriptor v2 requests require an +explicit resolver injected by a test or embedding application and fail closed +when no resolver is configured. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import math +import os +import signal +import socket +import stat +import struct +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +MAGIC = b"VCTM" +VERSION = 1 +EXTERNAL_VERSION = 2 +SCHEDULED_EXTERNAL_VERSION = 3 +SUPPORTED_VERSIONS = (VERSION, EXTERNAL_VERSION, SCHEDULED_EXTERNAL_VERSION) +REQUEST_KIND = 1 +RESPONSE_KIND = 2 +SCORING_SUM_QUERY_MAX_DOCUMENT_DOT = 1 +DTYPE_F32 = 1 +DTYPE_F16 = 2 + +HEADER = struct.Struct("<4sHHQQ") +REQUEST_FIXED = struct.Struct(" None: + super().__init__(message) + self.status = status + + +@dataclass(frozen=True) +class Limits: + max_request_bytes: int = 64 * 1024 * 1024 + max_batch_tokens: int = 1_000_000 + max_tensor_bytes: int = 1024 * 1024 * 1024 + max_candidates: int = 65_536 + + +@dataclass(frozen=True) +class ExternalTensorRequest: + model_contract_id: str + tensor_ref: str + rows: int + dimension: int + dtype: int + checksum: str + + +@dataclass(frozen=True) +class InlineTensorCandidate: + candidate_id: int + rows: int + payload: bytes + + +@dataclass(frozen=True) +class InlineTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + query_payload: bytes + candidates: tuple[InlineTensorCandidate, ...] + + +@dataclass(frozen=True) +class ExternalTensorCandidate: + candidate_id: int + descriptor: ExternalTensorRequest + + +@dataclass(frozen=True) +class ParsedExternalTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + model_contract_id: str + query_payload: bytes + candidates: tuple[ExternalTensorCandidate, ...] + scheduler_tenant: str = "__default__" + scheduler_priority: int = 0 + timeout_ms: int = 0 + + +ParsedRequest = InlineTensorRequest | ParsedExternalTensorRequest + + +# A reference resolver returns the exact row-major scalar bytes whose SHA-256 +# digest is stored in the descriptor. Production resolvers may adapt richer +# immutable object formats before returning this canonical tensor payload. +ExternalTensorResolver = Callable[[ExternalTensorRequest], bytes] + + +class Reader: + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.offset = 0 + + def take(self, count: int) -> bytes: + end = self.offset + count + if count < 0 or end > len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "truncated request") + chunk = self.payload[self.offset : end] + self.offset = end + return chunk + + def unpack(self, layout: struct.Struct) -> tuple: + return layout.unpack(self.take(layout.size)) + + def finish(self) -> None: + if self.offset != len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "trailing request bytes") + + +def checked_elements(rows: int, dimension: int) -> int: + if rows <= 0 or dimension <= 0: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor rows and dimension must be positive" + ) + elements = rows * dimension + if elements > (1 << 63) - 1: + raise SidecarError(STATUS_RESOURCE_LIMIT, "tensor shape is too large") + return elements + + +def dtype_size(dtype: int) -> int: + if dtype == DTYPE_F32: + return 4 + if dtype == DTYPE_F16: + return 2 + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + + +def checked_tensor_bytes(rows: int, dimension: int, dtype: int) -> int: + return checked_elements(rows, dimension) * dtype_size(dtype) + + +def validate_finite_tensor_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor byte length does not match its shape" + ) + code = "f" if dtype == DTYPE_F32 else "e" + try: + values = struct.iter_unpack(f"<{code}", payload) + if any(not math.isfinite(value[0]) for value in values): + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor contains non-finite value" + ) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + + +def read_text(reader: Reader, length: int, maximum: int, field: str) -> str: + if length <= 0 or length > maximum: + raise SidecarError(STATUS_RESOURCE_LIMIT, f"invalid {field} length") + try: + value = reader.take(length).decode("utf-8") + except UnicodeDecodeError as error: + raise SidecarError(STATUS_INVALID_REQUEST, f"{field} is not UTF-8") from error + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SidecarError( + STATUS_INVALID_REQUEST, f"{field} contains control characters" + ) + return value + + +def read_tensor( + reader: Reader, rows: int, dimension: int, dtype: int +) -> list[tuple[float, ...]]: + elements = checked_elements(rows, dimension) + if dtype == DTYPE_F32: + code = "f" + element_size = 4 + elif dtype == DTYPE_F16: + code = "e" + element_size = 2 + else: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + raw = reader.take(elements * element_size) + try: + values = struct.unpack(f"<{elements}{code}", raw) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + if not all(math.isfinite(value) for value in values): + raise SidecarError(STATUS_INVALID_REQUEST, "tensor contains non-finite value") + return [ + tuple(values[offset : offset + dimension]) + for offset in range(0, elements, dimension) + ] + + +def decode_resolved_tensor( + payload: bytes, request: ExternalTensorRequest +) -> list[tuple[float, ...]]: + expected_bytes = checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + if len(payload) != expected_bytes: + raise SidecarError( + STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + expected_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(expected_checksum, request.checksum): + raise SidecarError(STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch") + reader = Reader(payload) + tensor = read_tensor(reader, request.rows, request.dimension, request.dtype) + reader.finish() + return tensor + + +def tilemaxsim( + query: list[tuple[float, ...]], document: list[tuple[float, ...]] +) -> float: + if not query or not document: + raise SidecarError(STATUS_INVALID_REQUEST, "tensor must not be empty") + score = 0.0 + try: + for query_vector in query: + best = -math.inf + for document_vector in document: + dot = math.fsum( + left * right + for left, right in zip(query_vector, document_vector, strict=True) + ) + best = max(best, dot) + score += best + except (OverflowError, ValueError) as error: + raise SidecarError(STATUS_COMPUTE_ERROR, str(error)) from error + if not math.isfinite(score): + raise SidecarError(STATUS_COMPUTE_ERROR, "TileMaxSim result is non-finite") + return score + + +def success_response( + request_id: int, + results: Iterable[tuple[int, float]], + version: int = VERSION, +) -> bytes: + results = list(results) + body = bytearray(RESPONSE_FIXED.pack(0, len(results))) + for candidate_id, similarity in results: + body.extend(RESULT.pack(candidate_id, similarity)) + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def error_response( + request_id: int, + status_code: int, + message: str, + version: int = VERSION, +) -> bytes: + encoded = message.encode("utf-8", errors="replace")[:MAX_ERROR_BYTES] + body = ERROR_FIXED.pack(status_code or STATUS_COMPUTE_ERROR, len(encoded)) + encoded + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def validate_request_fixed( + dimension: int, + query_rows: int, + candidate_count: int, + dtype: int, + scoring: int, + reserved: int, + limits: Limits, +) -> None: + if dimension == 0 or dimension > 60_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid tensor dimension") + if query_rows == 0: + raise SidecarError(STATUS_INVALID_REQUEST, "query tensor is empty") + if candidate_count > limits.max_candidates: + raise SidecarError(STATUS_RESOURCE_LIMIT, "too many candidates") + dtype_size(dtype) + if scoring != SCORING_SUM_QUERY_MAX_DOCUMENT_DOT: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported scoring function") + if reserved != 0: + raise SidecarError(STATUS_INVALID_REQUEST, "reserved field must be zero") + if query_rows > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if checked_tensor_bytes(query_rows, dimension, dtype) > limits.max_tensor_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit") + + +def process_inline_request(reader: Reader, limits: Limits) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + candidates: list[tuple[int, list[tuple[float, ...]]]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append((candidate_id, read_tensor(reader, rows, dimension, dtype))) + reader.finish() + return [ + (candidate_id, tilemaxsim(query, document)) + for candidate_id, document in candidates + ] + + +def process_external_request( + reader: Reader, + limits: Limits, + resolver: ExternalTensorResolver | None, + version: int = EXTERNAL_VERSION, +) -> list[tuple[int, float]]: + if version == SCHEDULED_EXTERNAL_VERSION: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, priority, timeout_ms, tenant_length, + ) = reader.unpack(SCHEDULED_EXTERNAL_REQUEST_FIXED) + if not -100 <= priority <= 100 or not 1 <= timeout_ms <= 600_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid scheduler metadata") + else: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + tenant_length = 0 + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + if tenant_length: + read_text(reader, tenant_length, 256, "scheduler tenant") + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + descriptors: list[tuple[int, ExternalTensorRequest]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + descriptors.append( + ( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + + # Validate the complete control frame before performing any external I/O. + if resolver is None: + raise SidecarError( + STATUS_COMPUTE_ERROR, "external tensor resolver is not configured" + ) + results = [] + for candidate_id, descriptor in descriptors: + payload = resolver(descriptor) + if not isinstance(payload, bytes): + raise SidecarError( + STATUS_COMPUTE_ERROR, + "external tensor resolver returned a non-bytes value", + ) + document = decode_resolved_tensor(payload, descriptor) + results.append((candidate_id, tilemaxsim(query, document))) + return results + + +def parse_request_frame( + frame: bytes, + limits: Limits = Limits(), + *, + validate_finite: bool = True, +) -> ParsedRequest: + """Decode and validate a request without resolving or computing tensors. + + The production CUDA sidecar uses this function so the protocol oracle and + deployable executor share one strict wire decoder. All control data and the + inline query are validated before a v2 resolver performs external I/O. + """ + + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in SUPPORTED_VERSIONS: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids: set[int] = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + payload_bytes = checked_tensor_bytes(rows, dimension, dtype) + total_tokens += rows + total_tensor_bytes += payload_bytes + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + payload = reader.take(payload_bytes) + if validate_finite: + validate_finite_tensor_payload(payload, rows, dimension, dtype) + candidates.append(InlineTensorCandidate(candidate_id, rows, payload)) + reader.finish() + return InlineTensorRequest( + request_id, + dimension, + query_rows, + dtype, + query_payload, + tuple(candidates), + ) + + if version == SCHEDULED_EXTERNAL_VERSION: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, scheduler_priority, timeout_ms, tenant_length, + ) = reader.unpack(SCHEDULED_EXTERNAL_REQUEST_FIXED) + if not -100 <= scheduler_priority <= 100 or not 1 <= timeout_ms <= 600_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid scheduler metadata") + else: + ( + dimension, query_rows, candidate_count, dtype, scoring, reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + scheduler_priority = 0 + timeout_ms = 0 + tenant_length = 0 + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + scheduler_tenant = ( + read_text(reader, tenant_length, 256, "scheduler tenant") + if tenant_length + else "__default__" + ) + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append( + ExternalTensorCandidate( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + return ParsedExternalTensorRequest( + request_id, + dimension, + query_rows, + dtype, + model_contract_id, + query_payload, + tuple(candidates), + scheduler_tenant, + scheduler_priority, + timeout_ms, + ) + + +def process_frame( + frame: bytes, + limits: Limits = Limits(), + resolver: ExternalTensorResolver | None = None, +) -> bytes: + request_id = 0 + response_version = VERSION + try: + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if version in SUPPORTED_VERSIONS: + response_version = version + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in SUPPORTED_VERSIONS: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + results = process_inline_request(reader, limits) + else: + results = process_external_request(reader, limits, resolver, version) + return success_response(request_id, results, response_version) + except SidecarError as error: + return error_response(request_id, error.status, str(error), response_version) + except Exception as error: # Keep protocol failures inside the response boundary. + return error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + + +def receive_exact(connection: socket.socket, count: int) -> bytes: + chunks = bytearray() + while len(chunks) < count: + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise SidecarError( + STATUS_INVALID_REQUEST, "connection closed during request" + ) + chunks.extend(chunk) + return bytes(chunks) + + +def handle_connection( + connection: socket.socket, + limits: Limits, + resolver: ExternalTensorResolver | None = None, +) -> None: + request_id = 0 + response_version = VERSION + try: + header = receive_exact(connection, HEADER.size) + _, version, _, request_id, body_len = HEADER.unpack(header) + if version in SUPPORTED_VERSIONS: + response_version = version + if body_len > limits.max_request_bytes - HEADER.size: + response = error_response( + request_id, + STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + response_version, + ) + else: + body = receive_exact(connection, body_len) + response = process_frame(header + body, limits, resolver) + except SidecarError as error: + response = error_response( + request_id, error.status, str(error), response_version + ) + except Exception as error: + response = error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + connection.sendall(response) + + +def remove_stale_socket(path: Path) -> None: + try: + mode = path.lstat().st_mode + except FileNotFoundError: + return + if not stat.S_ISSOCK(mode): + raise RuntimeError(f"refusing to replace non-socket path: {path}") + path.unlink() + + +def serve( + socket_path: Path, + limits: Limits, + socket_mode: int = 0o600, + once: bool = False, + stop: threading.Event | None = None, + resolver: ExternalTensorResolver | None = None, +) -> None: + stop = stop or threading.Event() + remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(16) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + while not stop.is_set(): + try: + connection, _ = listener.accept() + except TimeoutError: + continue + with connection: + handle_connection(connection, limits, resolver) + if once: + break + finally: + listener.close() + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--max-request-bytes", type=int, default=64 * 1024 * 1024) + parser.add_argument("--max-batch-tokens", type=int, default=1_000_000) + parser.add_argument("--max-tensor-bytes", type=int, default=1024 * 1024 * 1024) + parser.add_argument("--max-candidates", type=int, default=65_536) + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + limits = Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + if ( + min( + limits.max_request_bytes, + limits.max_batch_tokens, + limits.max_tensor_bytes, + limits.max_candidates, + ) + <= 0 + ): + parser.error("all limits must be positive") + + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve(args.socket, limits, args.socket_mode, args.once, stop) + + +if __name__ == "__main__": + main() diff --git a/docs/MAXSIM_TENSOR_SOURCES.md b/docs/MAXSIM_TENSOR_SOURCES.md new file mode 100644 index 00000000..1e0ae37f --- /dev/null +++ b/docs/MAXSIM_TENSOR_SOURCES.md @@ -0,0 +1,162 @@ +# MaxSim Tensor-Source Registry + +## Scope + +The registry binds one physical `vchordrq` MaxSim index to the tensor contract +that Phase 3B will use for exact reranking. It stores metadata only. Tensor +credentials and tensor contents never belong in this catalog. + +The registry lifecycle, internal descriptor reader, and restricted Phase 3B +score executor are implemented. The reader validates model contract, shape, +dtype, reference, and checksum before the descriptor v2 IPC encoder accepts a +candidate. The internal Rust resolver calls `vchordrq_maxsim_source_info` and +never reads the private registry table directly. Same-heap bindings resolve +current names back to physical attribute numbers; independent descriptors are +read with one snapshot-consistent SQL join after candidate visibility is +resolved. Registering a source does not change ordinary `@#` execution. + +Search an external full-tensor source with: + +```sql +SELECT public_id, similarity +FROM vchordrq_maxsim_search( + 'visual_page_embeddings_maxsim_idx'::regclass, + $1::halfvec[], + 1024, -- bounded coarse page candidates + 20 -- final exact rows +) +ORDER BY similarity DESC, public_id; +``` + +The current function is deliberately restricted: + +- the registered source must use `external_ref` or `external_relation` storage; +- the named index must use `vector_maxsim_ops` or `halfvec_maxsim_ops` and the + query array type must match it exactly; +- `vchordrq.maxsim_backend` must be `gpu`, with an operations-controlled Unix + socket endpoint; +- `candidate_limit` is between 1 and 65,536 and `top_k` cannot exceed it; +- the caller must have table-level `SELECT` on the indexed heap and, for + `external_relation`, the descriptor relation; projection uses the caller's + MVCC snapshot and PostgreSQL row visibility; +- the caller selects the physical relation/partition. Arbitrary SQL predicates + would require an optional future CustomScan path and are not accepted as + strings; +- HOT root TIDs from the index are resolved through the table AM before the + descriptor query, while backend IDs remain opaque root-candidate IDs. + +The result contains the registered stable `bigint` public ID and positive exact +similarity. Equal scores are ordered by public ID. GPU/sidecar failures fail the +whole statement; this surface does not silently fall back to coarse scores. + +The bundled CUDA sidecar uses `sha256://` tensor references backed by +an operations-configured, per-model local content-addressed cache. GBrain may +populate that cache from its own storage layer; storage credentials and remote +destinations are not registry fields. See +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md) for the file layout, +runtime limits, and deployment contract. + +## Registration + +Only the index owner, a member of the owning role, or a superuser can register +or replace a binding. An independently owned descriptor relation must also be +owned by the caller (or a role it belongs to). Descriptor columns are resolved +to attribute numbers at registration time. + +`external_relation` is the recommended production layout. It avoids widening +and rewriting an already indexed multi-vector heap: the heap keeps its stable +public ID, model contract, and coarse array, while a compact relation keeps one +external descriptor per public ID. + +```sql +CREATE TABLE visual_page_tensor_descriptors ( + document_page_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +SELECT vchordrq_register_maxsim_source( + index_relation => 'visual_page_embeddings_maxsim_idx'::regclass, + model_contract_id => 'colqwen3.5@revision+preprocessing-hash', + storage => 'external_relation', + model_contract_column => 'model_contract_id'::name, + public_id_column => 'document_page_id'::name, + tensor_ref_column => 'tensor_ref'::name, + tensor_rows_column => 'tensor_rows'::name, + tensor_dim_column => 'tensor_dim'::name, + tensor_dtype_column => 'tensor_dtype'::name, + tensor_checksum_column => 'tensor_checksum'::name, + descriptor_relation => 'visual_page_tensor_descriptors'::regclass, + descriptor_public_id_column => 'document_page_id'::name +); +``` + +The index must be a valid, ready, single-key `vchordrq` index using one of the +four MaxSim opclasses. The model-contract and descriptor string columns must be +`NOT NULL text`; the public ID must be `NOT NULL bigint`; rows and dimension +must be `NOT NULL integer`. For `external_relation`, its public ID must have a +valid, non-partial, single-key unique index. Its public ID and five descriptor +columns must be distinct. Missing descriptor rows fail the complete search. + +`external_ref` remains available when all five descriptor columns already live +on the indexed heap. In that layout the two heap metadata columns and five +descriptor columns must all be distinct; omit the two descriptor-relation +arguments. + +For `heap_array`, omit all five external descriptor columns. The model-contract +and public-ID columns are still required: + +```sql +SELECT vchordrq_register_maxsim_source( + index_relation => 'visual_page_embeddings_maxsim_idx'::regclass, + model_contract_id => 'colqwen3.5@revision+preprocessing-hash', + storage => 'heap_array', + model_contract_column => 'model_contract_id'::name, + public_id_column => 'document_page_id'::name +); +``` + +Unregister explicitly with: + +```sql +SELECT vchordrq_unregister_maxsim_source( + 'visual_page_embeddings_maxsim_idx'::regclass +); +``` + +Resolve and revalidate a binding without reading tensor values: + +```sql +SELECT * +FROM vchordrq_maxsim_source_info( + 'visual_page_embeddings_maxsim_idx'::regclass +); +``` + +This function resolves current column names from stored attribute numbers, +rechecks the live index/opclass and every column type, and requires either index +ownership or table `SELECT` privilege. It is the metadata boundary consumed by +the Phase 3B executor; it does not grant access to row values. + +## DDL and Security Semantics + +- Direct DML on `_vchordrq_maxsim_sources` is revoked from `PUBLIC`. +- Registration functions use definer rights with the fixed search path + `pg_catalog, pg_temp`; ownership is checked against the session user and its + role memberships. +- Attribute-number binding survives column renames. +- An `sql_drop` event trigger removes the binding when its index, heap relation, + descriptor relation, or any bound column is dropped. +- A concurrent index replacement gets a new OID and therefore requires explicit + re-registration. This is fail-closed behavior. +- Raw OID bindings are intentionally not portable pg_dump data. Restore tooling + must register sources after restoring tables and indexes. +- Runtime metadata lookup revalidates the relation, opclass, and descriptor + types. Type-altering DDL that does not drop an attribute therefore makes + resolution fail closed until the binding is corrected or re-registered. +- The search API uses the caller's snapshot and permissions and preserves + PostgreSQL row visibility. The registry itself grants no table or tensor + access and does not define an application authorization or routing model. diff --git a/docs/TILEMAXSIM_CUDA_SIDECAR.md b/docs/TILEMAXSIM_CUDA_SIDECAR.md new file mode 100644 index 00000000..c1951caa --- /dev/null +++ b/docs/TILEMAXSIM_CUDA_SIDECAR.md @@ -0,0 +1,229 @@ +# VectorChord Native CUDA TileMaxSim Daemon + +## Scope + +`services/tilemaxsimd` is VectorChord's production CUDA execution service for +external TileMaxSim protocol v2/v3. PostgreSQL remains responsible for MVCC +visibility, hard filters, descriptor validation, and public-ID mapping. The +daemon owns only bounded tensor loading, host/GPU caches, scheduling, and exact +TileMaxSim execution. + +The daemon does not fetch S3 or HTTP objects and does not make authorization, +graph, Fact, community, or application-routing decisions. An application such +as GBrain publishes immutable tensors into a configured node-local shard root. + +The Python sidecar remains a reference and conformance implementation. New +production deployments should use the Rust/CUDA daemon. + +## Fail-closed resource contract + +TileMaxSim is opt-in. Do not start this service when no GPU memory has been +explicitly assigned. Every `--gpu-memory-gb GPU=GB` value is required, uses GiB +rather than bytes, and is reserved with one CUDA allocation during startup. +An invalid device, insufficient free memory, failed pinned-host allocation, or +failed CUDA stream creation terminates the process before either socket becomes +ready. + +The configured allocation is split into a persistent tensor arena and +`--gpu-workspace-gb`. The workspace must be smaller than every GPU allocation. +Host cache and aggregate in-flight request budgets are also configured in GiB. + +## Build the production image + +The build and runtime CUDA images are explicit deployment inputs: + +```shell +docker build \ + --build-arg CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04 \ + --build-arg CUDA_RUNTIME_IMAGE=nvidia/cuda:12.6.3-runtime-ubuntu24.04 \ + -f services/Dockerfile.tilemaxsimd \ + -t vectorchord-tilemaxsimd:local \ + . +``` + +Pin both images by digest in a production build. The container health check +expects the status socket at +`/run/vectorchord/tilemaxsim-status.sock`; mount the socket directory so the +PostgreSQL process can reach the protocol socket. + +## Start the daemon + +```shell +tilemaxsimd \ + --socket /run/vectorchord/tilemaxsim.sock \ + --status-socket /run/vectorchord/tilemaxsim-status.sock \ + --socket-mode 660 \ + --status-socket-mode 660 \ + --gpu-memory-gb 0=20 \ + --gpu-workspace-gb 2 \ + --host-cache-gb 8 \ + --io-pipeline overlap \ + --io-batch-gb 0.05 \ + --max-inflight-request-gb 1 \ + --contract-root 'MODEL_CONTRACT_ID=/var/lib/vectorchord/tensors' \ + --scheduler-policy fair-priority \ + --scheduler-quantum-io-gb 1 \ + --max-connections 256 \ + --max-queued-requests 128 \ + --max-tenant-queued-requests 16 +``` + +Multiple GPU assignments may be supplied. Device selection and the requested +allocation are never inferred from all currently free VRAM. + +## Tutti-inspired tensor I/O pipeline + +`--io-pipeline overlap` is the default for the opt-in CUDA daemon. It uses a +bounded one-batch resolver channel and separate CUDA upload/compute streams to +pipeline three stages: + +```text +immutable shard or L1 resolve N+1 + │ + ├── H2D upload N + │ + └── exact TileMaxSim N-1 +``` + +`--io-batch-gb` bounds the logical payload held by each resolver stage; the +default is 0.05 GiB. The channel has capacity one, so speculative prefetch +cannot grow without bound. GPU cache references keep the current compute pages +non-evictable, and a next batch that cannot acquire pages while computation is +active is deferred and retried after the current batch releases them. The +request therefore remains correct even when its working set exceeds L0. + +The cooperative scheduler also caps the L0-miss payload in one quantum with +`--scheduler-quantum-io-gb` in addition to candidate, token, and FMA limits. +This shortens the period for which a cold request can occupy the I/O/GPU path +before a higher-priority or under-served scheduling domain can run. + +Use `--io-pipeline serial` for a controlled fallback or an A/B comparison. +The current implementation retains ordinary CPU shard reads and pinned H2D +copies. It adopts Tutti's bounded asynchronous pipeline and I/O-aware +scheduling ideas, but it is not yet a GPU io_uring or GPUDirect Storage data +path. SSD and GPU traffic may share a PCIe root complex, so tune the batch on +target hardware and reject configurations that regress the application SLO. + +The example systemd unit and environment file are in `deploy/systemd`. The unit +is deliberately opt-in: enabling it is the operator's explicit TileMaxSim/GPU +configuration. It waits for readiness, reloads immutable shard indexes with +SIGHUP, restarts on failure, and gives SIGTERM shutdown time to stop accepts and +drain admitted work. + +## Health and metrics + +The status Unix socket serves: + +- `GET /livez`: the status server is alive; +- `GET /healthz`: the protocol listener and scheduler are ready; +- `GET /metrics`: bounded Prometheus resource, scheduler, cache, transfer, + storage, latency, timeout, and outcome metrics. GPU labels contain only the + configured numeric device/slot; application scheduling-domain names are + never exported. + +Use the image's dependency-free probe from systemd, Docker, or Kubernetes: + +```shell +tilemaxsimctl \ + --socket /run/vectorchord/tilemaxsim-status.sock \ + --wait-timeout-ms 30000 +``` + +Readiness is removed before shutdown. Unexpected scheduler or status-thread +exit makes the whole daemon fail, so a supervisor cannot keep routing requests +to a process whose CUDA worker has disappeared. + +The most useful production signals are: + +- `tilemaxsim_pending_requests` and `tilemaxsim_scheduler_queue_depth` for + saturation; +- `tilemaxsim_admission_rejections_total` and `tilemaxsim_timeouts_total` for + overload or an undersized deadline; +- `tilemaxsim_gpu_cache_events_total`, `tilemaxsim_gpu_h2d_bytes_total`, and + `tilemaxsim_storage_read_bytes_total` for cache churn and cold-load traffic; +- `tilemaxsim_io_pipeline_batches_total`, + `tilemaxsim_io_pipeline_resolve_bytes_total`, and + `tilemaxsim_io_pipeline_seconds_total` for resolve/upload/compute time and + time hidden by overlap; +- `tilemaxsim_gpu_cache_bytes` for free space, largest free extent, payload, + allocator waste, and pinned capacity; +- `tilemaxsim_host_cache_bytes` and `tilemaxsim_host_cache_events_total` for + the L1 cache; +- `tilemaxsim_request_duration_seconds`, `tilemaxsim_queue_duration_seconds`, + and `tilemaxsim_gpu_duration_seconds` for rate-derived mean latency; +- `tilemaxsim_requests_admitted_total{priority_class=...}` and + `tilemaxsim_scheduler_requeues_total` for priority/cooperative scheduling. + +At minimum, alert when readiness is zero, any admission-rejection rate is +nonzero under expected load, timeout ratio exceeds the application SLO, queue +depth remains above 80% of its limit, or cache admission rejections rise while +`largest_free_extent / free` is small. A rising miss/eviction/H2D rate with +stable traffic means the assigned GPU cache is too small or the access set has +poor locality; it is not a PostgreSQL HNSW symptom. + +## Immutable tensor storage + +Each model contract maps to one absolute shard root. A protocol descriptor uses +a content-addressed reference and matching checksum: + +```text +tensor_ref = sha256://0123...cdef +tensor_checksum = sha256:0123...cdef +tensor_rows = 747 +tensor_dim = 320 +tensor_dtype = float16 +``` + +Online writers publish canonical tensor bytes through VectorChord's publisher, +so application code never owns or duplicates the disk layout: + +```shell +install -d -m 0750 /var/lib/vectorchord/tensors +tilemaxsimctl publish-object \ + --root /var/lib/vectorchord/tensors \ + --rows 747 \ + --dimension 320 \ + --dtype float16 \ + --expected-sha256 0123...cdef < document.tensor.f16 +``` + +The command fsyncs a temporary object, publishes it by immutable hard link, and +returns the JSON descriptor. Repeating the same publication is idempotent. A +daemon that already has the contract root open can read the new content address +without restart or SIGHUP. SIGHUP is needed only for a newly published bulk +shard index. It never changes the configured contract-to-root mapping. The +daemon validates contract, digest, shape, dtype, length, and checksum before GPU +admission. + +## Batching, scheduling, and failure semantics + +- PostgreSQL splits one logical candidate set into protocol-bounded batches and + merges a deterministic global top-k. +- All batches share the caller's one logical deadline; a new IPC round trip does + not refresh it. +- Admission is bounded by connections, queued requests, per-scheduling-domain + queue depth, and aggregate frame bytes. +- Long batches re-enter the selected scheduler between candidate/token quanta. + This is cooperative preemption between CUDA kernels; a running kernel cannot + be interrupted safely. +- `fair-priority` preserves weighted fairness within the configured priority + band. Strict global priority is available with `--scheduler-policy priority`. +- Disconnects and deadlines are checked between quanta. CUDA, shard, checksum, + workspace, or timeout failure fails the logical request; PostgreSQL does not + expose partial results. + +Scheduling-domain strings affect latency ordering and cache quotas only. They +are not authorization evidence. + +## Verification + +```shell +cargo test --manifest-path services/tilemaxsimd/Cargo.toml --locked +python3 -m unittest -v services.test_tilemaxsim_rust_daemon +``` + +The Python integration suite exercises real Unix sockets and, when CUDA is +available, multi-GPU, resident-cache, oversized-working-set, concurrent-reader, +health, and score-equivalence paths. Application-level corpus recall and +end-to-end latency remain a release acceptance gate rather than an inference +from synthetic service tests. diff --git a/docs/TILEMAXSIM_IPC_V1.md b/docs/TILEMAXSIM_IPC_V1.md new file mode 100644 index 00000000..274258c2 --- /dev/null +++ b/docs/TILEMAXSIM_IPC_V1.md @@ -0,0 +1,160 @@ +# VectorChord TileMaxSim IPC Protocol v1 + +## Scope + +This protocol connects a PostgreSQL VectorChord backend to a local TileMaxSim +GPU sidecar. It carries full query/page tensors for Phase 3A. It is not a +network authentication protocol and must be exposed only through an +operations-controlled Unix socket. + +All integers and floating-point bit patterns are little-endian. The sidecar +must read and write exact frame lengths. Partial results are forbidden. + +## Common Frame Header + +Every frame starts with 24 bytes: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | bytes | ASCII magic `VCTM` | +| 4 | 2 | `u16` | protocol version, currently `1` | +| 6 | 2 | `u16` | message kind: `1` request, `2` response | +| 8 | 8 | `u64` | request ID, echoed unchanged by the response | +| 16 | 8 | `u64` | body length, excluding the 24-byte header | + +Unknown versions, message kinds, request IDs, and inconsistent lengths are +fatal protocol errors. + +## Rerank Request Body + +The fixed portion is 16 bytes: + +| Offset in body | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | `u32` | tensor dimension | +| 4 | 4 | `u32` | query row/token count | +| 8 | 4 | `u32` | candidate page count | +| 12 | 1 | `u8` | dtype: `1` f32, `2` IEEE f16 | +| 13 | 1 | `u8` | scoring: `1` = sum of query-token max document dot products | +| 14 | 2 | `u16` | reserved, must be zero | + +The fixed portion is followed by the query tensor in row-major order. Each +candidate then has: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | opaque request-local candidate ID | +| 4 | `u32` | page tensor row/token count | +| `rows * dimension * dtype_size` | bytes | row-major page tensor | + +Candidate IDs are assigned densely from zero in v1, but the sidecar must treat +them as opaque and echo them unchanged. Heap CTIDs and public GBrain page IDs +are not exposed to the sidecar. + +Every query and candidate tensor in one request has the same dtype and +dimension. v1 accepts f32 and f16 only. RaBitQ arrays must use CPU exact or a +future protocol version that defines their representation. + +## Successful Response Body + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | status `0` | +| 4 | `u32` | result count; must equal request candidate count | + +Each result then contains: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | echoed candidate ID | +| 4 | `f32` bits | positive MaxSim similarity | + +Each requested candidate must appear exactly once. Unknown IDs, duplicate IDs, +missing IDs, non-finite scores, and trailing bytes fail the entire SQL query. +VectorChord converts public positive similarity to its SQL distance convention +by negating the value before ascending sort. Equal exact distances are ordered +deterministically by the candidate's internal heap key after response IDs have +been mapped back inside the backend. + +## Error Response Body + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | nonzero sidecar-defined status | +| 4 | `u32` | UTF-8 error-message length | +| `length` | bytes | UTF-8 diagnostic, maximum 64 KiB | + +The status namespace is reserved for the sidecar implementation in v1. Any +nonzero status fails `gpu` mode. In `auto` mode it triggers a complete CPU exact +rerank of the original candidate set; no partial GPU scores are reused. + +## Limits and Timeouts + +Before connecting, VectorChord enforces: + +- `vchordrq.maxsim_candidate_limit`; +- `vchordrq.maxsim_gpu_max_batch_tokens`, including query and page tokens; +- `vchordrq.maxsim_gpu_max_batch_bytes`, including the request frame. + +`vchordrq.maxsim_gpu_timeout_ms` is one overall connect-plus-write-plus-read +deadline. Connection and backend I/O poll PostgreSQL interrupts at least every +50 ms. The response is also bounded from the expected candidate count and the +64 KiB error limit. + +v1 sends one bounded request. Splitting a candidate set into multiple GPU +batches is future work and must preserve one overall SQL deadline and +all-or-nothing result semantics. + +External immutable tensor descriptors use the separate Phase 3B +[`TILEMAXSIM_IPC_V2`](TILEMAXSIM_IPC_V2.md) contract. They are never smuggled +into a v1 inline-tensor frame. + +## Reference Sidecar + +`devtools/tilemaxsim_reference_sidecar.py` is a dependency-free, single-threaded +CPU protocol oracle for v1 and resolver-injected v2 tests. It is intended for +codec and end-to-end development only; it is not a production or GPU executor. +The production-oriented CUDA implementation is documented in +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md). + +Start one request/response cycle with: + +```text +python3 devtools/tilemaxsim_reference_sidecar.py \ + --socket /tmp/vectorchord-tilemaxsim.sock \ + --once +``` + +Then configure the PostgreSQL session as an administrator and select the GPU +backend in the query session: + +```sql +SET vchordrq.maxsim_gpu_endpoint = '/tmp/vectorchord-tilemaxsim.sock'; +SET vchordrq.maxsim_candidate_limit = 256; +SET vchordrq.maxsim_backend = 'gpu'; +``` + +The endpoint setting is `SUSET`. The reference sidecar creates its socket with +mode `0600` by default and refuses to replace a non-socket filesystem path. +Its default request limit is intentionally 64 MiB, lower than the extension's +configurable production ceiling. + +Run its protocol and live-socket tests with: + +```text +python3 -m unittest -v devtools/test_tilemaxsim_reference_sidecar.py +``` + +## Sidecar Conformance Cases + +A conforming implementation must be tested for: + +- f32 and f16 score equivalence with CPU exact; +- reordered response IDs; +- duplicate, missing, and unknown IDs; +- invalid magic/version/kind/request ID/body length; +- truncated and trailing data; +- NaN and infinite similarity; +- explicit error response; +- connection close during request and response; +- queue/compute duration exceeding the overall deadline. diff --git a/docs/TILEMAXSIM_IPC_V2.md b/docs/TILEMAXSIM_IPC_V2.md new file mode 100644 index 00000000..d5e211df --- /dev/null +++ b/docs/TILEMAXSIM_IPC_V2.md @@ -0,0 +1,173 @@ +# VectorChord TileMaxSim IPC Protocol v2 and v3 + +## Scope + +Version 2 carries an inline query tensor plus immutable external tensor +descriptors. It is the Phase 3B transport for a coarse/sketch `vchordrq` index +whose final score comes from a different full tensor. It does not change the +meaning of ordinary SQL `@#` and is not enabled by source registration alone. + +Version 3 is the backward-compatible scheduled form. It adds an authenticated +upstream scheduling domain, priority, and end-to-end timeout. These fields +control resource scheduling only; they are not authorization claims and never +replace PostgreSQL row visibility or application ACL checks. + +The transport is an operations-controlled Unix socket. It is not a network +authentication protocol. All integers and floating-point bit patterns are +little-endian. The sidecar must read and write exact frame lengths; partial +results are forbidden. + +## Common Header + +The 24-byte header is the same shape as v1: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | bytes | ASCII magic `VCTM` | +| 4 | 2 | `u16` | protocol version `2` | +| 6 | 2 | `u16` | message kind: `1` request, `2` response | +| 8 | 8 | `u64` | request ID, echoed unchanged | +| 16 | 8 | `u64` | body length, excluding the header | + +The response must echo the request version. Unknown versions, message kinds, +request IDs, or inconsistent lengths fail the complete request. + +## Version 3 Scheduling Extension + +For a version 3 request, the v2 fixed request body is immediately followed by +12 additional fixed bytes: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `i32` | priority from -100 through 100; higher values are more urgent | +| 4 | `u32` | end-to-end timeout in milliseconds, from 1 through 600000 | +| 4 | `u32` | scheduling-tenant UTF-8 byte length | + +The variable body then contains model-contract bytes, scheduling-tenant bytes, +the query tensor, and the candidate descriptors, in that order. The tenant is +nonempty, limited to 256 bytes, and cannot contain control characters. + +The timeout begins when the daemon accepts the connection, not when GPU work +starts. A server-side timeout may shorten it. Priority affects queue order only; +it cannot expand the request deadline, cache quota, candidate scope, or access +rights. Protocol v2 requests enter the default scheduling domain at priority +zero and use the server timeout. + +## External Rerank Request + +The fixed request body is 20 bytes: + +| Offset | Size | Type | Meaning | +| --- | ---: | --- | --- | +| 0 | 4 | `u32` | tensor dimension | +| 4 | 4 | `u32` | query row/token count | +| 8 | 4 | `u32` | candidate count | +| 12 | 1 | `u8` | dtype: `1` f32, `2` IEEE f16 | +| 13 | 1 | `u8` | scoring: `1` = sum of query-token max document dot products | +| 14 | 2 | `u16` | reserved, must be zero | +| 16 | 4 | `u32` | UTF-8 model-contract byte length | + +The fixed body is followed, in order, by: + +1. the model-contract UTF-8 bytes; +2. the inline query tensor in row-major order; +3. exactly `candidate_count` descriptor entries. + +Each descriptor entry is: + +| Size | Type | Meaning | +| ---: | --- | --- | +| 4 | `u32` | opaque request-local candidate ID | +| 4 | `u32` | full tensor row/token count | +| 4 | `u32` | tensor-reference UTF-8 byte length | +| 4 | `u32` | checksum UTF-8 byte length | +| variable | bytes | tensor reference | +| variable | bytes | checksum | + +All tensors in one request have the fixed body's dtype and dimension. The +backend validates each registered row descriptor against them before IPC. The +model contract and every descriptor field are nonempty, bounded UTF-8 without +control characters. Current bounds are 512 bytes for the contract and +checksum, and 4096 bytes for the reference. + +The checksum encoding is `sha256:` followed by 64 lowercase hexadecimal +digits. It covers the canonical row-major scalar payload after any immutable +object envelope is decoded. This makes verification independent of an object +store's metadata and compression. A resolver for formats such as safetensors +must decode the selected tensor, validate its declared shape/dtype, produce the +canonical payload, and then verify this digest before compute. + +Candidate IDs are opaque. PostgreSQL heap CTIDs, registered public document +IDs, application routing identifiers, and credentials are never encoded. +VectorChord maps response IDs back to heap keys inside the backend. + +## Resolution and Security + +The sidecar owns descriptor resolution. A production resolver must: + +- allowlist schemes, buckets/namespaces, and immutable reference forms; +- obtain credentials from operations-managed sidecar configuration, never from + PostgreSQL rows or this protocol; +- reject redirects or network destinations outside its allowlist; +- enforce the model contract, dtype, shape, and checksum before GPU access; +- bound queue, fetch, decode, host-to-device, and compute work by the one SQL + request deadline. + +The complete control frame must be parsed and validated before external I/O. +An absent resolver fails closed. The dependency-free reference sidecar exposes +v2 resolution only through an injected callback; its CLI deliberately has no +filesystem, HTTP, or object-store resolver. + +## Limits + +Both peers enforce: + +- candidate count; +- total query plus declared document tokens; +- total declared canonical tensor bytes; +- control-frame bytes; +- one overall connect/write/read deadline. + +`vchordrq.maxsim_gpu_max_batch_tokens` bounds total declared tokens. +`vchordrq.maxsim_gpu_max_batch_bytes` bounds both the control frame and total +declared canonical tensor bytes independently. A small descriptor frame cannot +therefore authorize unbounded object loads. + +Version 2 currently sends one bounded batch. Future splitting must retain one +overall deadline and all-or-nothing result semantics. + +## Response + +The success and error bodies are identical to v1 except that the header version +echoes 2 or 3. A success body begins with status `0` and a result count equal to the +number of accepted descriptors, followed by `(candidate_id u32, similarity +f32)` pairs. Every candidate appears exactly once. Unknown, duplicate, missing, +non-finite, or trailing results fail the whole SQL operation. + +Similarity is positive public MaxSim. VectorChord negates it to its ascending +distance convention and uses the internal heap key as a deterministic tie +breaker for the query execution. + +## Current Integration State + +The Rust descriptor source, strict v2/v3 encoder/decoder, and reference-sidecar +protocol oracle are implemented and unit tested. The internal runtime resolver +also consumes the privilege-aware SQL registry boundary and produces physical +attribute bindings. The restricted `vchordrq_maxsim_search` score API now sends +validated, SQL-visible external descriptors over v2 and returns exact positive +similarity. Ordinary `@#` scans keep their existing stored-array semantics. + +The native Rust/CUDA daemon additionally provides bounded global and per-tenant +admission, fair/priority/fair-priority policies, request aging, candidate/token +quanta, deadline and disconnect cancellation between CUDA launches, per-tenant +cache caps, optional reservations, and immutable-shard reload on `SIGHUP`. +An optional operations-only Unix status socket exposes HTTP `/healthz` and +Prometheus `/metrics`; it is separate from this binary scoring protocol. + +The dependency-free reference sidecar remains a protocol oracle. The separate +[`TILEMAXSIM_CUDA_SIDECAR`](TILEMAXSIM_CUDA_SIDECAR.md) implementation executes +v1/v2 on CUDA and provides a per-model, allowlisted, content-addressed local +resolver with checksum, resource-limit, deadline, backpressure, and structured +metric enforcement. GBrain remains responsible for populating that immutable +node-local cache from its storage system. Production acceptance still requires +the committed GBrain corpus recall, latency, concurrency, and failure matrix. diff --git a/licenses/LICENSE.AGPLv3 b/licenses/LICENSE.AGPLv3 new file mode 100644 index 00000000..0ad25db4 --- /dev/null +++ b/licenses/LICENSE.AGPLv3 @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published + by the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/licenses/LICENSE.ELv2 b/licenses/LICENSE.ELv2 new file mode 100644 index 00000000..efa8dfbe --- /dev/null +++ b/licenses/LICENSE.ELv2 @@ -0,0 +1,91 @@ +Elastic License 2.0 + +## Acceptance + +By using the software, you agree to all of the terms and conditions below. + +## Copyright License + +The licensor grants you a non-exclusive, royalty-free, worldwide, +non-sublicensable, non-transferable license to use, copy, distribute, make +available, and prepare derivative works of the software, in each case subject to +the limitations and conditions below. + +## Limitations + +You may not provide the software to third parties as a hosted or managed +service, where the service provides users with access to any substantial set of +the features or functionality of the software. + +You may not move, change, disable, or circumvent the license key functionality +in the software, and you may not remove or obscure any functionality in the +software that is protected by the license key. + +You may not alter, remove, or obscure any licensing, copyright, or other notices +of the licensor in the software. Any use of the licensor’s trademarks is subject +to applicable law. + +## Patents + +The licensor grants you a license, under any patent claims the licensor can +license, or becomes able to license, to make, have made, use, sell, offer for +sale, import and have imported the software, in each case subject to the +limitations and conditions in this license. This license does not cover any +patent claims that you cause to be infringed by modifications or additions to +the software. If you or your company make any written claim that the software +infringes or contributes to infringement of any patent, your patent license for +the software granted under these terms ends immediately. If your company makes +such a claim, your patent license ends immediately for work on behalf of your +company. + +## Notices + +You must ensure that anyone who gets a copy of any part of the software from you +also gets a copy of these terms. + +If you modify the software, you must include in any modified copies of the +software prominent notices stating that you have modified the software. + +## No Other Rights + +These terms do not imply any licenses other than those expressly granted in +these terms. + +## Termination + +If you use the software in violation of these terms, such use is not licensed, +and your licenses will automatically terminate. If the licensor provides you +with a notice of your violation, and you cease all violation of this license no +later than 30 days after you receive that notice, your licenses will be +reinstated retroactively. However, if you violate these terms after such +reinstatement, any additional violation of these terms will cause your licenses +to terminate automatically and permanently. + +## No Liability + +*As far as the law allows, the software comes as is, without any warranty or +condition, and the licensor will not be liable to you for any damages arising +out of these terms or the use or nature of the software, under any kind of +legal claim.* + +## Definitions + +The **licensor** is the entity offering these terms, and the **software** is the +software the licensor makes available under these terms, including any portion +of it. + +**you** refers to the individual or entity agreeing to these terms. + +**your company** is any legal entity, sole proprietorship, or other kind of +organization that you work for, plus all organizations that have control over, +are under the control of, or are under common control with that +organization. **control** means ownership of substantially all the assets of an +entity, or the power to direct its management and policies by vote, contract, or +otherwise. Control can be direct or indirect. + +**your licenses** are all the licenses granted to you for the software under +these terms. + +**use** means anything you do with the software requiring one of your licenses. + +**trademark** means trademarks, service marks, and similar rights. diff --git a/rabbithole.control b/rabbithole.control deleted file mode 100644 index 74b4cc1d..00000000 --- a/rabbithole.control +++ /dev/null @@ -1,7 +0,0 @@ -comment = 'rabbithole: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' -default_version = '@CARGO_VERSION@' -module_pathname = '$libdir/rabbithole' -relocatable = false -superuser = true -schema = rabbithole -requires = 'vectors' diff --git a/rust-toolchain.toml b/rust-toolchain.toml deleted file mode 100644 index edfb9578..00000000 --- a/rust-toolchain.toml +++ /dev/null @@ -1,2 +0,0 @@ -[toolchain] -channel = "nightly-2024-09-14" diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 00000000..c1578aaf --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1 @@ +imports_granularity = "Module" diff --git a/services/Dockerfile.postgres b/services/Dockerfile.postgres new file mode 100644 index 00000000..42a0db71 --- /dev/null +++ b/services/Dockerfile.postgres @@ -0,0 +1,36 @@ +ARG POSTGRES_VERSION=16 +ARG POSTGRES_RUNTIME_IMAGE=pgvector/pgvector:pg16@sha256:1d533553fefe4f12e5d80c7b80622ba0c382abb5758856f52983d8789179f0fb + +FROM rust:1.96.0-bookworm@sha256:5e2214abe154fe26e39f64488952e5c991eeed1d6d6da7cc8381ae83927f0cfc AS build + +ARG POSTGRES_VERSION +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates clang curl make postgresql-common \ + && /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + postgresql-server-dev-${POSTGRES_VERSION} \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /build/vectorchord +COPY . . +ENV CC=clang +RUN make PG_CONFIG=/usr/lib/postgresql/${POSTGRES_VERSION}/bin/pg_config build \ + && make PG_CONFIG=/usr/lib/postgresql/${POSTGRES_VERSION}/bin/pg_config \ + DESTDIR=/opt/vectorchord install + +FROM ${POSTGRES_RUNTIME_IMAGE} + +ARG POSTGRES_VERSION +COPY --from=build /opt/vectorchord/ / + +# The upstream PostgreSQL entrypoint only uses gosu to become the postgres +# account. Replace that call with Debian's maintained setpriv and remove the +# statically linked Go binary, whose embedded standard library otherwise leaves +# scanners unable to distinguish unreachable Go CVEs from exploitable code. +RUN postgres --version | grep -Eq "PostgreSQL\\) ${POSTGRES_VERSION}([.[:space:]]|$)" \ + && sed -ri \ + 's!exec gosu postgres "\$BASH_SOURCE" "\$@"!exec setpriv --reuid=postgres --regid=postgres --init-groups "\$BASH_SOURCE" "\$@"!' \ + /usr/local/bin/docker-entrypoint.sh \ + && ! grep -q 'exec gosu postgres' /usr/local/bin/docker-entrypoint.sh \ + && rm -f /usr/local/bin/gosu diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim new file mode 100644 index 00000000..2bf87252 --- /dev/null +++ b/services/Dockerfile.tilemaxsim @@ -0,0 +1,15 @@ +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +WORKDIR /opt/vectorchord + +COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py +COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py +COPY services/tilemaxsim_gpu_cache.py services/tilemaxsim_gpu_cache.py +COPY services/tilemaxsim_shard.py services/tilemaxsim_shard.py +COPY services/tilemaxsim_triton.py services/tilemaxsim_triton.py + +ENV PYTHONPATH=/opt/vectorchord \ + PYTHONUNBUFFERED=1 + +ENTRYPOINT ["python3", "-m", "services.tilemaxsim_cuda_sidecar"] diff --git a/services/Dockerfile.tilemaxsimd b/services/Dockerfile.tilemaxsimd new file mode 100644 index 00000000..6665366e --- /dev/null +++ b/services/Dockerfile.tilemaxsimd @@ -0,0 +1,52 @@ +ARG CUDA_DEVEL_IMAGE=nvidia/cuda:12.6.3-devel-ubuntu24.04@sha256:392c0df7b577ecae17a17f6ba7f2009c217bb4422f8431c053ae9af61a8c148a +ARG RUNTIME_IMAGE=ubuntu:24.04@sha256:4fbb8e6a8395de5a7550b33509421a2bafbc0aab6c06ba2cef9ebffbc7092d90 + +FROM ${CUDA_DEVEL_IMAGE} AS build + +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential ca-certificates curl \ + && rm -rf /var/lib/apt/lists/* +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal --default-toolchain 1.96.0 + +ENV PATH=/root/.cargo/bin:${PATH} +WORKDIR /build/tilemaxsimd +COPY services/tilemaxsimd/ . +RUN cargo build --release --locked + +FROM build AS test +RUN rustup component add clippy rustfmt \ + && cargo fmt --check \ + && cargo clippy --all-targets --locked -- -D warnings \ + && cargo test --locked + +FROM ${RUNTIME_IMAGE} AS runtime-base + +# nvcc links cudart statically for this daemon. The release binary needs only +# glibc/libstdc++ from Ubuntu; libcuda and the devices are injected by the +# NVIDIA container runtime. Keeping CUDA toolkits out of the final image avoids +# shipping an unused 1.4 GiB runtime and materially reduces its attack surface. +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get upgrade -y \ + && rm -rf /var/lib/apt/lists/* + +FROM runtime-base + +COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimd /usr/local/bin/tilemaxsimd +COPY --from=build /build/tilemaxsimd/target/release/tilemaxsimctl /usr/local/bin/tilemaxsimctl +RUN /usr/local/bin/tilemaxsimd --help >/dev/null \ + && /usr/local/bin/tilemaxsimctl --help >/dev/null \ + && groupadd --system vectorchord \ + && useradd --system --gid vectorchord --home-dir /nonexistent \ + --shell /usr/sbin/nologin vectorchord \ + && install -d -o vectorchord -g vectorchord -m 0750 \ + /run/vectorchord /var/lib/vectorchord + +USER vectorchord:vectorchord + +STOPSIGNAL SIGTERM +HEALTHCHECK --interval=10s --timeout=2s --start-period=5m --retries=3 \ + CMD ["/usr/local/bin/tilemaxsimctl", "--socket", "/run/vectorchord/tilemaxsim-status.sock"] + +ENTRYPOINT ["/usr/local/bin/tilemaxsimd"] diff --git a/services/benchmark_full_corpus_tilemaxsim.py b/services/benchmark_full_corpus_tilemaxsim.py new file mode 100644 index 00000000..42e5b0be --- /dev/null +++ b/services/benchmark_full_corpus_tilemaxsim.py @@ -0,0 +1,190 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark exact full-corpus TileMaxSim through the native CUDA daemon.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_tilemaxsim_ablation import encode_frame, request_round_trip + +MAX_BATCH_CANDIDATES = 65_536 +MAX_BATCH_TOKENS = 1_000_000 +MAX_BATCH_TENSOR_BYTES = 1024**3 + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def load_jsonl(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + records = [json.loads(line) for line in stream if line.strip()] + if not records: + raise ValueError(f"empty JSONL file: {path}") + return records + + +def descriptor_batches( + descriptors: list[dict[str, object]], query_rows: int +) -> list[tuple[int, list[dict[str, object]]]]: + batches: list[tuple[int, list[dict[str, object]]]] = [] + start = 0 + current: list[dict[str, object]] = [] + tokens = query_rows + tensor_bytes = query_rows * int(descriptors[0]["tensor_dim"]) * 2 + for descriptor in descriptors: + rows = int(descriptor["tensor_rows"]) + scalar_bytes = 2 if descriptor["tensor_dtype"] == "float16" else 4 + payload_bytes = rows * int(descriptor["tensor_dim"]) * scalar_bytes + would_overflow = current and ( + len(current) == MAX_BATCH_CANDIDATES + or tokens + rows > MAX_BATCH_TOKENS + or tensor_bytes + payload_bytes > MAX_BATCH_TENSOR_BYTES + ) + if would_overflow: + batches.append((start, current)) + start += len(current) + current = [] + tokens = query_rows + tensor_bytes = query_rows * int(descriptor["tensor_dim"]) * scalar_bytes + current.append(descriptor) + tokens += rows + tensor_bytes += payload_bytes + if current: + batches.append((start, current)) + return batches + + +def recall(expected: np.ndarray, actual: list[int], top_k: int) -> float: + return len(set(expected[:top_k].tolist()).intersection(actual[:top_k])) / top_k + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--query-dataset", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--gold", required=True, type=Path) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + descriptors = load_jsonl(args.descriptor_manifest) + queries = json.loads(args.query_dataset.read_text(encoding="utf-8")) + if not isinstance(queries, list) or not queries: + raise ValueError("query dataset must be a nonempty JSON array") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + queries = queries[: args.query_limit or None] + gold = np.load(args.gold, allow_pickle=False) + query_latencies: list[float] = [] + round_trip_latencies: list[float] = [] + batch_latencies: list[float] = [] + cases = [] + + for query_index, query_record in enumerate(queries): + query_id = query_record.get("query_id") + if not isinstance(query_id, str): + raise ValueError(f"query {query_index} has no query_id") + query = np.load(args.query_embeddings / f"{query_id}.npy", allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def sql_halfvec_array(tensor: np.ndarray) -> str: + if tensor.ndim != 2 or tensor.shape[1] <= 0 or tensor.shape[0] <= 0: + raise ValueError("query tensor must have shape [rows, dimension]") + if tensor.dtype not in (np.dtype("float16"), np.dtype("float32")): + raise ValueError("query tensor must use float16 or float32") + if not np.isfinite(tensor).all(): + raise ValueError("query tensor contains non-finite values") + vectors = [] + for row in tensor: + value = "[" + ",".join(format(float(item), ".8g") for item in row) + "]" + vectors.append("'" + value + "'::halfvec") + return "ARRAY[" + ",".join(vectors) + "]" + + +def load_manifest(path: Path) -> list[str]: + page_keys = [] + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + record = json.loads(line) + page_key = record.get("page_key") + if not isinstance(page_key, str) or not PAGE_KEY.fullmatch(page_key): + raise ValueError(f"invalid page_key at manifest line {line_number}") + page_keys.append(page_key) + if not page_keys: + raise ValueError("page manifest is empty") + return page_keys + + +def load_queries( + results_path: Path, embeddings: Path, limit: int +) -> list[tuple[str, Path]]: + results = json.loads(results_path.read_text(encoding="utf-8")) + cases = results.get("queries", {}).get("cases") + if not isinstance(cases, list) or not cases: + raise ValueError("results JSON has no queries.cases") + queries = [] + for case in cases[: limit or None]: + query_id = case.get("query_id") + if not isinstance(query_id, str) or not re.fullmatch(r"[0-9a-f]+", query_id): + raise ValueError("invalid query ID in results JSON") + path = embeddings / f"{query_id}.npy" + if not path.is_file(): + raise ValueError(f"missing query embedding: {path}") + queries.append((query_id, path)) + return queries + + +def recall(expected: list[str], actual: list[str], top_k: int) -> float: + expected_set = set(expected[:top_k]) + return len(expected_set.intersection(actual[:top_k])) / top_k + + +def candidate_recall(expected: list[str], candidates: list[str], top_k: int) -> float: + expected_set = set(expected[:top_k]) + return len(expected_set.intersection(candidates)) / top_k + + +def parse_profile(stderr: str) -> dict[str, int] | None: + profiles = [] + for line in stderr.splitlines(): + marker = line.find(PROFILE_PREFIX) + if marker < 0: + continue + payload = line[marker + len(PROFILE_PREFIX) :].strip() + profile = json.loads(payload) + if not isinstance(profile, dict) or profile.get("schema_version") != 1: + raise RuntimeError("unexpected MaxSim profile schema") + if any(not isinstance(value, int) for value in profile.values()): + raise RuntimeError("MaxSim profile values must be integers") + profiles.append(profile) + if len(profiles) > 1: + raise RuntimeError("query emitted multiple MaxSim profiles") + return profiles[0] if profiles else None + + +def execute_query( + psql_command: list[str], + table: str, + index: str | None, + query_sql: str, + endpoint: str, + probes: str, + refine: int, + candidate_limit: int, + timeout_ms: int, + profile: bool, +) -> tuple[list[str], float, dict[str, int] | None]: + if index is None: + search_sql = f""" +SELECT page_key +FROM {table} +ORDER BY embedding @# {query_sql} +LIMIT {candidate_limit}; +""" + else: + search_sql = f""" +SELECT p.page_key +FROM vchordrq_maxsim_search( + '{index}'::regclass, + {query_sql}, + {candidate_limit}, + {candidate_limit} + ) WITH ORDINALITY AS r(public_id, similarity, result_order) +JOIN {table} AS p ON p.id = r.public_id +ORDER BY r.result_order; +""" + sql = f""" +SET statement_timeout = '{timeout_ms}ms'; +SET enable_seqscan = off; +SET vchordrq.probes = '{probes}'; +SET vchordrq.maxsim_refine = {refine}; +SET vchordrq.maxsim_candidate_limit = {candidate_limit}; +SET vchordrq.maxsim_gpu_endpoint = '{endpoint}'; +SET vchordrq.maxsim_gpu_timeout_ms = {timeout_ms}; +SET vchordrq.maxsim_backend = 'gpu'; +SET vchordrq.maxsim_profile = {"on" if profile else "off"}; +{search_sql} +""" + started = time.perf_counter() + result = subprocess.run( + [*psql_command, "-X", "-q", "-A", "-t", "-v", "ON_ERROR_STOP=1"], + input=sql, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=max(5.0, timeout_ms / 1000.0 + 5.0), + check=False, + ) + latency_ms = (time.perf_counter() - started) * 1000.0 + if result.returncode != 0: + diagnostic = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"psql failed: {diagnostic}") + page_keys = [line.strip() for line in result.stdout.splitlines() if line.strip()] + invalid = [page_key for page_key in page_keys if not PAGE_KEY.fullmatch(page_key)] + if invalid: + raise RuntimeError(f"psql returned non-page-key output: {invalid[0]!r}") + if len(page_keys) != len(set(page_keys)): + raise RuntimeError("Phase 3 query returned duplicate page keys") + query_profile = parse_profile(result.stderr) + if profile and query_profile is None: + raise RuntimeError("query did not emit a MaxSim profile") + return page_keys, latency_ms, query_profile + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--results-json", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--gold", required=True, type=Path) + parser.add_argument("--psql-command", default="psql") + parser.add_argument("--table", default="kb_colqwen_pages") + parser.add_argument( + "--index", + help="use the Phase 3B external-tensor API with this MaxSim index", + ) + parser.add_argument("--endpoint", required=True) + parser.add_argument("--probes", default="8") + parser.add_argument("--refine", type=int, default=512) + parser.add_argument("--candidate-limit", type=positive_int, default=100) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--timeout-ms", type=positive_int, default=60000) + parser.add_argument( + "--profile", + action="store_true", + help="enable vchordrq.maxsim_profile and include phase metrics", + ) + args = parser.parse_args() + + if not IDENTIFIER.fullmatch(args.table): + parser.error( + "--table must be an unquoted SQL identifier, optionally schema-qualified" + ) + if args.index is not None and not IDENTIFIER.fullmatch(args.index): + parser.error( + "--index must be an unquoted SQL identifier, optionally schema-qualified" + ) + if not args.endpoint.startswith("/") or "'" in args.endpoint: + parser.error("--endpoint must be an absolute Unix-socket path without quotes") + if args.refine < 0: + parser.error("--refine must be nonnegative") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + if args.profile and args.index is None: + parser.error("--profile requires --index") + if not re.fullmatch(r"[0-9]+(,[0-9]+)*", args.probes): + parser.error("--probes must be a comma-separated list of nonnegative integers") + psql_command = shlex.split(args.psql_command) + if not psql_command: + parser.error("--psql-command must not be empty") + + page_keys = load_manifest(args.manifest) + queries = load_queries(args.results_json, args.query_embeddings, args.query_limit) + gold = np.load(args.gold, allow_pickle=False) + per_query = [] + latencies = [] + for query_number, (query_id, query_path) in enumerate(queries): + index_key = f"q{query_number}_idx" + if index_key not in gold: + raise ValueError(f"gold archive is missing {index_key}") + gold_indices = gold[index_key].astype(np.int64, copy=False).tolist() + if any(index < 0 or index >= len(page_keys) for index in gold_indices): + raise ValueError(f"gold archive {index_key} contains an invalid page index") + expected = [page_keys[index] for index in gold_indices] + query = np.load(query_path, allow_pickle=False) + actual, latency_ms, query_profile = execute_query( + psql_command, + args.table, + args.index, + sql_halfvec_array(query), + args.endpoint, + args.probes, + args.refine, + args.candidate_limit, + args.timeout_ms, + args.profile, + ) + latencies.append(latency_ms) + per_query.append( + { + "query_id": query_id, + "query_rows": int(query.shape[0]), + "returned_candidates": len(actual), + "latency_ms": round(latency_ms, 3), + "recall_at_10": recall(expected, actual, 10), + "recall_at_20": recall(expected, actual, 20), + "candidate_recall_at_10": candidate_recall(expected, actual, 10), + "candidate_recall_at_20": candidate_recall(expected, actual, 20), + "top_10": actual[:10], + **({"profile": query_profile} if query_profile is not None else {}), + } + ) + + output = { + "benchmark": ( + "gbrain_vectorchord_phase3b_external_v1" + if args.index is not None + else "gbrain_vectorchord_phase3a_v1" + ), + "corpus_pages": len(page_keys), + "query_count": len(per_query), + "configuration": { + "table": args.table, + "index": args.index, + "probes": args.probes, + "maxsim_refine": args.refine, + "candidate_limit": args.candidate_limit, + }, + "metrics": { + "recall_at_10": sum(item["recall_at_10"] for item in per_query) + / len(per_query), + "recall_at_20": sum(item["recall_at_20"] for item in per_query) + / len(per_query), + "candidate_recall_at_10": sum( + item["candidate_recall_at_10"] for item in per_query + ) + / len(per_query), + "candidate_recall_at_20": sum( + item["candidate_recall_at_20"] for item in per_query + ) + / len(per_query), + }, + "latency_ms": { + "mean": sum(latencies) / len(latencies), + "p50": percentile(latencies, 0.50), + "p95": percentile(latencies, 0.95), + "p99": percentile(latencies, 0.99), + "max": max(latencies), + }, + "per_query": per_query, + } + profiles = [item["profile"] for item in per_query if "profile" in item] + if profiles: + profile_keys = sorted(set.intersection(*(set(profile) for profile in profiles))) + output["profile_summary"] = { + key: { + "mean": sum(profile[key] for profile in profiles) / len(profiles), + "p50": percentile([profile[key] for profile in profiles], 0.50), + "p95": percentile([profile[key] for profile in profiles], 0.95), + "max": max(profile[key] for profile in profiles), + } + for key in profile_keys + if key != "schema_version" + } + print(json.dumps(output, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_gbrain_scoped_tilemaxsim.py b/services/benchmark_gbrain_scoped_tilemaxsim.py new file mode 100644 index 00000000..96c4d646 --- /dev/null +++ b/services/benchmark_gbrain_scoped_tilemaxsim.py @@ -0,0 +1,252 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark real CUDA TileMaxSim over a GBrain-style lexical/structured scope.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import statistics +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_full_corpus_tilemaxsim import descriptor_batches +from services.benchmark_tilemaxsim_ablation import encode_frame, request_round_trip + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def load_jsonl(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + return [json.loads(line) for line in stream if line.strip()] + + +def select_method(report: dict[str, object], method_name: str) -> dict[str, object]: + methods = report.get("methods") + if not isinstance(methods, list): + raise ValueError("candidate report has no methods array") + for method in methods: + if isinstance(method, dict) and method.get("method") == method_name: + return method + raise ValueError(f"candidate method {method_name!r} was not found") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--page-manifest", required=True, type=Path) + parser.add_argument("--query-dataset", required=True, type=Path) + parser.add_argument("--corpus-json", required=True, type=Path) + parser.add_argument("--query-embeddings", required=True, type=Path) + parser.add_argument("--candidate-report", required=True, type=Path) + parser.add_argument("--candidate-method", default="es_bm25") + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--query-limit", type=int, default=0) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + raw_descriptors = load_jsonl(args.descriptor_manifest) + page_metadata = { + str(record["page_key"]): record for record in load_jsonl(args.page_manifest) + } + descriptors = [ + {**descriptor, **page_metadata.get(str(descriptor.get("page_key")), {})} + for descriptor in raw_descriptors + ] + queries = json.loads(args.query_dataset.read_text(encoding="utf-8")) + corpus_records = json.loads(args.corpus_json.read_text(encoding="utf-8")) + candidate_report = json.loads(args.candidate_report.read_text(encoding="utf-8")) + method = select_method(candidate_report, args.candidate_method) + per_query = method.get("per_query_results") + if not isinstance(queries, list) or not queries: + raise ValueError("query dataset must be a nonempty JSON array") + if not isinstance(corpus_records, list): + raise ValueError("corpus JSON must be an array") + if not isinstance(per_query, dict): + raise ValueError("candidate method has no per_query_results object") + if args.query_limit < 0: + parser.error("--query-limit must be nonnegative") + queries = queries[: args.query_limit or None] + + descriptors_by_doc: dict[str, list[dict[str, object]]] = {} + descriptors_by_doc_page: dict[tuple[str, int], list[dict[str, object]]] = {} + for descriptor in descriptors: + doc_name = descriptor.get("doc_name") + if isinstance(doc_name, str): + descriptors_by_doc.setdefault(doc_name, []).append(descriptor) + descriptors_by_doc_page.setdefault( + (doc_name, int(descriptor.get("page_no", 0))), [] + ).append(descriptor) + chunks_by_id = { + str(record["_id"]): record["_source"] for record in corpus_records + } + + cases: list[dict[str, object]] = [] + scope_latencies: list[float] = [] + cuda_latencies: list[float] = [] + end_to_end_latencies: list[float] = [] + for query_index, query_record in enumerate(queries): + query_id = query_record.get("query_id") + gold_doc = query_record.get("gold_doc_name") + if not isinstance(query_id, str) or not isinstance(gold_doc, str): + raise ValueError(f"query {query_index} lacks query_id or gold_doc_name") + raw_candidates = per_query.get(query_id, []) + if not isinstance(raw_candidates, list): + raise ValueError(f"candidate list for {query_id} is not an array") + + scope_started = time.perf_counter() + candidate_docs: list[str] = [] + seen_docs: set[str] = set() + for row in raw_candidates: + doc_name = row.get("doc_name") if isinstance(row, dict) else None + if isinstance(doc_name, str) and doc_name not in seen_docs: + seen_docs.add(doc_name) + candidate_docs.append(doc_name) + scoped_descriptors: list[dict[str, object]] = [] + seen_pages: set[str] = set() + for row in raw_candidates: + if not isinstance(row, dict): + continue + chunk_id = str(row.get("chunk_id", "")) + doc_name = row.get("doc_name") + source = chunks_by_id.get(chunk_id) + content = source.get("content_with_weight", "") if isinstance(source, dict) else "" + page_numbers = { + int(value) for value in re.findall(r"::(\d+)", str(content)) + } + mapped = [ + descriptor + for page_no in sorted(page_numbers) + for descriptor in descriptors_by_doc_page.get((str(doc_name), page_no), []) + ] + # Preserve recall when an imported chunk has no page markers. + if not mapped and isinstance(doc_name, str): + mapped = descriptors_by_doc.get(doc_name, []) + for descriptor in mapped: + page_key = str(descriptor.get("page_key", "")) + if page_key and page_key not in seen_pages: + seen_pages.add(page_key) + scoped_descriptors.append(descriptor) + scope_ms = (time.perf_counter() - scope_started) * 1000 + scope_latencies.append(scope_ms) + + query = np.load(args.query_embeddings / f"{query_id}.npy", allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" None: + self.model_contract_id = model_contract_id + self.records = self._load_records(source_manifest, descriptor_manifest) + self.cache = sidecar.PayloadCache(cache_bytes) + + @staticmethod + def _json_lines(path: Path) -> list[dict[str, object]]: + records = [] + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: expected a JSON object") + records.append(record) + if not records: + raise ValueError(f"{path}: manifest is empty") + return records + + @classmethod + def _load_records( + cls, source_manifest: Path, descriptor_manifest: Path + ) -> dict[str, TensorRecord]: + source_root = source_manifest.resolve(strict=True).parent + sources: dict[str, tuple[Path, int, int]] = {} + for record in cls._json_lines(source_manifest): + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if ( + not isinstance(page_key, str) + or not isinstance(relative, str) + or not isinstance(rows, int) + or rows <= 0 + or not isinstance(dimension, int) + or dimension <= 0 + ): + raise ValueError("source manifest contains an invalid tensor record") + path = (source_root / relative).resolve(strict=True) + try: + path.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes source root: {relative}" + ) from error + if page_key in sources: + raise ValueError(f"duplicate source page_key: {page_key}") + sources[page_key] = (path, rows, dimension) + + resolved: dict[str, TensorRecord] = {} + seen_pages = set() + for descriptor in cls._json_lines(descriptor_manifest): + page_key = descriptor.get("page_key") + tensor_ref = descriptor.get("tensor_ref") + rows = descriptor.get("tensor_rows") + dimension = descriptor.get("tensor_dim") + dtype_name = descriptor.get("tensor_dtype") + checksum = descriptor.get("tensor_checksum") + source = sources.get(page_key) if isinstance(page_key, str) else None + if source is None: + raise ValueError(f"descriptor has unknown page_key: {page_key!r}") + source_path, source_rows, source_dimension = source + if rows != source_rows or dimension != source_dimension: + raise ValueError(f"descriptor shape disagrees for page {page_key}") + dtype = { + "float16": protocol.DTYPE_F16, + "float32": protocol.DTYPE_F32, + }.get(dtype_name) + if dtype is None: + raise ValueError(f"descriptor has invalid dtype for page {page_key}") + if not isinstance(tensor_ref, str) or not tensor_ref.startswith( + "sha256://" + ): + raise ValueError( + f"descriptor has invalid tensor_ref for page {page_key}" + ) + digest = tensor_ref.removeprefix("sha256://") + if ( + len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + or checksum != f"sha256:{digest}" + ): + raise ValueError(f"descriptor has invalid SHA-256 for page {page_key}") + tensor_record = TensorRecord(source_path, rows, dimension, dtype, checksum) + previous = resolved.get(tensor_ref) + if previous is not None: + if ( + previous.rows != tensor_record.rows + or previous.dimension != tensor_record.dimension + or previous.dtype != tensor_record.dtype + or previous.checksum != tensor_record.checksum + ): + raise ValueError(f"conflicting duplicate tensor_ref: {tensor_ref}") + else: + resolved[tensor_ref] = tensor_record + seen_pages.add(page_key) + if seen_pages != set(sources): + raise ValueError("descriptor manifest does not cover the source manifest") + return resolved + + def resolve( + self, request: protocol.ExternalTensorRequest + ) -> sidecar.ResolvedPayload: + if request.model_contract_id != self.model_contract_id: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract is not allowlisted by this benchmark resolver", + ) + record = self.records.get(request.tensor_ref) + if record is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference is not present in the benchmark manifest", + ) + if ( + request.rows != record.rows + or request.dimension != record.dimension + or request.dtype != record.dtype + or not hmac.compare_digest(request.checksum, record.checksum) + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor descriptor disagrees with the benchmark manifest", + ) + key = ( + request.tensor_ref, + request.rows, + request.dimension, + request.dtype, + ) + cached = self.cache.get(key) + if cached is not None: + return sidecar.ResolvedPayload(cached, True) + + dtype_name = "float16" if record.dtype == protocol.DTYPE_F16 else "float32" + tensor = canonical_tensor(record.path, record.rows, record.dimension) + if tensor.dtype != np.dtype(dtype_name): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "NPY dtype disagrees with the benchmark descriptor", + ) + payload = memoryview(tensor).cast("B").tobytes() + actual_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(actual_checksum, record.checksum): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "NPY payload checksum disagrees with the benchmark descriptor", + ) + self.cache.put(key, payload) + return sidecar.ResolvedPayload(payload, False) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=sidecar.parse_mode, default=0o600) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--model-contract", required=True) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument( + "--max-request-bytes", type=sidecar.positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument( + "--max-batch-tokens", type=sidecar.positive_int, default=1_000_000 + ) + parser.add_argument( + "--max-tensor-bytes", type=sidecar.positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=sidecar.positive_int, default=65_536) + parser.add_argument( + "--max-device-bytes", + type=sidecar.positive_int, + default=8 * 1024 * 1024 * 1024, + ) + parser.add_argument( + "--cache-bytes", type=sidecar.nonnegative_int, default=2 * 1024 * 1024 * 1024 + ) + parser.add_argument( + "--request-timeout-ms", type=sidecar.positive_int, default=60_000 + ) + parser.add_argument("--max-inflight", type=sidecar.positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=sidecar.positive_int, default=1) + parser.add_argument("--backlog", type=sidecar.positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + if not args.model_contract: + parser.error("--model-contract must not be empty") + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = NpyManifestResolver( + args.model_contract, + args.manifest, + args.descriptor_manifest, + args.cache_bytes, + ) + metrics = sidecar.JsonMetrics() + engine = sidecar.TorchTileMaxsimEngine( + args.device, + args.max_device_bytes, + args.allow_tf32, + args.max_cuda_inflight, + ) + service = sidecar.TileMaxsimService( + limits, resolver, engine, args.request_timeout_ms, metrics + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + sidecar.serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + ) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_postgres_single_vector.py b/services/benchmark_postgres_single_vector.py new file mode 100644 index 00000000..3041190d --- /dev/null +++ b/services/benchmark_postgres_single_vector.py @@ -0,0 +1,336 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Benchmark PostgreSQL/pgvector HNSW with original text-embedding vectors.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import shlex +import statistics +import subprocess +import time +from pathlib import Path + +import numpy as np + +IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*(?:\.[A-Za-z_][A-Za-z0-9_$]*)?$") + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)] + + +def summary(samples: list[float]) -> dict[str, float | int]: + return { + "count": len(samples), + "mean": statistics.fmean(samples), + "p50": percentile(samples, 0.50), + "p95": percentile(samples, 0.95), + "p99": percentile(samples, 0.99), + "max": max(samples), + } + + +def vector_text(vector: np.ndarray) -> str: + return "[" + ",".join(format(float(value), ".8g") for value in vector) + "]" + + +def run_psql(command: list[str], sql: str) -> str: + result = subprocess.run( + [*command, "-X", "-q", "-A", "-t", "-v", "ON_ERROR_STOP=1"], + input=sql, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if result.returncode != 0: + diagnostic = result.stderr.strip() or result.stdout.strip() + raise RuntimeError(f"psql failed: {diagnostic}") + return result.stdout + + +def prepare_table( + command: list[str], table: str, vectors: np.ndarray, m: int, ef_construction: int +) -> float: + dimension = int(vectors.shape[1]) + process = subprocess.Popen( + [*command, "-X", "-q", "-v", "ON_ERROR_STOP=1"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + assert process.stdin is not None + started = time.perf_counter() + process.stdin.write( + f"DROP TABLE IF EXISTS {table};\n" + f"CREATE TABLE {table} (id integer PRIMARY KEY, embedding vector({dimension}) NOT NULL);\n" + f"COPY {table} (id, embedding) FROM STDIN;\n" + ) + for index, vector in enumerate(vectors): + process.stdin.write(f"{index}\t{vector_text(vector)}\n") + process.stdin.write( + "\\.\n" + f"CREATE INDEX ON {table} USING hnsw (embedding vector_ip_ops) " + f"WITH (m={m}, ef_construction={ef_construction});\n" + f"ANALYZE {table};\n" + ) + process.stdin.close() + stdout = process.stdout.read() if process.stdout is not None else "" + stderr = process.stderr.read() if process.stderr is not None else "" + return_code = process.wait() + if return_code != 0: + raise RuntimeError(f"psql prepare failed: {stderr.strip() or stdout.strip()}") + return (time.perf_counter() - started) * 1000 + + +def execute_queries( + command: list[str], table: str, queries: np.ndarray, ef_search: int, top_k: int +) -> list[tuple[int, float, list[int]]]: + dimension = int(queries.shape[1]) + query_values = ",\n".join( + f"({index}, '{vector_text(vector)}'::vector({dimension}))" + for index, vector in enumerate(queries) + ) + sql = f""" +SET hnsw.ef_search = {ef_search}; +CREATE TEMP TABLE single_vector_queries ( + query_index integer PRIMARY KEY, + embedding vector({dimension}) NOT NULL +) ON COMMIT PRESERVE ROWS; +INSERT INTO single_vector_queries VALUES +{query_values}; +CREATE TEMP TABLE single_vector_results ( + query_index integer PRIMARY KEY, + latency_ms double precision NOT NULL, + ids integer[] NOT NULL +) ON COMMIT PRESERVE ROWS; +DO $bench$ +DECLARE + q record; + started timestamptz; + result_ids integer[]; +BEGIN + -- Warm PostgreSQL buffers and the HNSW path before recording latency. + FOR q IN SELECT * FROM single_vector_queries ORDER BY query_index LOOP + EXECUTE 'SELECT array_agg(id ORDER BY distance, id) FROM ( + SELECT id, embedding <#> $1 AS distance + FROM {table} + ORDER BY embedding <#> $1 + LIMIT {top_k} + ) ranked' + INTO result_ids USING q.embedding; + END LOOP; + FOR q IN SELECT * FROM single_vector_queries ORDER BY query_index LOOP + started := clock_timestamp(); + EXECUTE 'SELECT array_agg(id ORDER BY distance, id) FROM ( + SELECT id, embedding <#> $1 AS distance + FROM {table} + ORDER BY embedding <#> $1 + LIMIT {top_k} + ) ranked' + INTO result_ids USING q.embedding; + INSERT INTO single_vector_results VALUES ( + q.query_index, + extract(epoch FROM clock_timestamp() - started) * 1000.0, + result_ids + ); + END LOOP; +END +$bench$; +SELECT query_index, latency_ms, array_to_string(ids, ',') + FROM single_vector_results + ORDER BY query_index; +""" + rows = [] + for line in run_psql(command, sql).splitlines(): + if not line.strip(): + continue + query_index, latency_ms, ids = line.split("|") + rows.append( + ( + int(query_index), + float(latency_ms), + [int(value) for value in ids.split(",")], + ) + ) + return rows + + +def recall(expected: np.ndarray, actual: list[int], top_k: int) -> float: + return len(set(expected[:top_k].tolist()).intersection(actual[:top_k])) / top_k + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--page-vectors", type=Path) + parser.add_argument("--query-vectors", type=Path) + parser.add_argument("--gold", type=Path) + parser.add_argument( + "--corpus-json", + type=Path, + help="RAG/GBrain corpus JSON whose _source.q_1024_vec is the stored text embedding", + ) + parser.add_argument( + "--query-dataset", + type=Path, + help="query JSON whose q_1024_vec is the original text query embedding", + ) + parser.add_argument("--psql-command", default="psql") + parser.add_argument("--table", default="public.tilemaxsim_single_vector_bench") + parser.add_argument("--prepare", action="store_true") + parser.add_argument("--m", type=int, default=16) + parser.add_argument("--ef-construction", type=int, default=64) + parser.add_argument("--ef-search", type=int, default=40) + parser.add_argument("--top-k", type=int, default=20) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + if not IDENTIFIER.fullmatch(args.table): + parser.error("--table must be an unquoted SQL identifier") + if min(args.m, args.ef_construction, args.ef_search, args.top_k) <= 0: + parser.error("HNSW parameters and top-k must be positive") + command = shlex.split(args.psql_command) + if not command: + parser.error("--psql-command must not be empty") + + corpus_chunk_ids: list[str] | None = None + corpus_doc_names: list[str] | None = None + query_records: list[dict[str, object]] | None = None + if args.corpus_json or args.query_dataset: + if not args.corpus_json or not args.query_dataset: + parser.error("--corpus-json and --query-dataset must be used together") + corpus_records = json.loads(args.corpus_json.read_text(encoding="utf-8")) + query_records = json.loads(args.query_dataset.read_text(encoding="utf-8")) + if not isinstance(corpus_records, list) or not isinstance(query_records, list): + raise ValueError("corpus and query JSON inputs must be arrays") + page_vectors = np.asarray( + [record["_source"]["q_1024_vec"] for record in corpus_records], + dtype=np.float32, + ) + query_vectors = np.asarray( + [record["q_1024_vec"] for record in query_records], + dtype=np.float32, + ) + corpus_chunk_ids = [str(record["_id"]) for record in corpus_records] + corpus_doc_names = [str(record["_source"]["docnm_kwd"]) for record in corpus_records] + else: + if not args.page_vectors or not args.query_vectors or not args.gold: + parser.error( + "use --corpus-json/--query-dataset or all of " + "--page-vectors/--query-vectors/--gold" + ) + page_vectors = np.load(args.page_vectors, mmap_mode="r", allow_pickle=False) + query_vectors = np.load(args.query_vectors, mmap_mode="r", allow_pickle=False) + if page_vectors.ndim != 2 or query_vectors.ndim != 2: + raise ValueError("page and query vectors must be rank-2 arrays") + if page_vectors.shape[1] != query_vectors.shape[1]: + raise ValueError("page and query dimensions differ") + prepare_ms = ( + prepare_table( + command, args.table, page_vectors, args.m, args.ef_construction + ) + if args.prepare + else None + ) + query_rows = execute_queries( + command, args.table, query_vectors, args.ef_search, args.top_k + ) + cases = [] + if query_records is not None and corpus_chunk_ids is not None and corpus_doc_names is not None: + for query_index, latency_ms, ids in query_rows: + query_record = query_records[query_index] + gold_chunk = str(query_record["gold_chunk_id"]) + gold_doc = str(query_record["gold_doc_name"]) + ranked_chunks = [corpus_chunk_ids[index] for index in ids] + ranked_docs = [corpus_doc_names[index] for index in ids] + cases.append( + { + "query_index": query_index, + "query_id": query_record.get("query_id"), + "latency_ms": latency_ms, + "chunk_hit_at_1": gold_chunk in ranked_chunks[:1], + "chunk_hit_at_10": gold_chunk in ranked_chunks[:10], + "chunk_hit_at_20": gold_chunk in ranked_chunks[:20], + "doc_hit_at_1": gold_doc in ranked_docs[:1], + "doc_hit_at_10": gold_doc in ranked_docs[:10], + "doc_hit_at_20": gold_doc in ranked_docs[:20], + "top_20_chunk_ids": ranked_chunks[:20], + "top_20_doc_names": ranked_docs[:20], + } + ) + else: + assert args.gold is not None + gold = np.load(args.gold, allow_pickle=False) + for query_index, latency_ms, ids in query_rows: + expected = gold[f"q{query_index}_idx"].astype(np.int64, copy=False) + cases.append( + { + "query_index": query_index, + "latency_ms": latency_ms, + "recall_at_10": recall(expected, ids, 10), + "recall_at_20": recall(expected, ids, 20), + "top_20": ids[:20], + } + ) + latencies = [item["latency_ms"] for item in cases] + report = { + "benchmark": "postgres_pgvector_hnsw_text_embedding_v2", + "corpus_vectors": int(page_vectors.shape[0]), + "query_vectors": int(query_vectors.shape[0]), + "dimension": int(page_vectors.shape[1]), + "configuration": { + "table": args.table, + "m": args.m, + "ef_construction": args.ef_construction, + "ef_search": args.ef_search, + "top_k": args.top_k, + }, + "prepare_ms": prepare_ms, + "latency_ms": summary(latencies), + "quality": ( + { + key: statistics.fmean(float(item[key]) for item in cases) + for key in ( + "chunk_hit_at_1", + "chunk_hit_at_10", + "chunk_hit_at_20", + "doc_hit_at_1", + "doc_hit_at_10", + "doc_hit_at_20", + ) + } + if query_records is not None + else { + "mean_recall_at_10": statistics.fmean( + item["recall_at_10"] for item in cases + ), + "mean_recall_at_20": statistics.fmean( + item["recall_at_20"] for item in cases + ), + } + ), + "cases": cases, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_tilemaxsim_ablation.py b/services/benchmark_tilemaxsim_ablation.py new file mode 100644 index 00000000..d7ab0ab5 --- /dev/null +++ b/services/benchmark_tilemaxsim_ablation.py @@ -0,0 +1,765 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Run storage, cache, H2D, and native-daemon TileMaxSim ablations.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import select +import socket +import statistics +import subprocess +import sys +import tempfile +import time +from collections import OrderedDict +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import decode_response +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver, PayloadCache +from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + GpuTensorLoad, +) + + +class _LegacyBuddyAllocator: + """Former power-of-two allocator retained only as an ablation baseline.""" + + def __init__(self, capacity: int, block_bytes: int = 256 * 1024) -> None: + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + self.capacity = self.block_count * block_bytes + self._free: dict[int, set[tuple[int, int, int]]] = {} + self._allocated: dict[int, tuple[int, int, int]] = {} + start = 0 + remaining = self.block_count + while remaining: + order = remaining.bit_length() - 1 + size = 1 << order + self._free.setdefault(order, set()).add((start, start, order)) + start += size + remaining -= size + + @property + def free_bytes(self) -> int: + return sum( + len(items) * (1 << order) * self.block_bytes + for order, items in self._free.items() + ) + + def allocation_bytes(self, payload_bytes: int) -> int: + raw = math.ceil(payload_bytes / self.block_bytes) + return (1 << (raw - 1).bit_length()) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.allocation_bytes(payload_bytes) // self.block_bytes + order = required.bit_length() - 1 + available_order = next( + ( + candidate + for candidate in range(order, self.block_count.bit_length()) + if self._free.get(candidate) + ), + None, + ) + if available_order is None: + return None + start, root_start, root_order = self._free[available_order].pop() + while available_order > order: + available_order -= 1 + buddy = start + (1 << available_order) + self._free.setdefault(available_order, set()).add( + (buddy, root_start, root_order) + ) + self._allocated[start] = (order, root_start, root_order) + return tuple(range(start, start + required)) + + def release(self, blocks: tuple[int, ...]) -> None: + start = blocks[0] + order, root_start, root_order = self._allocated.pop(start) + while order < root_order: + buddy = root_start + ((start - root_start) ^ (1 << order)) + item = (buddy, root_start, root_order) + free = self._free.setdefault(order, set()) + if item not in free: + break + free.remove(item) + start = min(start, buddy) + order += 1 + self._free.setdefault(order, set()).add((start, root_start, root_order)) + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + return ordered[min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)] + + +def load_records(path: Path) -> list[dict[str, object]]: + with path.open(encoding="utf-8") as stream: + return [json.loads(line) for line in stream if line.strip()] + + +def request(record: dict[str, object], contract: str) -> protocol.ExternalTensorRequest: + dtype = ( + protocol.DTYPE_F16 + if record["tensor_dtype"] == "float16" + else protocol.DTYPE_F32 + ) + return protocol.ExternalTensorRequest( + contract, + str(record["tensor_ref"]), + int(record["tensor_rows"]), + int(record["tensor_dim"]), + dtype, + str(record["tensor_checksum"]), + ) + + +def evict_paths(paths: list[Path]) -> None: + if not hasattr(os, "posix_fadvise"): + return + for path in paths: + descriptor = os.open(path, os.O_RDONLY | os.O_CLOEXEC) + try: + os.posix_fadvise(descriptor, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(descriptor) + + +def storage_ablation( + selected: list[dict[str, object]], + contract: str, + legacy_root: Path, + shard_root: Path, +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + legacy_paths = [ + legacy_root + / str(record["tensor_checksum"])[7:9] + / f"{str(record['tensor_checksum'])[7:]}.bin" + for record in selected + ] + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + sequential = [resolver.resolve(item) for item in requests] + finally: + resolver.close() + sequential_ms = (time.perf_counter() - started) * 1000 + + evict_paths(legacy_paths) + resolver = ContentAddressedResolver({contract: legacy_root}, 0) + started = time.perf_counter() + try: + legacy_batch = resolver.resolve_many(requests) + finally: + resolver.close() + legacy_batch_ms = (time.perf_counter() - started) * 1000 + + evict_paths(shard_paths) + resolver = ContentAddressedResolver({contract: shard_root}, 0) + started = time.perf_counter() + try: + shard_batch = resolver.resolve_many(requests) + shard_status = resolver.status() + finally: + resolver.close() + shard_batch_ms = (time.perf_counter() - started) * 1000 + expected = [item.payload for item in sequential] + if expected != [item.payload for item in legacy_batch] or expected != [ + item.payload for item in shard_batch + ]: + raise RuntimeError("storage ablation payloads disagree") + return { + "candidates": len(selected), + "logical_bytes": sum(len(item) for item in expected), + "legacy_sequential_ms": sequential_ms, + "legacy_batch_ms": legacy_batch_ms, + "shard_batch_ms": shard_batch_ms, + "shard_batch_read_calls": shard_status["batch_read_calls"], + "shard_batch_read_bytes": shard_status["batch_read_bytes"], + } + + +def h2d_ablation( + selected: list[dict[str, object]], contract: str, shard_root: Path, device: int +) -> dict[str, object]: + requests = [request(record, contract) for record in selected] + resolver = ContentAddressedResolver({contract: shard_root}, 0) + try: + payloads = resolver.resolve_many(requests) + keys = [resolver.key(item) for item in requests] + finally: + resolver.close() + total_bytes = 768 * 1024**2 + workspace_bytes = 256 * 1024**2 + + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + started = time.perf_counter() + handles = [] + for key, item, resolved in zip(keys, requests, payloads, strict=True): + handle, _ = cache.acquire( + key, + item.rows, + item.dimension, + item.dtype, + lambda payload=resolved.payload: payload, + ) + handles.append(handle) + torch.cuda.synchronize(device) + sequential_ms = (time.perf_counter() - started) * 1000 + sequential_status = pool.status()[0] + for handle in handles: + cache.release(handle) + finally: + pool.close() + + pool = GpuResourcePool( + [GpuArenaSpec(f"cuda:{device}", total_bytes)], workspace_bytes + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + loads = [ + GpuTensorLoad(key, item.rows, item.dimension, item.dtype, resolved.payload) + for key, item, resolved in zip(keys, requests, payloads, strict=True) + ] + started = time.perf_counter() + batch = cache.acquire_many(loads) + torch.cuda.synchronize(device) + batch_ms = (time.perf_counter() - started) * 1000 + batch_status = pool.status()[0] + if batch.bypassed or batch.deferred: + raise RuntimeError("H2D ablation cache is undersized") + for handle in batch.handles: + assert handle is not None + cache.release(handle) + finally: + pool.close() + return { + "candidates": len(selected), + "logical_bytes": sum(len(item.payload) for item in payloads), + "per_tensor_h2d_ms": sequential_ms, + "batch_h2d_ms": batch_ms, + "per_tensor_copy_batches": sequential_status["h2d_batches"], + "batch_copy_batches": batch_status["h2d_batches"], + "batch_copy_calls": batch_status["h2d_copy_calls"], + } + + +def allocator_ablation(records: list[dict[str, object]]) -> dict[str, object]: + sizes = [int(record["canonical_bytes"]) for record in records] + rng = random.Random(991) + capacity = 256 * 1024**2 + events: list[tuple[str, int, int]] = [] + abstract_live: list[int] = [] + next_identifier = 0 + for _ in range(20_000): + if abstract_live and rng.random() < 0.48: + index = rng.randrange(len(abstract_live)) + identifier = abstract_live.pop(index) + events.append(("release", identifier, 0)) + else: + size = rng.choice(sizes) + identifier = next_identifier + next_identifier += 1 + abstract_live.append(identifier) + events.append(("allocate", identifier, size)) + + def run_trace( + allocator: FreeExtentAllocator | FixedBlockAllocator | _LegacyBuddyAllocator, + ) -> dict[str, object]: + live: dict[int, tuple[int, ...] | tuple[int, int]] = {} + failures = 0 + fragmentation_failures = 0 + requested_bytes = 0 + allocated_bytes = 0 + started = time.perf_counter() + for operation, identifier, size in events: + if operation == "release": + allocation = live.pop(identifier, None) + if allocation is None: + continue + if isinstance(allocator, FreeExtentAllocator): + allocator.release(*allocation) + else: + allocator.release(allocation) + continue + required = allocator.allocation_bytes(size) + allocation = allocator.allocate(size) + if allocation is None: + failures += 1 + fragmentation_failures += int(allocator.free_bytes >= required) + else: + live[identifier] = allocation + requested_bytes += size + allocated_bytes += ( + allocation[1] + if isinstance(allocator, FreeExtentAllocator) + else len(allocation) * allocator.block_bytes + ) + return { + "allocation_failures": failures, + "fragmentation_failures": fragmentation_failures, + "internal_waste_ratio": ( + (allocated_bytes - requested_bytes) / allocated_bytes + ), + "trace_ms": (time.perf_counter() - started) * 1000, + } + + extent = FreeExtentAllocator(capacity) + legacy = _LegacyBuddyAllocator(capacity) + page_runs = FixedBlockAllocator(capacity) + legacy_full_bytes = sum(legacy.allocation_bytes(size) for size in sizes) + page_run_full_bytes = sum(page_runs.allocation_bytes(size) for size in sizes) + return { + "operations": 20_000, + "capacity_bytes": capacity, + "exact_byte_extents": run_trace(extent), + "legacy_power_of_two_buddy": { + "block_bytes": legacy.block_bytes, + "full_corpus_allocated_bytes": legacy_full_bytes, + **run_trace(legacy), + }, + "segregated_page_runs": { + "block_bytes": page_runs.block_bytes, + "full_corpus_allocated_bytes": page_run_full_bytes, + "full_corpus_space_saved_bytes": legacy_full_bytes - page_run_full_bytes, + **run_trace(page_runs), + }, + } + + +class LegacyLru: + def __init__(self, maximum_bytes: int) -> None: + self.maximum_bytes = maximum_bytes + self.bytes = 0 + self.entries: OrderedDict[str, bytes] = OrderedDict() + + def access(self, key: str, size: int) -> bool: + if key in self.entries: + self.entries.move_to_end(key) + return True + payload = bytes(size) + if size <= self.maximum_bytes: + self.entries[key] = payload + self.bytes += size + while self.bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.bytes -= len(evicted) + return False + + +def policy_ablation(records: list[dict[str, object]]) -> dict[str, object]: + rng = random.Random(77) + universe = records[:5000] + sizes = { + str(record["tensor_ref"]): int(record["canonical_bytes"]) for record in universe + } + hot = list(sizes)[:100] + cold = list(sizes)[100:] + # Warm every hot object before injecting scans so TinyLFU admission is + # measured directly rather than starting from an empty cache. + trace = hot * 5 + for cycle in range(30): + trace.extend(rng.choices(hot, weights=range(len(hot), 0, -1), k=1000)) + start = cycle * 100 % len(cold) + trace.extend(cold[start : start + 300]) + # The budget holds exactly the hot set but not the scan tail. + budget = sum(sizes[key] for key in hot) + lru = LegacyLru(budget) + gdsf = PayloadCache(budget) + lru_hits = 0 + gdsf_hits = 0 + for key in trace: + size = sizes[key] + lru_hits += int(lru.access(key, size)) + cached = gdsf.get((key,)) + gdsf_hits += int(cached is not None) + if cached is None: + gdsf.put((key,), bytes(size)) + return { + "accesses": len(trace), + "budget_bytes": budget, + "lru_hit_ratio": lru_hits / len(trace), + "tinylfu_gdsf_hit_ratio": gdsf_hits / len(trace), + "tinylfu_gdsf_status": gdsf.status(), + } + + +def encode_frame( + records: list[dict[str, object]], contract: str, query: np.ndarray, request_id: int +) -> bytes: + dtype = protocol.DTYPE_F16 if query.dtype == np.dtype(" tuple[float, list[tuple[int, float]]]: + started = time.perf_counter() + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + elapsed = (time.perf_counter() - started) * 1000 + _, status, parsed = decode_response(response) + if status != 0 or not isinstance(parsed, list): + raise RuntimeError(f"daemon request failed: {parsed}") + if len(parsed) != len(protocol.parse_request_frame(frame).candidates): + raise RuntimeError("daemon returned an incomplete result") + return elapsed, parsed + + +def daemon_ablation( + records: list[dict[str, object]], + contract: str, + shard_root: Path, + rust_binary: Path, + device: int, + repeats: int, + gpu_block_kib: int = 32, +) -> dict[str, object]: + rng = np.random.default_rng(31) + query = rng.standard_normal((44, int(records[0]["tensor_dim"]))).astype(" dict[str, object]: + shard_paths = sorted((shard_root / "shards").glob("*.vts")) + results = {} + with tempfile.TemporaryDirectory() as directory: + commands = { + "python_triton": [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + f"{directory}/python-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--gpu-block-kib", + str(gpu_block_kib), + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + "--request-timeout-ms", + "20000", + ], + "rust_cuda": [ + os.fspath(rust_binary), + "--socket", + f"{directory}/rust-resident.sock", + "--gpu-memory-gb", + f"{device}=20", + "--gpu-workspace-gb", + "2", + "--gpu-block-kib", + str(gpu_block_kib), + "--host-cache-gb", + "0.1", + "--contract-root", + f"{contract}={shard_root}", + "--gpu-cache-mode", + "resident", + "--resident-manifest", + f"{contract}={descriptor_manifest}", + "--prewarm-batch-size", + "256", + ], + } + for name, command in commands.items(): + os.sync() + evict_paths(shard_paths) + started = time.perf_counter() + process = subprocess.Popen( + command, + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + lines = [] + prewarm_event = None + ready_event = None + deadline = time.monotonic() + 300 + try: + assert process.stdout is not None + while time.monotonic() < deadline: + readable, _, _ = select.select([process.stdout], [], [], 1.0) + if readable: + line = process.stdout.readline() + if line: + lines.append(line) + if line.startswith("{"): + event = json.loads(line) + if event.get("event") in ( + "tilemaxsim_prewarm_complete", + "tilemaxsim_rust_prewarm_complete", + ): + prewarm_event = event + if event.get("event") in ( + "tilemaxsim_ready", + "tilemaxsim_rust_ready", + ): + ready_event = event + break + if process.poll() is not None: + break + if ready_event is None: + remainder, _ = process.communicate(timeout=5) + raise RuntimeError( + f"{name} resident prewarm failed: {''.join(lines)}{remainder}" + ) + results[name] = { + "process_to_ready_ms": (time.perf_counter() - started) * 1000, + "prewarm_reported_ms": prewarm_event.get("elapsed_ms") + if prewarm_event + else None, + "cache": ready_event.get("gpu_cache", ready_event.get("cache")), + } + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=15) + process.communicate(timeout=5) + return results + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--legacy-root", required=True, type=Path) + parser.add_argument("--shard-root", required=True, type=Path) + parser.add_argument("--rust-binary", required=True, type=Path) + parser.add_argument("--contract", default="benchmark@1") + parser.add_argument("--device", required=True, type=int) + parser.add_argument("--sample-size", type=int, default=100) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--gpu-block-kib", type=int, default=32) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--full-prewarm", action="store_true") + args = parser.parse_args() + records = load_records(args.descriptor_manifest) + selected = random.Random(20260714).sample(records, args.sample_size) + report = { + "corpus": { + "records": len(records), + "logical_bytes": sum(int(record["canonical_bytes"]) for record in records), + "sample_size": len(selected), + }, + "storage": storage_ablation( + selected, args.contract, args.legacy_root, args.shard_root + ), + "h2d": h2d_ablation(selected, args.contract, args.shard_root, args.device), + "allocator": allocator_ablation(records), + "policy": policy_ablation(records), + "daemon": daemon_ablation( + selected, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + args.repeats, + args.gpu_block_kib, + ), + } + if args.full_prewarm: + report["full_resident_prewarm"] = prewarm_ablation( + args.descriptor_manifest, + args.contract, + args.shard_root, + args.rust_binary, + args.device, + args.gpu_block_kib, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as stream: + json.dump(report, stream, indent=2, sort_keys=True) + stream.write("\n") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py new file mode 100644 index 00000000..31799d26 --- /dev/null +++ b/services/benchmark_tilemaxsim_cuda.py @@ -0,0 +1,148 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time + +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import TorchTileMaxsimEngine, positive_int + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def canonical_payload(tensor: torch.Tensor, dtype: int) -> bytes: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return tensor.to(dtype=scalar_dtype).contiguous().numpy().tobytes() + + +def normalized_tensor( + shape: tuple[int, ...], generator: torch.Generator +) -> torch.Tensor: + tensor = torch.randn(shape, dtype=torch.float32, generator=generator) + return torch.nn.functional.normalize(tensor, dim=-1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", choices=("f16", "f32"), default="f16") + parser.add_argument("--dimension", type=positive_int, default=320) + parser.add_argument("--query-rows", type=positive_int, default=32) + parser.add_argument("--document-rows", type=positive_int, default=747) + parser.add_argument("--candidates", type=positive_int, default=128) + parser.add_argument("--warmup", type=positive_int, default=3) + parser.add_argument("--iterations", type=positive_int, default=10) + parser.add_argument("--seed", type=int, default=20260713) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + dtype = protocol.DTYPE_F32 if args.dtype == "f32" else protocol.DTYPE_F16 + generator = torch.Generator(device="cpu").manual_seed(args.seed) + query = normalized_tensor((args.query_rows, args.dimension), generator) + document_tensors = normalized_tensor( + (args.candidates, args.document_rows, args.dimension), generator + ) + query_payload = canonical_payload(query, dtype) + documents = [ + ( + candidate_id, + args.document_rows, + canonical_payload(document_tensors[candidate_id], dtype), + ) + for candidate_id in range(args.candidates) + ] + del document_tensors + + engine = TorchTileMaxsimEngine( + args.device, args.max_device_bytes, args.allow_tf32, 1 + ) + if engine.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(engine.device) + + total_samples: list[float] = [] + queue_samples: list[float] = [] + compute_samples: list[float] = [] + score_checksum = 0.0 + for iteration in range(args.warmup + args.iterations): + started = time.perf_counter() + results, queue_ms, compute_ms = engine.score( + query_payload, + args.query_rows, + args.dimension, + dtype, + documents, + time.monotonic() + 300, + lambda: False, + ) + total_ms = (time.perf_counter() - started) * 1000.0 + if iteration >= args.warmup: + total_samples.append(total_ms) + queue_samples.append(queue_ms) + compute_samples.append(compute_ms) + score_checksum = math.fsum(score for _, score in results) + + output = { + "benchmark": "tilemaxsim_cuda_synthetic_v1", + "device": str(engine.device), + "device_name": ( + torch.cuda.get_device_name(engine.device) + if engine.device.type == "cuda" + else "cpu" + ), + "torch_version": torch.__version__, + "dtype": args.dtype, + "dimension": args.dimension, + "query_rows": args.query_rows, + "document_rows": args.document_rows, + "candidates": args.candidates, + "candidate_tokens": args.candidates * args.document_rows, + "seed": args.seed, + "warmup": args.warmup, + "iterations": args.iterations, + "allow_tf32": args.allow_tf32, + "max_device_bytes": args.max_device_bytes, + "latency_ms": { + "mean": round(statistics.fmean(total_samples), 3), + "p50": round(percentile(total_samples, 0.50), 3), + "p95": round(percentile(total_samples, 0.95), 3), + "p99": round(percentile(total_samples, 0.99), 3), + "queue_mean": round(statistics.fmean(queue_samples), 3), + "compute_mean": round(statistics.fmean(compute_samples), 3), + }, + "score_checksum": score_checksum, + } + if engine.device.type == "cuda": + output["cuda_peak_allocated_bytes"] = torch.cuda.max_memory_allocated( + engine.device + ) + output["cuda_peak_reserved_bytes"] = torch.cuda.max_memory_reserved( + engine.device + ) + print(json.dumps(output, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/benchmark_tilemaxsim_io_pipeline.py b/services/benchmark_tilemaxsim_io_pipeline.py new file mode 100644 index 00000000..df2eeaac --- /dev/null +++ b/services/benchmark_tilemaxsim_io_pipeline.py @@ -0,0 +1,241 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Compare serial and overlapped tensor I/O on the native CUDA daemon.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import tempfile +import time +from pathlib import Path + +import numpy as np + +from services.benchmark_tilemaxsim_ablation import ( + encode_frame, + evict_paths, + request_round_trip, +) + + +def load_records(path: Path, maximum_candidates: int) -> list[dict[str, object]]: + records: list[dict[str, object]] = [] + tokens = 0 + logical_bytes = 0 + with path.open(encoding="utf-8") as stream: + for line in stream: + if not line.strip(): + continue + record = json.loads(line) + rows = int(record["tensor_rows"]) + payload_bytes = int(record["canonical_bytes"]) + if records and ( + len(records) >= maximum_candidates + or tokens + rows > 1_000_000 + or logical_bytes + payload_bytes > 1024**3 + ): + break + records.append(record) + tokens += rows + logical_bytes += payload_bytes + if len(records) < 2: + raise ValueError("the benchmark needs at least two protocol-valid tensors") + return records + + +def run_mode( + *, + mode: str, + trial: int, + binary: Path, + shard_root: Path, + contract: str, + device: int, + gpu_memory_gb: str, + workspace_gb: str, + host_cache_gb: str, + io_batch_gb: str, + frame: bytes, + shard_paths: list[Path], + directory: Path, +) -> tuple[float, list[tuple[int, float]], dict[str, object]]: + evict_paths(shard_paths) + socket_path = directory / f"{mode}-{trial}.sock" + command = [ + os.fspath(binary), + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}={gpu_memory_gb}", + "--gpu-workspace-gb", + workspace_gb, + "--host-cache-gb", + host_cache_gb, + "--io-pipeline", + mode, + "--io-batch-gb", + io_batch_gb, + "--contract-root", + f"{contract}={shard_root}", + "--request-timeout-ms", + "60000", + "--socket-io-timeout-ms", + "60000", + "--once", + ] + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(3000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + if process.poll() is not None or not socket_path.exists(): + output, _ = process.communicate(timeout=5) + raise RuntimeError(f"{mode} daemon failed to start: {output}") + latency_ms, scores = request_round_trip(socket_path, frame) + output, _ = process.communicate(timeout=70) + if process.returncode != 0: + raise RuntimeError(f"{mode} daemon failed: {output}") + events = [ + json.loads(line) for line in output.splitlines() if line.startswith("{") + ] + request_event = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + return latency_ms, scores, request_event["cache"]["io_pipeline"] + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--query-embedding", required=True, type=Path) + parser.add_argument("--shard-root", required=True, type=Path) + parser.add_argument("--rust-binary", required=True, type=Path) + parser.add_argument("--contract", required=True) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--candidates", type=int, default=1000) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--gpu-memory-gb", default="1") + parser.add_argument("--workspace-gb", default="0.2") + parser.add_argument("--host-cache-gb", default="0.1") + parser.add_argument("--io-batch-gb", default="0.05") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + if args.candidates < 2 or args.trials < 1: + parser.error("candidates must be >= 2 and trials must be positive") + + records = load_records(args.descriptor_manifest, args.candidates) + query = np.load(args.query_embedding, allow_pickle=False) + if query.dtype != np.dtype("float16"): + query = query.astype(" np.ndarray: + tensor = np.load(path, mmap_mode="r", allow_pickle=False) + if tensor.ndim != 2 or tensor.shape != (expected_rows, expected_dim): + raise ValueError( + f"{path}: expected shape {(expected_rows, expected_dim)}, got {tensor.shape}" + ) + if tensor.dtype == np.dtype("float16"): + little_dtype = np.dtype(" str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + if path.stat().st_size != len(payload): + raise ValueError(f"existing cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"existing cache payload has wrong checksum: {path}") + return + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + try: + with os.fdopen(descriptor, "wb", closefd=True) as stream: + stream.write(payload) + stream.flush() + if fsync: + os.fsync(stream.fileno()) + try: + os.link(temporary_name, path) + except FileExistsError: + if path.stat().st_size != len(payload): + raise ValueError(f"concurrent cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"concurrent cache payload has wrong checksum: {path}") + if fsync: + directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + +def prepare_record( + record: dict[str, object], + source_root: Path, +) -> tuple[dict[str, object], memoryview]: + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if not isinstance(page_key, str) or not page_key: + raise ValueError("manifest record has no page_key") + if not isinstance(relative, str) or not relative: + raise ValueError(f"manifest page {page_key} has no embedding_file") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"manifest page {page_key} has invalid n_tokens") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"manifest page {page_key} has invalid dim") + source = (source_root / relative).resolve(strict=True) + try: + source.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes the source root: {relative}" + ) from error + tensor = canonical_tensor(source, rows, dimension) + payload = memoryview(tensor).cast("B") + digest = hashlib.sha256(payload).hexdigest() + dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" + return { + "page_key": page_key, + "tensor_ref": f"sha256://{digest}", + "tensor_rows": rows, + "tensor_dim": dimension, + "tensor_dtype": dtype_name, + "tensor_checksum": f"sha256:{digest}", + "canonical_bytes": len(payload), + }, payload + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + """Publish one legacy per-tensor file. + + Kept for backwards compatibility and migration tests. New cache builds use + immutable shards by default. + """ + + descriptor, payload = prepare_record(record, source_root) + digest = str(descriptor["tensor_checksum"]).removeprefix("sha256:") + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + return descriptor + + +def shard_size_gb(value: str) -> int: + try: + parsed = int(Decimal(value) * 1024**3) + except (InvalidOperation, ValueError) as error: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") from error + if parsed <= 0: + raise argparse.ArgumentTypeError("shard size must be a positive number of GB") + return parsed + + +def bounded_parallel_map(executor, function, items, maximum_pending: int): + """Preserve input order without retaining a corpus-sized future backlog.""" + + iterator = iter(items) + pending = deque() + for _ in range(maximum_pending): + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + break + while pending: + future = pending.popleft() + yield future.result() + try: + pending.append(executor.submit(function, next(iterator))) + except StopIteration: + pass + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--cache-root", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument( + "--storage-format", + choices=("shards", "files"), + default="shards", + help="publish immutable shards (default) or legacy per-tensor files", + ) + parser.add_argument( + "--shard-size-gb", + type=shard_size_gb, + default=DEFAULT_SHARD_BYTES, + help="target immutable shard size in GB", + ) + parser.add_argument("--no-fsync", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + source_root = args.manifest.resolve(strict=True).parent + cache_root = args.cache_root.resolve() + if not cache_root.is_absolute(): + parser.error("--cache-root must be absolute") + records = [] + with args.manifest.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError( + f"invalid JSON at manifest line {line_number}" + ) from error + if not isinstance(record, dict): + raise ValueError(f"manifest line {line_number} is not an object") + records.append(record) + if not records: + raise ValueError("manifest is empty") + + cache_root.mkdir(parents=True, exist_ok=True) + descriptors = [] + shard_writer = None + if args.storage_format == "shards" and not args.dry_run: + shard_writer = ImmutableShardWriter( + cache_root, + target_bytes=args.shard_size_gb, + fsync=not args.no_fsync, + ) + try: + with ThreadPoolExecutor(max_workers=args.workers) as workers: + if args.storage_format == "files": + results = ( + (item, None) + for item in bounded_parallel_map( + workers, + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + args.workers * 2, + ) + ) + else: + results = bounded_parallel_map( + workers, + lambda record: prepare_record(record, source_root), + records, + args.workers * 2, + ) + for completed, (item, payload) in enumerate(results, 1): + descriptors.append(item) + if shard_writer is not None: + assert payload is not None + digest = str(item["tensor_checksum"]).removeprefix("sha256:") + shard_writer.add( + digest, + payload, + int(item["tensor_rows"]), + int(item["tensor_dim"]), + str(item["tensor_dtype"]), + ) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + shard_index = shard_writer.finish() if shard_writer is not None else None + finally: + if shard_writer is not None: + shard_writer.close() + + output_parent = args.descriptor_manifest.parent + output_parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{args.descriptor_manifest.name}.", + suffix=".tmp", + dir=output_parent, + text=True, + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + for item in descriptors: + stream.write(json.dumps(item, separators=(",", ":"), sort_keys=True)) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, args.descriptor_manifest) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + total_bytes = sum(int(item["canonical_bytes"]) for item in descriptors) + print( + json.dumps( + { + "pages": len(descriptors), + "canonical_bytes": total_bytes, + "cache_root": os.fspath(cache_root), + "descriptor_manifest": os.fspath(args.descriptor_manifest), + "storage_format": args.storage_format, + "shard_index": os.fspath(shard_index) if shard_index else None, + "dry_run": args.dry_run, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..fb4b88e7 --- /dev/null +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -0,0 +1,586 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import hashlib +import json +import os +import socket +import stat +import struct +import subprocess +import sys +import tempfile +import threading +import time +import unittest +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import ( + decode_response, + external_request_frame, + request_frame, +) +from services import tilemaxsim_cuda_sidecar as cuda_sidecar +from services.build_tilemaxsim_tensor_cache import process_record + + +class CapturingMetrics(cuda_sidecar.JsonMetrics): + def __init__(self) -> None: + super().__init__() + self.events: list[dict[str, object]] = [] + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + self.events.append(fields.copy()) + + +def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: + digest = hashlib.sha256(payload).hexdigest() + directory = root / digest[:2] + directory.mkdir(parents=True, exist_ok=True) + (directory / f"{digest}.bin").write_bytes(payload) + return f"sha256://{digest}", f"sha256:{digest}" + + +class CudaSidecarTest(unittest.TestCase): + def test_cli_without_explicit_gpu_memory_keeps_tilemaxsim_disabled(self) -> None: + with tempfile.TemporaryDirectory() as directory: + socket_path = Path(directory) / "disabled.sock" + completed = subprocess.run( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + ], + cwd=Path(__file__).resolve().parents[1], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(completed.returncode, 2) + self.assertIn("TileMaxSim is disabled", completed.stderr) + self.assertFalse(socket_path.exists()) + + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_root = root / "source" + cache_root = root / "cache" + source_root.mkdir() + tensor = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + cuda_sidecar.validate_finite_payload( + struct.pack("<2e", 1.0, 0.0), 1, 2, protocol.DTYPE_F16 + ) + with self.assertRaisesRegex(protocol.SidecarError, "non-finite"): + cuda_sidecar.validate_finite_payload( + struct.pack("<2f", 1.0, float("nan")), + 1, + 2, + protocol.DTYPE_F32, + ) + + def test_content_addressed_resolver_validates_and_caches(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + try: + request = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + checksum, + ) + first = resolver.resolve(request) + second = resolver.resolve(request) + self.assertEqual(first.payload, payload) + self.assertFalse(first.cache_hit) + self.assertTrue(second.cache_hit) + + bad = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + "sha256:" + "0" * 64, + ) + with self.assertRaisesRegex(protocol.SidecarError, "disagree"): + resolver.resolve(bad) + finally: + resolver.close() + + def test_host_payload_cache_evicts_to_its_byte_budget(self) -> None: + cache = cuda_sidecar.PayloadCache(6) + cache.put(("first",), b"1234") + cache.put(("second",), b"5678") + self.assertIsNone(cache.get(("first",))) + self.assertEqual(cache.get(("second",)), b"5678") + self.assertEqual(cache.current_bytes, 4) + + def test_content_addressed_resolver_rejects_symlink(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) / "root" + root.mkdir() + payload = struct.pack(" None: + query = [[1.0, 0.0], [0.0, 1.0]] + candidates = [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5], [0.25, 0.25]]), + ] + frame = request_frame(41, protocol.DTYPE_F32, query, candidates) + parsed = protocol.parse_request_frame(frame) + self.assertIsInstance(parsed, protocol.InlineTensorRequest) + assert isinstance(parsed, protocol.InlineTensorRequest) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ] + # 64 bytes fits one candidate but not both, exercising internal + # all-or-nothing device chunking. + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 64, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + documents, + time.monotonic() + 2, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + self.assertEqual(results, oracle) + + def test_compute_capacity_wait_uses_overall_deadline(self) -> None: + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024, False, 1) + self.assertTrue(engine.compute_slots.acquire(blocking=False)) + try: + started = time.monotonic() + with self.assertRaisesRegex(protocol.SidecarError, "CUDA capacity"): + engine.score( + struct.pack(" None: + query = [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]] + candidates = [ + (7, [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]]), + (2, [[0.5, 0.5, 0.0], [-0.5, 0.25, 1.0]]), + ] + frame = request_frame(52, protocol.DTYPE_F16, query, candidates) + parsed = protocol.parse_request_frame(frame) + assert isinstance(parsed, protocol.InlineTensorRequest) + engine = cuda_sidecar.TorchTileMaxsimEngine("cuda:0", 1024 * 1024, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ], + time.monotonic() + 5, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + assert isinstance(oracle, list) + self.assertEqual([item[0] for item in results], [item[0] for item in oracle]) + for (_, actual), (_, expected) in zip(results, oracle, strict=True): + self.assertAlmostEqual(actual, expected, places=5) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_cli_gb_allocation_serves_tilemaxsim(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 70, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + socket_path = directory_path / "resident.sock" + process = subprocess.Popen( + [ + sys.executable, + "-m", + "services.tilemaxsim_cuda_sidecar", + "--socket", + os.fspath(socket_path), + "--gpu-memory-gb", + f"{device}=0.05", + "--gpu-workspace-gb", + "0.02", + "--host-cache-gb", + "0.01", + "--contract-root", + f"model@1={root}", + "--once", + ], + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + self.assertTrue(socket_path.exists()) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + self.assertEqual(decode_response(response)[1:], (0, [(9, 2.0)])) + self.assertIn('"event":"tilemaxsim_ready"', output) + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_v2_gpu_resident_hit_does_not_resolve_payload_again(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + digest = checksum.removeprefix("sha256:") + frame, _ = external_request_frame( + 71, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, [[1.0, 0.0], [0.0, 1.0]])], + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + stream_engine = cuda_sidecar.TorchTileMaxsimEngine( + f"cuda:{device}", 16 * 1024 * 1024, False, 1 + ) + resident_engine = cuda_sidecar.ResidentTorchTileMaxsimEngine( + pool, 16 * 1024 * 1024, False, 1 + ) + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + stream_engine, + 2000, + metrics, + cache, + resident_engine, + pin_gpu_entries=True, + ) + client, server = socket.socketpair() + try: + first = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + (root / digest[:2] / f"{digest}.bin").unlink() + second = service.process_frame( + frame, server, time.monotonic() + 2, None + ) + finally: + client.close() + server.close() + self.assertEqual(decode_response(first)[1:], (0, [(9, 2.0)])) + self.assertEqual(decode_response(second)[1:], (0, [(9, 2.0)])) + requests = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(requests[0]["gpu_cache_misses"], 1) + self.assertEqual(requests[1]["gpu_cache_hits"], 1) + finally: + resolver.close() + pool.close() + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_lru_gpu_cache_streams_request_larger_than_its_arena(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + rows, dimension = 480, 320 + first_tensor = np.zeros((rows, dimension), dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + manifest = root / "descriptors.jsonl" + manifest.write_text( + json.dumps( + { + "page_key": "page-1", + "tensor_ref": tensor_ref, + "tensor_rows": 2, + "tensor_dim": 2, + "tensor_dtype": "float16", + "tensor_checksum": checksum, + "canonical_bytes": len(payload), + } + ) + + "\n", + encoding="utf-8", + ) + device = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + pool = cuda_sidecar.GpuResourcePool( + [cuda_sidecar.GpuArenaSpec(f"cuda:{device}", 32 * 1024 * 1024)], + 16 * 1024 * 1024, + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 0) + try: + cache = cuda_sidecar.GpuTensorCache(pool, allow_eviction=False) + metrics = CapturingMetrics() + cuda_sidecar.prewarm_resident_cache( + [("model@1", manifest)], resolver, cache, metrics + ) + status = cache.status() + self.assertEqual(status["entries"], 1) + self.assertEqual(status["pinned_entries"], 1) + self.assertEqual( + metrics.events[-1]["event"], "tilemaxsim_prewarm_complete" + ) + finally: + resolver.close() + pool.close() + + def test_v2_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + tensor = [[1.0, 0.0], [0.0, 1.0]] + payload = struct.pack("<4e", *sum(tensor, [])) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 61, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, tensor)], + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + metrics = CapturingMetrics() + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024 * 1024, False, 1), + 2000, + metrics, + ) + socket_path = directory_path / "tilemaxsim.sock" + stop = threading.Event() + thread = threading.Thread( + target=cuda_sidecar.serve, + args=(socket_path, 0o600, 4, 2, service, stop), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if socket_path.exists(): + break + time.sleep(0.01) + else: + self.fail("CUDA sidecar socket was not created") + self.assertEqual(stat.S_IMODE(socket_path.stat().st_mode), 0o600) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + thread.join(timeout=3) + resolver.close() + self.assertFalse(thread.is_alive()) + _, status, results = decode_response(response) + self.assertEqual(status, 0) + self.assertEqual(results, [(9, 2.0)]) + request_events = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(len(request_events), 1) + self.assertEqual(request_events[0]["source"], "content_addressed") + self.assertEqual(request_events[0]["status"], "ok") + + +if __name__ == "__main__": + unittest.main() diff --git a/services/test_tilemaxsim_gpu_cache.py b/services/test_tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..2fd82cc4 --- /dev/null +++ b/services/test_tilemaxsim_gpu_cache.py @@ -0,0 +1,372 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import time +import unittest + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ResidentTorchTileMaxsimEngine +from services.tilemaxsim_gpu_cache import ( + FixedBlockAllocator, + FreeExtentAllocator, + GpuArenaSpec, + GpuResourcePool, + GpuTensorCache, + GpuTensorLoad, + parse_gpu_memory_gb, + parse_memory_gb, +) + + +def available_device() -> str: + index = max( + range(torch.cuda.device_count()), + key=lambda candidate: torch.cuda.mem_get_info(candidate)[0], + ) + return f"cuda:{index}" + + +class GpuCacheUnitTest(unittest.TestCase): + def test_public_memory_configuration_uses_gb(self) -> None: + self.assertEqual(parse_memory_gb("20"), 20 * 1024**3) + self.assertEqual(parse_memory_gb("0.5"), 512 * 1024**2) + self.assertEqual( + parse_gpu_memory_gb("2=12"), + GpuArenaSpec("cuda:2", 12 * 1024**3), + ) + self.assertEqual( + parse_gpu_memory_gb("cuda:2=12.5"), + GpuArenaSpec("cuda:2", int(12.5 * 1024**3)), + ) + with self.assertRaisesRegex(ValueError, "GPU=GB"): + parse_gpu_memory_gb("cuda:0") + with self.assertRaisesRegex(ValueError, "byte suffixes"): + parse_gpu_memory_gb("0=20GiB") + + def test_extent_allocator_coalesces_released_ranges(self) -> None: + allocator = FreeExtentAllocator(4096, alignment=256) + first = allocator.allocate(300) + second = allocator.allocate(700) + self.assertEqual(first, (0, 512)) + self.assertEqual(second, (512, 768)) + assert first is not None and second is not None + allocator.release(*first) + allocator.release(*second) + self.assertEqual(allocator.extents, [(0, 4096)]) + + def test_fixed_block_allocator_uses_exact_runs_and_coalesces(self) -> None: + allocator = FixedBlockAllocator(8 * 256, block_bytes=256) + first = allocator.allocate(300) + second = allocator.allocate(300) + third = allocator.allocate(700) + self.assertEqual(first, (0, 1)) + self.assertEqual(second, (2, 3)) + self.assertEqual(third, (4, 5, 6)) + self.assertEqual(allocator.free_bytes, 256) + assert first is not None and second is not None and third is not None + allocator.release(second) + allocator.release(first) + self.assertEqual(allocator.largest_free_extent, 4 * 256) + allocator.release(third) + self.assertEqual(allocator.largest_free_extent, 8 * 256) + + def test_fixed_block_allocator_reuses_best_fit_run(self) -> None: + allocator = FixedBlockAllocator(12 * 256, block_bytes=256) + first = allocator.allocate(2 * 256) + separator = allocator.allocate(256) + second = allocator.allocate(3 * 256) + tail = allocator.allocate(6 * 256) + assert first is not None and separator is not None + assert second is not None and tail is not None + allocator.release(first) + allocator.release(second) + reused = allocator.allocate(3 * 256) + self.assertEqual(reused, second) + self.assertEqual(allocator.largest_free_extent, 2 * 256) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_pool_rejects_budget_larger_than_currently_free_memory(self) -> None: + device = available_device() + free_bytes, _ = torch.cuda.mem_get_info(torch.device(device)) + with self.assertRaisesRegex(RuntimeError, "cannot acquire"): + GpuResourcePool( + [GpuArenaSpec(device, free_bytes + 1024 * 1024)], + 1024 * 1024, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_gpu_cache_evicts_only_released_entries(self) -> None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 16 * 1024 * 1024)], 8 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + rows, dimension = 8192, 320 + payload = np.ones((rows, dimension), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 32 * 1024 * 1024)], 16 * 1024 * 1024 + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + query = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = available_device() + pool = GpuResourcePool([GpuArenaSpec(device, 4 * 1024 * 1024)], 2 * 1024 * 1024) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + first = np.ones((128, 320), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 768 * 1024)], + 512 * 1024, + block_bytes=256 * 1024, + ) + try: + cache = GpuTensorCache(pool, allow_eviction=True) + tensor = np.ones((128, 320), dtype=" None: + device = available_device() + pool = GpuResourcePool( + [GpuArenaSpec(device, 64 * 1024 * 1024)], 32 * 1024 * 1024 + ) + try: + generator = np.random.default_rng(7) + query = generator.standard_normal((44, 320)).astype("= 2, + "two CUDA devices are unavailable", + ) + def test_resident_engine_scores_shards_on_multiple_gpus(self) -> None: + pool = GpuResourcePool( + [ + GpuArenaSpec("cuda:0", 32 * 1024 * 1024), + GpuArenaSpec("cuda:1", 32 * 1024 * 1024), + ], + 16 * 1024 * 1024, + ) + try: + cache = GpuTensorCache(pool, allow_eviction=False) + identity = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" Path: + return Path(__file__).parent / "tilemaxsimd" / "target" / "release" / "tilemaxsimd" + + def run_daemon( + self, + devices: list[int], + documents: list[np.ndarray] | None = None, + query: np.ndarray | None = None, + gpu_memory_gb: str = "0.05", + workspace_gb: str = "0.02", + resident: bool = False, + scheduled: bool = False, + tcp: bool = False, + io_pipeline: str = "overlap", + io_batch_gb: str = "0.25", + ) -> tuple[str, list[tuple[int, float]]]: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + if documents is None: + documents = [ + np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + socket_path = root / "tilemaxsimd.sock" + completed = subprocess.run( + [ + os.fspath(binary), + "--socket", + os.fspath(socket_path), + "--contract-root", + f"model@1={root}", + ], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + self.assertNotEqual(completed.returncode, 0) + self.assertIn("--gpu-memory-gb", completed.stderr) + self.assertFalse(socket_path.exists()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_unavailable_configured_gpu_fails_before_socket_ready(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + writer = ImmutableShardWriter( + root, target_bytes=4096, alignment=256, fsync=False + ) + payload = np.asarray([[1.0, 0.0]], dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + self.run_daemon([device]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_io_overlap_matches_serial_scores_and_reports_pipeline_cycles(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + serial_output, serial_results = self.run_daemon( + [device], io_pipeline="serial", io_batch_gb="0.000000004" + ) + overlap_output, overlap_results = self.run_daemon( + [device], io_pipeline="overlap", io_batch_gb="0.000000004" + ) + self.assertEqual(overlap_results, serial_results) + serial_events = [ + json.loads(line) for line in serial_output.splitlines() if line.startswith("{") + ] + overlap_events = [ + json.loads(line) + for line in overlap_output.splitlines() + if line.startswith("{") + ] + serial_cache = next( + event["cache"] + for event in serial_events + if event.get("event") == "tilemaxsim_rust_request" + ) + overlap_cache = next( + event["cache"] + for event in overlap_events + if event.get("event") == "tilemaxsim_rust_request" + ) + self.assertEqual(serial_cache["io_pipeline"]["mode"], "serial") + self.assertEqual(overlap_cache["io_pipeline"]["mode"], "overlap") + self.assertEqual(serial_cache["io_pipeline"]["overlap_cycles"], 0) + self.assertGreaterEqual(overlap_cache["io_pipeline"]["overlap_cycles"], 1) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_tcp_scoring_round_trip_matches_protocol_oracle(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, results = self.run_daemon([device], tcp=True, scheduled=True) + self.assertIn('"listen":"127.0.0.1:', output) + self.assertEqual([candidate for candidate, _ in results], [11, 12]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_exact_scores_are_repeatable_across_cuda_contexts(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + generator = np.random.default_rng(20260715) + document = generator.normal(size=(79, 64)).astype(" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + document = np.asarray([[1.0, 0.0]], dtype=" None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + object_root = root / "objects-root" + object_root.mkdir() + socket_path = root / "tilemaxsimd.sock" + process = subprocess.Popen( + [ + os.fspath(binary), + "--socket", os.fspath(socket_path), + "--gpu-memory-gb", f"{device}=0.05", + "--gpu-workspace-gb", "0.02", + "--host-cache-gb", "0.01", + "--contract-root", f"model@1={object_root}", + "--once", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + document = np.asarray( + [[1.0, 0.0], [0.0, 1.0]], dtype=" None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, _ = self.run_daemon([device], scheduled=True) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + request = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + self.assertNotIn("tenant", request) + self.assertRegex(request["tenant_hash"], r"^[0-9a-f]{16}$") + self.assertEqual(request["priority"], 17) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.device_count() >= 2, + "two CUDA devices are unavailable", + ) + def test_multi_gpu_scheduler_uploads_and_scores_on_each_device(self) -> None: + output, _ = self.run_daemon([0, 1]) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + request_event = next( + event for event in events if event.get("event") == "tilemaxsim_rust_request" + ) + devices = request_event["cache"]["devices"] + self.assertEqual(len(devices), 2) + self.assertEqual([device["h2d_batches"] for device in devices], [1, 1]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_one_request_larger_than_gpu_cache_is_scored_in_chunks(self) -> None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + rows, dimension = 480, 320 + first = np.zeros((rows, dimension), dtype=" None: + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + output, _ = self.run_daemon([device], resident=True) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + prewarm = next( + event + for event in events + if event.get("event") == "tilemaxsim_rust_prewarm_complete" + ) + self.assertEqual(prewarm["entries"], 2) + self.assertEqual(prewarm["cache"]["devices"][0]["gpu_pinned_entries"], 2) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_concurrent_readers_do_not_block_fair_priority_scheduler(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" int: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.settimeout(10) + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + request_id, status, _ = decode_response(response) + self.assertEqual(status, 0) + return request_id + + with ThreadPoolExecutor(max_workers=len(frames)) as executor: + completed = list(executor.map(call, frames)) + self.assertEqual(set(completed), set(range(1_000, 1_000 + len(frames)))) + slow.close() + slow = None + process.terminate() + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + events = [ + json.loads(line) + for line in output.splitlines() + if line.startswith("{") + ] + processed = [ + event["priority"] + for event in events + if event.get("event") == "tilemaxsim_rust_request" + ] + # All public priorities share the default fair-priority band. + # Urgency breaks the first tie, then equal-cost tenants + # alternate instead of one tenant draining its whole queue. + self.assertEqual(processed, [9, 3, 8, 1, 7, 0, 5, -2]) + finally: + if slow is not None: + slow.close() + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_overload_rejects_one_tenant_without_crashing_daemon(self) -> None: + binary = self._release_binary() + if not binary.exists(): + self.skipTest("release tilemaxsimd binary has not been built") + device = max( + range(torch.cuda.device_count()), + key=lambda index: torch.cuda.mem_get_info(index)[0], + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shard_root = root / "cache" + shard_root.mkdir() + document = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" tuple[int, int]: + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_bytes = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_bytes) + request_id, status, _ = decode_response(response) + return request_id, status + + first = None + try: + for _ in range(1000): + if socket_path.exists() or process.poll() is not None: + break + time.sleep(0.01) + self.assertIsNone(process.poll()) + first = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + first.settimeout(10) + first.connect(os.fspath(socket_path)) + first.sendall(frames[0]) + time.sleep(0.05) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as second: + second.settimeout(10) + second.connect(os.fspath(socket_path)) + second.sendall(frames[1]) + self.assertEqual(receive(second), (2_002, 2)) + self.assertEqual(receive(first), (2_001, 0)) + self.assertIsNone(process.poll()) + finally: + if first is not None: + first.close() + if process.poll() is None: + process.terminate() + output, _ = process.communicate(timeout=10) + self.assertEqual(process.returncode, 0, output) + + +if __name__ == "__main__": + unittest.main() diff --git a/services/test_tilemaxsim_shard.py b/services/test_tilemaxsim_shard.py new file mode 100644 index 00000000..591e4954 --- /dev/null +++ b/services/test_tilemaxsim_shard.py @@ -0,0 +1,181 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +from __future__ import annotations + +import hashlib +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import ContentAddressedResolver +from services.tilemaxsim_shard import ImmutableShardWriter, load_index + + +class ImmutableShardTest(unittest.TestCase): + def test_builder_defaults_to_shards_and_publishes_atomically(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + source.mkdir() + manifest = source / "pages.jsonl" + records = [] + for index in range(3): + tensor = np.full((2, 4), index + 1, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensors = [ + np.arange(32, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensor = np.eye(4, dtype=" None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = np.ones((2, 4), dtype=" None: + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor byte length does not match its shape", + ) + scalar_dtype = " None: + self.maximum_bytes = maximum_bytes + self.current_bytes = 0 + self.entries: OrderedDict[tuple[object, ...], tuple[bytes, float]] = ( + OrderedDict() + ) + self.lock = threading.Lock() + self.sketch = TinyLfuSketch() + self.inflation = 0.0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.admission_rejections = 0 + + def get(self, key: tuple[object, ...]) -> bytes | None: + if self.maximum_bytes == 0: + return None + with self.lock: + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is not None: + payload, _ = entry + priority = self.inflation + frequency / len(payload) + self.entries[key] = (payload, priority) + self.entries.move_to_end(key) + self.hits += 1 + return payload + self.misses += 1 + return None + + def put(self, key: tuple[object, ...], payload: bytes) -> None: + if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: + return + with self.lock: + previous = self.entries.pop(key, None) + if previous is not None: + self.current_bytes -= len(previous[0]) + frequency = max(1, self.sketch.estimate(key)) + candidate_priority = self.inflation + frequency / len(payload) + while self.current_bytes + len(payload) > self.maximum_bytes: + victim_key, (victim_payload, victim_priority) = min( + self.entries.items(), key=lambda item: item[1][1] + ) + if candidate_priority < victim_priority: + if previous is not None: + self.entries[key] = previous + self.current_bytes += len(previous[0]) + self.admission_rejections += 1 + return + self.entries.pop(victim_key) + self.current_bytes -= len(victim_payload) + self.inflation = max(self.inflation, victim_priority) + candidate_priority = self.inflation + frequency / len(payload) + self.evictions += 1 + self.entries[key] = (payload, candidate_priority) + self.current_bytes += len(payload) + + def status(self) -> dict[str, object]: + with self.lock: + return { + "entries": len(self.entries), + "bytes": self.current_bytes, + "maximum_bytes": self.maximum_bytes, + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + } + + +class ContentAddressedResolver: + """Resolve ``sha256://`` inside an allowlisted model cache root. + + A payload with digest ``abcdef...`` is stored as + ``/ab/abcdef....bin``. Directory and file symlinks are + rejected with ``openat(O_NOFOLLOW)``. The digest in the reference, the + registered checksum, the exact byte length, and the file content must all + agree before a payload is returned. + """ + + def __init__( + self, + roots: dict[str, Path], + cache_bytes: int, + verify_full_shards: bool = False, + ) -> None: + self.root_fds: dict[str, int] = {} + self.shard_roots: dict[str, _OpenShardRoot] = {} + self.batch_read_calls = 0 + self.batch_read_bytes = 0 + self.batch_lock = threading.Lock() + self.verify_full_shards = verify_full_shards + try: + for contract, path in roots.items(): + root_fd = os.open( + os.fspath(path), + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + ) + self.root_fds[contract] = root_fd + self._open_shard_root(contract, root_fd) + except Exception: + self.close() + raise + self.cache = PayloadCache(cache_bytes) + + def _open_shard_root(self, contract: str, root_fd: int) -> None: + try: + index_fd = os.open( + INDEX_NAME, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + except FileNotFoundError: + return + try: + with os.fdopen(index_fd, "r", encoding="utf-8", closefd=True) as stream: + index = parse_index(json.load(stream)) + except Exception: + raise + shard_directory_fd = -1 + shard_fds: dict[str, int] = {} + try: + shard_directory_fd = os.open( + "shards", + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + for relative, shard in index.shards.items(): + name = relative.removeprefix("shards/") + descriptor = os.open( + name, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=shard_directory_fd, + ) + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size != shard.size: + os.close(descriptor) + raise ValueError(f"immutable shard metadata disagrees: {relative}") + shard_fds[relative] = descriptor + self.shard_roots[contract] = _OpenShardRoot( + index, + shard_directory_fd, + shard_fds, + {name: threading.Lock() for name in shard_fds}, + set(), + ) + except Exception: + for descriptor in shard_fds.values(): + os.close(descriptor) + if shard_directory_fd >= 0: + os.close(shard_directory_fd) + raise + + def close(self) -> None: + for root in self.shard_roots.values(): + for descriptor in root.shard_fds.values(): + os.close(descriptor) + os.close(root.directory_fd) + self.shard_roots.clear() + for descriptor in self.root_fds.values(): + os.close(descriptor) + self.root_fds.clear() + + @staticmethod + def _digest(request: protocol.ExternalTensorRequest) -> str: + prefix = "sha256://" + if not request.tensor_ref.startswith(prefix): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "unsupported tensor reference; expected sha256://", + ) + digest = request.tensor_ref[len(prefix) :] + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "invalid content-addressed tensor reference", + ) + if not hmac.compare_digest(request.checksum, f"sha256:{digest}"): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference and checksum disagree", + ) + return digest + + @staticmethod + def _read_exact_file(root_fd: int, digest: str, expected_bytes: int) -> bytes: + directory_fd = -1 + payload_fd = -1 + try: + directory_fd = os.open( + digest[:2], + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + payload_fd = os.open( + f"{digest}.bin", + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=directory_fd, + ) + metadata = os.fstat(payload_fd) + if not stat.S_ISREG(metadata.st_mode): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor is not a regular file", + ) + if metadata.st_size != expected_bytes: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + chunks = bytearray() + while len(chunks) < expected_bytes: + chunk = os.read( + payload_fd, min(1024 * 1024, expected_bytes - len(chunks)) + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file ended early", + ) + chunks.extend(chunk) + if os.read(payload_fd, 1): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file grew during read", + ) + return bytes(chunks) + except FileNotFoundError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "content-addressed tensor is missing" + ) from error + except OSError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + f"content-addressed tensor read failed: {error.strerror}", + ) from error + finally: + if payload_fd >= 0: + os.close(payload_fd) + if directory_fd >= 0: + os.close(directory_fd) + + def key(self, request: protocol.ExternalTensorRequest) -> tuple[object, ...]: + root_fd = self.root_fds.get(request.model_contract_id) + if root_fd is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract has no configured tensor cache root", + ) + digest = self._digest(request) + key = ( + request.model_contract_id, + digest, + request.rows, + request.dimension, + request.dtype, + ) + return key + + @staticmethod + def _validate_shard_entry( + request: protocol.ExternalTensorRequest, digest: str, entry: ShardEntry + ) -> None: + dtype_name = "float32" if request.dtype == protocol.DTYPE_F32 else "float16" + if ( + entry.digest != digest + or entry.rows != request.rows + or entry.dimension != request.dimension + or entry.dtype != dtype_name + or entry.length + != protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "shard entry disagrees with the tensor descriptor", + ) + + @staticmethod + def _verify_shard(root: _OpenShardRoot, name: str) -> None: + if name in root.verified: + return + lock = root.verification_locks[name] + with lock: + if name in root.verified: + return + shard = root.index.shards[name] + expected = shard.checksum.removeprefix("sha256:") + digest = hashlib.sha256() + offset = 0 + descriptor = root.shard_fds[name] + while offset < shard.size: + chunk = os.pread( + descriptor, min(8 * 1024 * 1024, shard.size - offset), offset + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "immutable tensor shard ended early", + ) + digest.update(chunk) + offset += len(chunk) + if not hmac.compare_digest(digest.hexdigest(), expected): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "immutable tensor shard checksum mismatch", + ) + root.verified.add(name) + + @staticmethod + def _coalesced_ranges( + entries: list[tuple[tuple[object, ...], ShardEntry]], + maximum_gap: int = 64 * 1024, + maximum_span: int = 8 * 1024 * 1024, + ) -> list[tuple[int, int, list[tuple[tuple[object, ...], ShardEntry]]]]: + ranges = [] + current: list[tuple[tuple[object, ...], ShardEntry]] = [] + start = 0 + end = 0 + for item in sorted(entries, key=lambda value: value[1].offset): + entry = item[1] + entry_end = entry.offset + entry.length + if current and ( + entry.offset - end > maximum_gap or entry_end - start > maximum_span + ): + ranges.append((start, end, current)) + current = [] + if not current: + start = entry.offset + end = entry_end + else: + end = max(end, entry_end) + current.append(item) + if current: + ranges.append((start, end, current)) + return ranges + + def _read_shard_range( + self, + root: _OpenShardRoot, + shard_name: str, + start: int, + end: int, + entries: list[tuple[tuple[object, ...], ShardEntry]], + ) -> dict[tuple[object, ...], bytes]: + payload = os.pread(root.shard_fds[shard_name], end - start, start) + if len(payload) != end - start: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "immutable tensor shard ended early" + ) + with self.batch_lock: + self.batch_read_calls += 1 + self.batch_read_bytes += len(payload) + result = {} + for key, entry in entries: + tensor = payload[entry.offset - start : entry.offset - start + entry.length] + actual = hashlib.sha256(tensor).hexdigest() + if not hmac.compare_digest(actual, entry.digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved shard tensor checksum mismatch", + ) + dtype = ( + protocol.DTYPE_F32 if entry.dtype == "float32" else protocol.DTYPE_F16 + ) + validate_finite_payload(tensor, entry.rows, entry.dimension, dtype) + result[key] = tensor + return result + + def resolve_many( + self, requests: list[protocol.ExternalTensorRequest] + ) -> list[ResolvedPayload]: + if not requests: + return [] + keys = [self.key(request) for request in requests] + payloads: dict[tuple[object, ...], bytes] = {} + hits: dict[tuple[object, ...], bool] = {} + missing: dict[tuple[object, ...], protocol.ExternalTensorRequest] = {} + for key, request in zip(keys, requests, strict=True): + cached = self.cache.get(key) + if cached is not None: + payloads[key] = cached + hits[key] = True + elif key not in missing: + missing[key] = request + hits[key] = False + + shard_groups: dict[ + tuple[str, str], list[tuple[tuple[object, ...], ShardEntry]] + ] = {} + legacy: list[tuple[tuple[object, ...], protocol.ExternalTensorRequest]] = [] + for key, request in missing.items(): + contract = request.model_contract_id + shard_root = self.shard_roots.get(contract) + if shard_root is None: + legacy.append((key, request)) + continue + digest = str(key[1]) + entry = shard_root.index.entries.get(digest) + if entry is None: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "content-addressed tensor is missing from the shard index", + ) + self._validate_shard_entry(request, digest, entry) + shard_groups.setdefault((contract, entry.shard), []).append((key, entry)) + + jobs: list[ + tuple[ + _OpenShardRoot, + str, + int, + int, + list[tuple[tuple[object, ...], ShardEntry]], + ] + ] = [] + for (contract, shard_name), entries in shard_groups.items(): + root = self.shard_roots[contract] + if self.verify_full_shards: + self._verify_shard(root, shard_name) + for start, end, grouped in self._coalesced_ranges(entries): + jobs.append((root, shard_name, start, end, grouped)) + if jobs: + with ThreadPoolExecutor(max_workers=min(8, len(jobs))) as workers: + for resolved in workers.map( + lambda job: self._read_shard_range(*job), jobs + ): + payloads.update(resolved) + + def read_legacy( + item: tuple[tuple[object, ...], protocol.ExternalTensorRequest], + ) -> tuple[tuple[object, ...], bytes]: + key, request = item + digest = str(key[1]) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + payload = self._read_exact_file( + self.root_fds[request.model_contract_id], digest, expected_bytes + ) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor checksum mismatch", + ) + validate_finite_payload( + payload, request.rows, request.dimension, request.dtype + ) + return key, payload + + if legacy: + with ThreadPoolExecutor(max_workers=min(8, len(legacy))) as workers: + for key, payload in workers.map(read_legacy, legacy): + payloads[key] = payload + for key in missing: + self.cache.put(key, payloads[key]) + return [ResolvedPayload(payloads[key], hits[key]) for key in keys] + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + return self.resolve_many([request])[0] + + def status(self) -> dict[str, object]: + with self.batch_lock: + return { + "shard_contracts": len(self.shard_roots), + "verified_shards": sum( + len(root.verified) for root in self.shard_roots.values() + ), + "batch_read_calls": self.batch_read_calls, + "batch_read_bytes": self.batch_read_bytes, + } + + +class TorchTileMaxsimEngine: + def __init__( + self, + device_name: str, + max_device_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.device = torch.device(device_name) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.type not in ("cuda", "cpu"): + raise RuntimeError("device must be CUDA or CPU") + self.max_device_bytes = max_device_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + if self.device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + left = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + _ = left @ left + torch.cuda.synchronize(self.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + # bytearray gives torch a writable, owned buffer; clone detaches the + # resulting tensor before that temporary buffer leaves scope. + tensor = torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + if scalar_dtype == torch.float32: + return tensor.clone() + return tensor.to(dtype=torch.float32) + + def _groups( + self, + query_rows: int, + dimension: int, + documents: list[tuple[int, int, bytes]], + ) -> Iterable[list[tuple[int, int, bytes]]]: + query_bytes = query_rows * dimension * 4 + group: list[tuple[int, int, bytes]] = [] + group_rows = 0 + for document in documents: + rows = document[1] + next_rows = group_rows + rows + # Device residency includes the f32 query, f32 documents, and the + # q-by-total-document-token similarity matrix. + required = ( + query_bytes + next_rows * dimension * 4 + query_rows * next_rows * 4 + ) + if required > self.max_device_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = query_bytes + rows * dimension * 4 + query_rows * rows * 4 + if required > self.max_device_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one candidate exceeds the CUDA device-byte limit", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, int, bytes]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + results: list[tuple[int, float]] = [] + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + try: + with torch.inference_mode(): + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(self.device) + for group in self._groups(query_rows, dimension, documents): + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request peer disconnected" + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + cpu_documents = [ + self._cpu_tensor(payload, rows, dimension, dtype) + for _, rows, payload in group + ] + document_device = torch.cat(cpu_documents).to(self.device) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, rows, _ in group: + scores.append( + similarities[:, offset : offset + rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += rows + host_scores = torch.stack(scores).to(device="cpu").tolist() + for (candidate_id, _, _), score in zip( + group, host_scores, strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + finally: + self.compute_slots.release() + return ( + results, + queue_ms, + (time.monotonic() - compute_started) * 1000.0, + ) + + +class ResidentTorchTileMaxsimEngine: + """Score tensors already owned by one or more process GPU arenas.""" + + def __init__( + self, + pool: GpuResourcePool, + max_workspace_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.pool = pool + self.device = pool.primary_device + self.max_workspace_bytes = max_workspace_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + for arena in pool.arenas: + left = torch.zeros((1, 1), dtype=torch.float32, device=arena.device) + _ = left @ left + for arena in pool.arenas: + torch.cuda.synchronize(arena.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + + def _groups( + self, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + ) -> Iterable[list[tuple[int, GpuTensorHandle]]]: + scalar_bytes = 4 if dtype == protocol.DTYPE_F32 else 2 + query_bytes = query_rows * dimension * scalar_bytes + group: list[tuple[int, GpuTensorHandle]] = [] + group_rows = 0 + for document in documents: + rows = document[1].rows + next_rows = group_rows + rows + # The resident document remains inside the arena. torch.cat makes + # one device-local contiguous scoring view; the other temporaries + # are the query and q-by-document-token similarity matrix. + required = ( + query_bytes + + next_rows * dimension * scalar_bytes + + query_rows * next_rows * scalar_bytes + ) + if required > self.max_workspace_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = ( + query_bytes + + rows * dimension * scalar_bytes + + query_rows * rows * scalar_bytes + ) + if required > self.max_workspace_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one resident candidate exceeds the configured GPU workspace", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, GpuTensorHandle]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + by_device: dict[str, list[tuple[int, GpuTensorHandle]]] = {} + for document in documents: + handle = document[1] + if handle.dimension != dimension or handle.dtype != dtype: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resident tensor contract disagrees with the query", + ) + by_device.setdefault(str(handle.device), []).append(document) + + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + pending: list[tuple[list[tuple[int, GpuTensorHandle]], torch.Tensor]] = [] + try: + with torch.inference_mode(): + for arena in self.pool.arenas: + device_documents = by_device.get(str(arena.device), []) + if not device_documents: + continue + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request peer disconnected", + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(arena.device) + if dtype == protocol.DTYPE_F16 and ragged_tilemaxsim_fp16: + offsets = torch.tensor( + [ + handle.offset_bytes // 2 + for _, handle in device_documents + ], + dtype=torch.int64, + device=arena.device, + ) + rows = torch.tensor( + [handle.rows for _, handle in device_documents], + dtype=torch.int32, + device=arena.device, + ) + assert arena.storage is not None + device_scores = ragged_tilemaxsim_fp16( + query_device, + arena.storage.view(torch.float16), + offsets, + rows, + max(handle.rows for _, handle in device_documents), + ) + pending.append((device_documents, device_scores)) + continue + for group in self._groups( + query_rows, dimension, dtype, device_documents + ): + document_device = torch.cat( + [handle.tensor() for _, handle in group] + ) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, handle in group: + scores.append( + similarities[:, offset : offset + handle.rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += handle.rows + pending.append((group, torch.stack(scores))) + + results: list[tuple[int, float]] = [] + for group, device_scores in pending: + host_scores = device_scores.to(device="cpu", dtype=torch.float32) + for (candidate_id, _), score in zip( + group, host_scores.tolist(), strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + for arena in self.pool.arenas: + torch.cuda.synchronize(arena.device) + finally: + self.compute_slots.release() + return results, queue_ms, (time.monotonic() - compute_started) * 1000.0 + + +class JsonMetrics: + def __init__(self) -> None: + self.lock = threading.Lock() + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + print(json.dumps(fields, separators=(",", ":"), sort_keys=True), flush=True) + + +class TileMaxsimService: + def __init__( + self, + limits: protocol.Limits, + resolver: ContentAddressedResolver, + engine: TorchTileMaxsimEngine, + request_timeout_ms: int, + metrics: JsonMetrics, + gpu_cache: GpuTensorCache | None = None, + resident_engine: ResidentTorchTileMaxsimEngine | None = None, + pin_gpu_entries: bool = False, + ) -> None: + self.limits = limits + self.resolver = resolver + self.engine = engine + self.request_timeout_seconds = request_timeout_ms / 1000.0 + self.metrics = metrics + self.gpu_cache = gpu_cache + self.resident_engine = resident_engine + self.pin_gpu_entries = pin_gpu_entries + if (gpu_cache is None) != (resident_engine is None): + raise ValueError( + "GPU cache and resident engine must be configured together" + ) + + @staticmethod + def _peer_disconnected(connection: socket.socket) -> bool: + poller = select.poll() + poller.register(connection, select.POLLHUP | select.POLLERR | select.POLLNVAL) + return bool(poller.poll(0)) + + @staticmethod + def _receive_exact_until( + connection: socket.socket, count: int, deadline: float + ) -> bytes: + chunks = bytearray() + while len(chunks) < count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired during socket read") + connection.settimeout(remaining) + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "connection closed during request", + ) + chunks.extend(chunk) + return bytes(chunks) + + def process_frame( + self, + frame: bytes, + connection: socket.socket, + deadline: float, + peer_credentials: tuple[int, int, int] | None, + ) -> bytes: + request_id = 0 + version = protocol.VERSION + started = time.monotonic() + metrics: dict[str, object] = {"event": "tilemaxsim_request"} + resident_documents: list[tuple[int, GpuTensorHandle]] = [] + if peer_credentials is not None: + metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( + peer_credentials + ) + try: + if len(frame) >= protocol.HEADER.size: + _, wire_version, _, request_id, _ = protocol.HEADER.unpack_from(frame) + if wire_version in protocol.SUPPORTED_VERSIONS: + version = wire_version + request = protocol.parse_request_frame( + frame, self.limits, validate_finite=False + ) + metrics.update( + request_id=request.request_id, + protocol_version=version, + query_rows=request.query_rows, + dimension=request.dimension, + candidate_count=len(request.candidates), + ) + validate_finite_payload( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + ) + resolve_started = time.monotonic() + cache_hits = 0 + gpu_cache_hits = 0 + gpu_cache_misses = 0 + gpu_chunks = 0 + resident_results: list[tuple[int, float]] = [] + resident_queue_ms = 0.0 + resident_compute_ms = 0.0 + document_tokens = 0 + documents: list[tuple[int, int, bytes]] = [] + + def flush_resident_documents() -> None: + nonlocal gpu_chunks, resident_queue_ms, resident_compute_ms + if not resident_documents: + return + assert self.resident_engine is not None + try: + batch_results, batch_queue_ms, batch_compute_ms = ( + self.resident_engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + resident_documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + resident_results.extend(batch_results) + resident_queue_ms += batch_queue_ms + resident_compute_ms += batch_compute_ms + gpu_chunks += 1 + finally: + assert self.gpu_cache is not None + for _, resident_handle in resident_documents: + self.gpu_cache.release(resident_handle) + resident_documents.clear() + + if isinstance(request, protocol.InlineTensorRequest): + for candidate in request.candidates: + validate_finite_payload( + candidate.payload, + candidate.rows, + request.dimension, + request.dtype, + ) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in request.candidates + ] + document_tokens = sum( + candidate.rows for candidate in request.candidates + ) + metrics["source"] = "inline" + else: + metrics["source"] = "content_addressed" + document_tokens = sum( + candidate.descriptor.rows for candidate in request.candidates + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + if self.gpu_cache is None: + descriptors = [ + candidate.descriptor for candidate in request.candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many(descriptors) + else: + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in descriptors + ] + for candidate, resolved in zip( + request.candidates, resolved_payloads, strict=True + ): + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) + ) + else: + keys = [ + self.resolver.key(candidate.descriptor) + for candidate in request.candidates + ] + probed, miss_indices = self.gpu_cache.probe_many(keys) + for candidate, handle in zip( + request.candidates, probed, strict=True + ): + if handle is not None: + resident_documents.append((candidate.candidate_id, handle)) + gpu_cache_hits = len(request.candidates) - len(miss_indices) + gpu_cache_misses = len(miss_indices) + missing_candidates = [ + request.candidates[index] for index in miss_indices + ] + missing_descriptors = [ + candidate.descriptor for candidate in missing_candidates + ] + if hasattr(self.resolver, "resolve_many"): + resolved_payloads = self.resolver.resolve_many( + missing_descriptors + ) + else: + resolved_payloads = [ + self.resolver.resolve(descriptor) + for descriptor in missing_descriptors + ] + cache_hits += sum(item.cache_hit for item in resolved_payloads) + pending = [ + ( + candidate, + GpuTensorLoad( + key, + candidate.descriptor.rows, + candidate.descriptor.dimension, + candidate.descriptor.dtype, + resolved.payload, + self.pin_gpu_entries, + ), + False, + ) + for candidate, key, resolved in zip( + missing_candidates, + (keys[index] for index in miss_indices), + resolved_payloads, + strict=True, + ) + ] + admission_rejections = 0 + while pending: + batch = self.gpu_cache.acquire_many( + [item[1] for item in pending], + enforce_admission=( + not self.pin_gpu_entries + and not any(item[2] for item in pending) + ), + record_access=False, + count_stats=False, + ) + for (candidate, load, _force), handle in zip( + pending, batch.handles, strict=True + ): + if handle is not None: + resident_documents.append( + (candidate.candidate_id, handle) + ) + for index in batch.bypassed: + candidate, load, _force = pending[index] + stream_bytes = ( + request.query_rows * request.dimension * 4 + + load.rows * request.dimension * 4 + + request.query_rows * load.rows * 4 + ) + if stream_bytes <= self.engine.max_device_bytes: + documents.append( + (candidate.candidate_id, load.rows, load.payload) + ) + admission_rejections += len(batch.bypassed) + deferred = [pending[index] for index in batch.deferred] + forced = [ + (pending[index][0], pending[index][1], True) + for index in batch.bypassed + if ( + request.query_rows * request.dimension * 4 + + pending[index][1].rows * request.dimension * 4 + + request.query_rows * pending[index][1].rows * 4 + > self.engine.max_device_bytes + ) + ] + made_progress = len(deferred) < len(pending) + if resident_documents: + flush_resident_documents() + made_progress = True + if deferred and not made_progress: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "GPU block cache cannot make progress with its configured capacity", + ) + pending = forced + deferred + if resident_documents: + flush_resident_documents() + metrics["gpu_admission_rejections"] = admission_rejections + metrics["cache_hits"] = cache_hits + metrics["host_cache_hits"] = cache_hits + metrics["gpu_cache_hits"] = gpu_cache_hits + metrics["gpu_cache_misses"] = gpu_cache_misses + metrics["gpu_chunks"] = gpu_chunks + metrics["resolve_ms"] = round( + max( + 0.0, + (time.monotonic() - resolve_started) * 1000.0 + - resident_queue_ms + - resident_compute_ms, + ), + 3, + ) + metrics["document_tokens"] = document_tokens + if self.gpu_cache is not None and isinstance( + request, protocol.ParsedExternalTensorRequest + ): + results = resident_results + queue_ms = resident_queue_ms + compute_ms = resident_compute_ms + if documents: + stream_results, stream_queue_ms, stream_compute_ms = ( + self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + ) + results.extend(stream_results) + queue_ms += stream_queue_ms + compute_ms += stream_compute_ms + metrics["streamed_candidates"] = len(documents) + else: + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + metrics["queue_ms"] = round(queue_ms, 3) + metrics["compute_ms"] = round(compute_ms, 3) + metrics["status"] = "ok" + return protocol.success_response(request.request_id, results, version) + except protocol.SidecarError as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, error.status, str(error), version + ) + except torch.OutOfMemoryError: + metrics.update(status="error", error_class="CudaOutOfMemory") + return protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "CUDA out of memory", + version, + ) + except Exception as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + f"TileMaxSim compute failed: {error}", + version, + ) + finally: + if self.gpu_cache is not None: + for _, handle in resident_documents: + self.gpu_cache.release(handle) + metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) + self.metrics.emit(metrics) + + def handle(self, connection: socket.socket) -> None: + deadline = time.monotonic() + self.request_timeout_seconds + request_id = 0 + version = protocol.VERSION + peer_credentials = None + if hasattr(socket, "SO_PEERCRED"): + try: + raw_credentials = connection.getsockopt( + socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i") + ) + peer_credentials = struct.unpack("3i", raw_credentials) + except OSError: + pass + try: + header = self._receive_exact_until( + connection, protocol.HEADER.size, deadline + ) + _, wire_version, _, request_id, body_len = protocol.HEADER.unpack(header) + if wire_version in protocol.SUPPORTED_VERSIONS: + version = wire_version + if body_len > self.limits.max_request_bytes - protocol.HEADER.size: + response = protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + version, + ) + else: + body = self._receive_exact_until(connection, body_len, deadline) + response = self.process_frame( + header + body, connection, deadline, peer_credentials + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired before socket write") + connection.settimeout(remaining) + connection.sendall(response) + except (BrokenPipeError, ConnectionResetError): + return + except (TimeoutError, socket.timeout): + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during socket I/O", + version, + ) + ) + except OSError: + pass + except Exception as error: + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + str(error), + version, + ) + ) + except OSError: + pass + + +def serve( + socket_path: Path, + socket_mode: int, + backlog: int, + max_inflight: int, + service: TileMaxsimService, + stop: threading.Event, + once: bool = False, +) -> None: + protocol.remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + slots = threading.BoundedSemaphore(max_inflight) + workers = ThreadPoolExecutor( + max_workers=max_inflight, thread_name_prefix="tilemaxsim" + ) + active = 0 + active_lock = threading.Lock() + + def handle(connection: socket.socket) -> None: + nonlocal active + try: + with connection: + service.handle(connection) + finally: + with active_lock: + active -= 1 + slots.release() + + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(backlog) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + ready: dict[str, object] = { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + if service.gpu_cache is not None: + ready["gpu_cache"] = service.gpu_cache.status() + service.metrics.emit(ready) + accepted = 0 + while not stop.is_set(): + if not slots.acquire(timeout=0.25): + continue + try: + connection, _ = listener.accept() + except TimeoutError: + slots.release() + continue + with active_lock: + active += 1 + current_active = active + service.metrics.emit( + {"event": "tilemaxsim_accept", "inflight": current_active} + ) + workers.submit(handle, connection) + accepted += 1 + if once and accepted == 1: + break + finally: + listener.close() + workers.shutdown(wait=True, cancel_futures=False) + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def memory_gb(value: str) -> int: + try: + return parse_memory_gb(value) + except ValueError as error: + raise argparse.ArgumentTypeError(str(error)) from error + + +def gpu_memory_gb(value: str) -> GpuArenaSpec: + try: + return parse_gpu_memory_gb(value) + except (ValueError, RuntimeError) as error: + raise argparse.ArgumentTypeError(str(error)) from error + + +def contract_roots( + values: list[str], parser: argparse.ArgumentParser +) -> dict[str, Path]: + roots = {} + for value in values: + if "=" not in value: + parser.error("--contract-root must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error("--contract-root must contain a nonempty ID and absolute path") + if contract in roots: + parser.error(f"duplicate --contract-root for {contract!r}") + roots[contract] = path + return roots + + +def contract_manifests( + values: list[str], parser: argparse.ArgumentParser +) -> list[tuple[str, Path]]: + manifests = [] + for value in values: + if "=" not in value: + parser.error("--resident-manifest must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error( + "--resident-manifest must contain a nonempty ID and absolute path" + ) + manifests.append((contract, path)) + return manifests + + +def prewarm_resident_cache( + manifests: list[tuple[str, Path]], + resolver: ContentAddressedResolver, + gpu_cache: GpuTensorCache, + metrics: JsonMetrics, + batch_size: int = 256, +) -> None: + completed = 0 + loaded_bytes = 0 + started = time.monotonic() + pending: list[tuple[protocol.ExternalTensorRequest, int]] = [] + + def flush() -> None: + nonlocal completed, loaded_bytes + if not pending: + return + keys = [resolver.key(request) for request, _ in pending] + probed, miss_indices = gpu_cache.probe_many(keys) + acquired = [handle for handle in probed if handle is not None] + try: + missing = [pending[index][0] for index in miss_indices] + resolved = resolver.resolve_many(missing) + loads = [ + GpuTensorLoad( + keys[index], + request.rows, + request.dimension, + request.dtype, + payload.payload, + True, + ) + for index, request, payload in zip( + miss_indices, missing, resolved, strict=True + ) + ] + batch = gpu_cache.acquire_many( + loads, + enforce_admission=False, + record_access=False, + count_stats=False, + ) + if ( + batch.bypassed + or batch.deferred + or any(handle is None for handle in batch.handles) + ): + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "resident manifest exceeds the configured GPU block arenas", + ) + acquired.extend(handle for handle in batch.handles if handle is not None) + finally: + for handle in acquired: + gpu_cache.release(handle) + completed += len(pending) + loaded_bytes += sum(size for _, size in pending) + if completed % 1000 < len(pending): + metrics.emit( + { + "event": "tilemaxsim_prewarm_progress", + "entries": completed, + "logical_bytes": loaded_bytes, + } + ) + pending.clear() + + for contract, path in manifests: + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(record, dict): + raise ValueError(f"{path}:{line_number}: record must be an object") + dtype_name = record.get("tensor_dtype") + if dtype_name == "float16": + dtype = protocol.DTYPE_F16 + elif dtype_name == "float32": + dtype = protocol.DTYPE_F32 + else: + raise ValueError(f"{path}:{line_number}: unsupported tensor_dtype") + tensor_ref = record.get("tensor_ref") + checksum = record.get("tensor_checksum") + rows = record.get("tensor_rows") + dimension = record.get("tensor_dim") + if not isinstance(tensor_ref, str) or not tensor_ref: + raise ValueError(f"{path}:{line_number}: invalid tensor_ref") + if not isinstance(checksum, str) or not checksum: + raise ValueError(f"{path}:{line_number}: invalid tensor_checksum") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_rows") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"{path}:{line_number}: invalid tensor_dim") + request = protocol.ExternalTensorRequest( + contract, tensor_ref, rows, dimension, dtype, checksum + ) + expected_bytes = protocol.checked_tensor_bytes(rows, dimension, dtype) + declared_bytes = record.get("canonical_bytes") + if declared_bytes is not None and declared_bytes != expected_bytes: + raise ValueError( + f"{path}:{line_number}: canonical_bytes disagrees with shape" + ) + pending.append((request, expected_bytes)) + if len(pending) >= batch_size: + flush() + flush() + if completed == 0: + raise ValueError("resident manifests contain no tensor descriptors") + metrics.emit( + { + "event": "tilemaxsim_prewarm_complete", + "entries": completed, + "logical_bytes": loaded_bytes, + "elapsed_ms": round((time.monotonic() - started) * 1000.0, 3), + "gpu_cache": gpu_cache.status(), + } + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument( + "--gpu-cache-mode", + choices=("lru", "resident"), + default="lru", + help="use evictable GPU arenas or pin a full descriptor manifest", + ) + parser.add_argument( + "--gpu-memory-gb", + action="append", + type=gpu_memory_gb, + default=[], + metavar="GPU=GB", + help="repeatable strict allocation, for example 1=20; enables TileMaxSim", + ) + parser.add_argument( + "--gpu-workspace-gb", + type=memory_gb, + default=2 * 1024**3, + help="per-GPU portion of the configured GB reserved for scoring temporaries", + ) + parser.add_argument( + "--gpu-block-kib", + type=positive_int, + default=32, + help="base page size inside each preallocated GPU tensor arena", + ) + parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--resident-manifest", + action="append", + default=[], + metavar="MODEL_CONTRACT_ID=/ABSOLUTE/PATH", + ) + parser.add_argument( + "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument("--max-batch-tokens", type=positive_int, default=1_000_000) + parser.add_argument( + "--max-tensor-bytes", type=positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=positive_int, default=65_536) + parser.add_argument( + "--host-cache-gb", + type=memory_gb, + default=8 * 1024**3, + help="decoded host-memory tensor cache size in GB", + ) + parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) + parser.add_argument("--max-inflight", type=positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--prewarm-batch-size", type=positive_int, default=256) + parser.add_argument("--backlog", type=positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument( + "--verify-full-shards", + action="store_true", + help="verify complete immutable shard digests lazily in addition to every tensor digest", + ) + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + + roots = contract_roots(args.contract_root, parser) + manifests = contract_manifests(args.resident_manifest, parser) + if not args.gpu_memory_gb: + parser.error( + "TileMaxSim is disabled until at least one --gpu-memory-gb GPU=GB is configured" + ) + if args.gpu_cache_mode == "resident" and not manifests: + parser.error( + "--gpu-cache-mode resident requires at least one --resident-manifest" + ) + if args.gpu_cache_mode == "lru" and manifests: + parser.error("--resident-manifest is valid only in resident mode") + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = ContentAddressedResolver( + roots, args.host_cache_gb, args.verify_full_shards + ) + metrics = JsonMetrics() + pool: GpuResourcePool | None = None + try: + pool = GpuResourcePool( + args.gpu_memory_gb, + args.gpu_workspace_gb, + args.gpu_block_kib * 1024, + ) + engine = TorchTileMaxsimEngine( + str(pool.primary_device), + args.gpu_workspace_gb, + args.allow_tf32, + args.max_cuda_inflight, + ) + gpu_cache = GpuTensorCache(pool, allow_eviction=args.gpu_cache_mode == "lru") + resident_engine = ResidentTorchTileMaxsimEngine( + pool, + args.gpu_workspace_gb, + args.allow_tf32, + args.max_cuda_inflight, + ) + if args.gpu_cache_mode == "resident": + prewarm_resident_cache( + manifests, + resolver, + gpu_cache, + metrics, + args.prewarm_batch_size, + ) + service = TileMaxsimService( + limits, + resolver, + engine, + args.request_timeout_ms, + metrics, + gpu_cache, + resident_engine, + pin_gpu_entries=args.gpu_cache_mode == "resident", + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + args.once, + ) + finally: + if pool is not None: + pool.close() + resolver.close() + + +if __name__ == "__main__": + main() diff --git a/services/tilemaxsim_gpu_cache.py b/services/tilemaxsim_gpu_cache.py new file mode 100644 index 00000000..0c9b5ef2 --- /dev/null +++ b/services/tilemaxsim_gpu_cache.py @@ -0,0 +1,971 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Process-owned GPU arenas and a bounded tensor cache for TileMaxSim.""" + +from __future__ import annotations + +import re +import threading +from bisect import bisect_left, insort +from collections import OrderedDict +from dataclasses import dataclass +from decimal import Decimal, InvalidOperation +from hashlib import blake2b +from math import ceil +from typing import Callable + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol + + +_GPU_MEMORY_GB = re.compile(r"^(?:cuda:)?([0-9]+)=([0-9]+(?:\.[0-9]+)?)$") +_MEMORY_GB = re.compile(r"^[0-9]+(?:\.[0-9]+)?$") +GIB = 1024**3 +# A 32 KiB page keeps the allocator metadata small while avoiding the 8.82% +# rounding loss observed with the former 256 KiB default on 472--483 KiB +# document tensors. The native daemon exposes the same default via +# ``--gpu-block-kib``. +DEFAULT_BLOCK_BYTES = 32 * 1024 +DEFAULT_STAGING_BYTES = 64 * 1024 * 1024 + + +@dataclass(frozen=True) +class GpuArenaSpec: + device: str + total_bytes: int + + +def parse_gpu_memory_gb(value: str) -> GpuArenaSpec: + """Parse the public ``GPU=GB`` configuration into an internal byte budget.""" + + match = _GPU_MEMORY_GB.fullmatch(value.strip()) + if match is None: + raise ValueError( + "GPU memory must be GPU=GB, for example 1=20; byte suffixes are not accepted" + ) + raw_index, raw_gb = match.groups() + try: + total_bytes = int(Decimal(raw_gb) * GIB) + except InvalidOperation as error: + raise ValueError("GPU memory GB value is invalid") from error + if total_bytes <= 0: + raise ValueError("GPU memory GB value must be positive") + return GpuArenaSpec(f"cuda:{int(raw_index)}", total_bytes) + + +def parse_memory_gb(value: str) -> int: + if _MEMORY_GB.fullmatch(value.strip()) is None: + raise ValueError("memory size must be a positive number of GB") + try: + total_bytes = int(Decimal(value.strip()) * GIB) + except InvalidOperation as error: + raise ValueError("memory size must be a positive number of GB") from error + if total_bytes <= 0: + raise ValueError("memory size must be a positive number of GB") + return total_bytes + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +class FreeExtentAllocator: + """Best-fit allocator for contiguous slices of a preallocated byte arena.""" + + def __init__(self, capacity: int, alignment: int = 256) -> None: + if capacity <= 0: + raise ValueError("arena capacity must be positive") + self.capacity = capacity + self.alignment = alignment + self.extents: list[tuple[int, int]] = [(0, capacity)] + + @property + def free_bytes(self) -> int: + return sum(length for _, length in self.extents) + + @property + def largest_free_extent(self) -> int: + return max((length for _, length in self.extents), default=0) + + def allocation_bytes(self, payload_bytes: int) -> int: + return _align_up(payload_bytes, self.alignment) + + def allocate(self, payload_bytes: int) -> tuple[int, int] | None: + required = self.allocation_bytes(payload_bytes) + choices = [ + (length, index, start) + for index, (start, length) in enumerate(self.extents) + if length >= required + ] + if not choices: + return None + length, index, start = min(choices) + if length == required: + self.extents.pop(index) + else: + self.extents[index] = (start + required, length - required) + return start, required + + def release(self, start: int, length: int) -> None: + if start < 0 or length <= 0 or start + length > self.capacity: + raise ValueError("released extent is outside the arena") + self.extents.append((start, length)) + self.extents.sort() + merged: list[tuple[int, int]] = [] + for extent_start, extent_length in self.extents: + if merged and merged[-1][0] + merged[-1][1] == extent_start: + previous_start, previous_length = merged[-1] + merged[-1] = (previous_start, previous_length + extent_length) + else: + merged.append((extent_start, extent_length)) + self.extents = merged + + +class FixedBlockAllocator: + """Best-fit page-run allocator over one process-owned GPU arena. + + Tensors keep the dense layout required by the TileMaxSim kernel, but their + runs are rounded only to the base page instead of the next power of two. + Free runs are indexed by both address and exact size. Allocation therefore + finds the smallest suitable run, while release coalesces adjacent runs. + """ + + def __init__(self, capacity: int, block_bytes: int = DEFAULT_BLOCK_BYTES) -> None: + if capacity <= 0 or block_bytes <= 0: + raise ValueError("arena capacity and block size must be positive") + if block_bytes % 256: + raise ValueError("GPU block size must be 256-byte aligned") + self.block_bytes = block_bytes + self.block_count = capacity // block_bytes + if self.block_count == 0: + raise ValueError("arena must contain at least one GPU block") + self.capacity = self.block_count * block_bytes + self._free_by_start: dict[int, int] = {} + self._free_by_size: dict[int, set[int]] = {} + self._starts: list[int] = [] + self._sizes: list[int] = [] + self._allocated: dict[int, int] = {} + self._free_blocks = 0 + self._add_free_run(0, self.block_count) + + def _add_free_run(self, start: int, blocks: int) -> None: + if blocks <= 0 or start < 0 or start + blocks > self.block_count: + raise ValueError("free GPU run is outside the arena") + if start in self._free_by_start: + raise ValueError("duplicate free GPU run") + self._free_by_start[start] = blocks + insort(self._starts, start) + bucket = self._free_by_size.get(blocks) + if bucket is None: + bucket = set() + self._free_by_size[blocks] = bucket + insort(self._sizes, blocks) + bucket.add(start) + self._free_blocks += blocks + + def _remove_free_run(self, start: int, blocks: int) -> None: + if self._free_by_start.get(start) != blocks: + raise ValueError("free GPU run directory is inconsistent") + del self._free_by_start[start] + start_index = bisect_left(self._starts, start) + if start_index == len(self._starts) or self._starts[start_index] != start: + raise ValueError("free GPU address index is inconsistent") + self._starts.pop(start_index) + bucket = self._free_by_size[blocks] + bucket.remove(start) + if not bucket: + del self._free_by_size[blocks] + size_index = bisect_left(self._sizes, blocks) + if size_index == len(self._sizes) or self._sizes[size_index] != blocks: + raise ValueError("free GPU size index is inconsistent") + self._sizes.pop(size_index) + self._free_blocks -= blocks + + @property + def free_blocks(self) -> int: + return self._free_blocks + + @property + def free_bytes(self) -> int: + return self.free_blocks * self.block_bytes + + @property + def largest_free_extent(self) -> int: + return self._sizes[-1] * self.block_bytes if self._sizes else 0 + + def blocks_for(self, payload_bytes: int) -> int: + if payload_bytes <= 0: + raise ValueError("payload size must be positive") + return ceil(payload_bytes / self.block_bytes) + + def allocation_bytes(self, payload_bytes: int) -> int: + return self.blocks_for(payload_bytes) * self.block_bytes + + def allocate(self, payload_bytes: int) -> tuple[int, ...] | None: + required = self.blocks_for(payload_bytes) + size_index = bisect_left(self._sizes, required) + if size_index == len(self._sizes): + return None + available = self._sizes[size_index] + start = min(self._free_by_size[available]) + self._remove_free_run(start, available) + if available > required: + self._add_free_run(start + required, available - required) + self._allocated[start] = required + return tuple(range(start, start + required)) + + def release(self, blocks: tuple[int, ...]) -> None: + if ( + not blocks + or len(set(blocks)) != len(blocks) + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + ): + raise ValueError("released GPU blocks are invalid") + allocated_blocks = self._allocated.pop(blocks[0], None) + if allocated_blocks is None: + raise ValueError("GPU block was released more than once") + if len(blocks) != allocated_blocks: + raise ValueError("released GPU page run has the wrong size") + start = blocks[0] + length = allocated_blocks + insertion = bisect_left(self._starts, start) + if insertion: + previous_start = self._starts[insertion - 1] + previous_length = self._free_by_start[previous_start] + if previous_start + previous_length == start: + self._remove_free_run(previous_start, previous_length) + start = previous_start + length += previous_length + next_start = start + length + next_length = self._free_by_start.get(next_start) + if next_length is not None: + self._remove_free_run(next_start, next_length) + length += next_length + self._add_free_run(start, length) + + +class TinyLfuSketch: + """A small aging count-min sketch for cache admission and GDSF frequency.""" + + def __init__(self, width: int = 4096, depth: int = 4) -> None: + if width <= 0 or depth <= 0: + raise ValueError("TinyLFU dimensions must be positive") + self.width = width + self.depth = depth + self.tables = [[0] * width for _ in range(depth)] + self.samples = 0 + self.reset_at = width * 10 + + def _indices(self, key: tuple[object, ...]) -> tuple[int, ...]: + digest = blake2b(repr(key).encode("utf-8"), digest_size=16).digest() + return tuple( + int.from_bytes(digest[row * 4 : row * 4 + 4], "little") % self.width + for row in range(self.depth) + ) + + def increment(self, key: tuple[object, ...]) -> int: + indices = self._indices(key) + for row, index in enumerate(indices): + if self.tables[row][index] < 65535: + self.tables[row][index] += 1 + self.samples += 1 + estimate = min(self.tables[row][index] for row, index in enumerate(indices)) + if self.samples >= self.reset_at: + for table in self.tables: + for index, value in enumerate(table): + table[index] = value // 2 + self.samples //= 2 + return max(1, estimate) + + def estimate(self, key: tuple[object, ...]) -> int: + return min( + self.tables[row][index] for row, index in enumerate(self._indices(key)) + ) + + +class GpuArena: + """A CUDA byte buffer acquired atomically during process startup.""" + + def __init__( + self, + spec: GpuArenaSpec, + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: + self.device = torch.device(spec.device) + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.index is None or self.device.index >= torch.cuda.device_count(): + raise RuntimeError(f"configured CUDA device is unavailable: {spec.device}") + if workspace_bytes <= 0 or workspace_bytes >= spec.total_bytes: + raise RuntimeError( + f"{spec.device} allocation must exceed its TileMaxSim workspace" + ) + self.total_bytes = spec.total_bytes + self.workspace_bytes = workspace_bytes + raw_capacity = spec.total_bytes - workspace_bytes + self.allocator = FixedBlockAllocator(raw_capacity, block_bytes) + self.capacity = self.allocator.capacity + self.reserved_workspace_bytes = spec.total_bytes - self.capacity + if self.capacity <= 0: + raise RuntimeError(f"{spec.device} has no aligned tensor-cache capacity") + self.storage: torch.Tensor | None = None + self.host_staging: torch.Tensor | None = None + self.copy_stream: torch.cuda.Stream | None = None + self.h2d_batches = 0 + self.h2d_copy_calls = 0 + self.h2d_bytes = 0 + + with torch.cuda.device(self.device): + free_bytes, device_bytes = torch.cuda.mem_get_info(self.device) + self.device_bytes = device_bytes + if free_bytes < spec.total_bytes: + raise RuntimeError( + f"cannot acquire {spec.total_bytes} bytes on {spec.device}: " + f"only {free_bytes} bytes are free" + ) + try: + self.storage = torch.empty( + self.capacity, dtype=torch.uint8, device=self.device + ) + # Reserve the remaining configured budget in PyTorch's CUDA + # allocator. Releasing this temporary tensor leaves the block + # in the process-owned caching allocator for TileMaxSim + # workspaces instead of returning it to another process. + workspace = torch.empty( + self.reserved_workspace_bytes, + dtype=torch.uint8, + device=self.device, + ) + torch.cuda.synchronize(self.device) + del workspace + staging_bytes = min(self.capacity, DEFAULT_STAGING_BYTES) + staging_bytes = max( + self.allocator.block_bytes, + staging_bytes + // self.allocator.block_bytes + * self.allocator.block_bytes, + ) + self.host_staging = torch.empty( + staging_bytes, dtype=torch.uint8, pin_memory=True + ) + # Fault and pin every staging page during startup so the first + # cache-miss batch does not pay a request-path NUMA/page cost. + self.host_staging.zero_() + self.copy_stream = torch.cuda.Stream(device=self.device) + except Exception: + self.storage = None + self.host_staging = None + self.copy_stream = None + torch.cuda.empty_cache() + raise + + def tensor( + self, + blocks: tuple[int, ...], + payload_bytes: int, + rows: int, + dimension: int, + dtype: int, + ) -> torch.Tensor: + if ( + self.storage is None + or self.host_staging is None + or self.copy_stream is None + ): + raise RuntimeError("GPU arena is closed") + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + block_bytes = self.allocator.block_bytes + if all(right == left + 1 for left, right in zip(blocks, blocks[1:])): + raw = self.storage.narrow(0, blocks[0] * block_bytes, payload_bytes) + else: + raw = torch.cat( + [ + self.storage.narrow(0, block * block_bytes, block_bytes) + for block in blocks + ] + ).narrow(0, 0, payload_bytes) + return raw.view(scalar_dtype).reshape(rows, dimension) + + def copy_many_from_host(self, items: list[tuple[tuple[int, ...], bytes]]) -> None: + if self.storage is None: + raise RuntimeError("GPU arena is closed") + if not items: + return + block_bytes = self.allocator.block_bytes + runs: list[tuple[int, int, bytes]] = [] + for blocks, payload in items: + if ( + not blocks + or blocks != tuple(range(blocks[0], blocks[0] + len(blocks))) + or len(payload) > len(blocks) * block_bytes + ): + raise RuntimeError("GPU upload requires one valid contiguous page run") + runs.append((blocks[0], len(blocks), payload)) + runs.sort(key=lambda item: item[0]) + staging = self.host_staging + staging_array = staging.numpy() + staging_bytes = staging.numel() + copy_calls = 0 + with torch.cuda.device(self.device): + stream = self.copy_stream + run_index = 0 + while run_index < len(runs): + first_block, run_blocks, payload = runs[run_index] + allocation_bytes = run_blocks * block_bytes + if allocation_bytes > staging_bytes: + source_offset = 0 + while source_offset < len(payload): + length = min(staging_bytes, len(payload) - source_offset) + staging_array[:length] = np.frombuffer( + payload, + dtype=np.uint8, + count=length, + offset=source_offset, + ) + with torch.cuda.stream(stream): + self.storage.narrow( + 0, + first_block * block_bytes + source_offset, + length, + ).copy_(staging.narrow(0, 0, length), non_blocking=True) + stream.synchronize() + copy_calls += 1 + source_offset += length + run_index += 1 + continue + + batch_first_block = first_block + expected_block = first_block + staging_offset = 0 + while run_index < len(runs): + start_block, count, current_payload = runs[run_index] + current_allocation = count * block_bytes + if ( + start_block != expected_block + or staging_offset + current_allocation > staging_bytes + ): + break + payload_end = staging_offset + len(current_payload) + staging_array[staging_offset:payload_end] = np.frombuffer( + current_payload, dtype=np.uint8 + ) + allocation_end = staging_offset + current_allocation + staging_array[payload_end:allocation_end] = 0 + staging_offset = allocation_end + expected_block += count + run_index += 1 + with torch.cuda.stream(stream): + self.storage.narrow( + 0, batch_first_block * block_bytes, staging_offset + ).copy_( + staging.narrow(0, 0, staging_offset), + non_blocking=True, + ) + copy_calls += 1 + stream.synchronize() + self.h2d_batches += 1 + self.h2d_copy_calls += copy_calls + self.h2d_bytes += sum(len(payload) for _, payload in items) + + def copy_from_host(self, blocks: tuple[int, ...], payload: bytes) -> None: + self.copy_many_from_host([(blocks, payload)]) + + def status(self) -> dict[str, object]: + return { + "device": str(self.device), + "allocated_gb": round(self.total_bytes / GIB, 3), + "tensor_capacity_gb": round(self.capacity / GIB, 3), + "workspace_gb": round(self.workspace_bytes / GIB, 3), + "allocated_bytes": self.total_bytes, + "tensor_capacity_bytes": self.capacity, + "workspace_bytes": self.workspace_bytes, + "tensor_free_bytes": self.allocator.free_bytes, + "largest_free_extent_bytes": self.allocator.largest_free_extent, + "block_bytes": self.allocator.block_bytes, + "block_count": self.allocator.block_count, + "free_blocks": self.allocator.free_blocks, + "h2d_batches": self.h2d_batches, + "h2d_copy_calls": self.h2d_copy_calls, + "h2d_bytes": self.h2d_bytes, + "host_staging_bytes": self.host_staging.numel() + if self.host_staging is not None + else 0, + } + + def close(self) -> None: + if self.storage is None: + return + with torch.cuda.device(self.device): + self.storage = None + self.host_staging = None + self.copy_stream = None + torch.cuda.empty_cache() + + +class GpuResourcePool: + """Own all configured CUDA allocations or fail without a partial pool.""" + + def __init__( + self, + specs: list[GpuArenaSpec], + workspace_bytes: int, + block_bytes: int = DEFAULT_BLOCK_BYTES, + ) -> None: + if not specs: + raise RuntimeError("at least one GPU allocation is required") + devices = [spec.device for spec in specs] + if len(devices) != len(set(devices)): + raise RuntimeError("each CUDA device may be configured only once") + self.arenas: list[GpuArena] = [] + try: + for spec in specs: + self.arenas.append(GpuArena(spec, workspace_bytes, block_bytes)) + except Exception: + self.close() + raise + + @property + def primary_device(self) -> torch.device: + return self.arenas[0].device + + def status(self) -> list[dict[str, object]]: + return [arena.status() for arena in self.arenas] + + def close(self) -> None: + for arena in self.arenas: + arena.close() + self.arenas.clear() + + +@dataclass +class _GpuCacheEntry: + key: tuple[object, ...] + arena: GpuArena + blocks: tuple[int, ...] + allocated_bytes: int + payload_bytes: int + rows: int + dimension: int + dtype: int + references: int = 0 + pinned: bool = False + priority: float = 0.0 + + +@dataclass(frozen=True) +class GpuTensorHandle: + entry: _GpuCacheEntry + + @property + def arena(self) -> GpuArena: + return self.entry.arena + + @property + def device(self) -> torch.device: + return self.entry.arena.device + + @property + def rows(self) -> int: + return self.entry.rows + + @property + def dimension(self) -> int: + return self.entry.dimension + + @property + def dtype(self) -> int: + return self.entry.dtype + + @property + def payload_bytes(self) -> int: + return self.entry.payload_bytes + + @property + def offset_bytes(self) -> int: + return self.entry.blocks[0] * self.entry.arena.allocator.block_bytes + + @property + def block_ids(self) -> tuple[int, ...]: + return self.entry.blocks + + @property + def block_bytes(self) -> int: + return self.entry.arena.allocator.block_bytes + + def tensor(self) -> torch.Tensor: + return self.entry.arena.tensor( + self.entry.blocks, + self.entry.payload_bytes, + self.entry.rows, + self.entry.dimension, + self.entry.dtype, + ) + + +@dataclass(frozen=True) +class GpuTensorLoad: + key: tuple[object, ...] + rows: int + dimension: int + dtype: int + payload: bytes + pin: bool = False + + +@dataclass(frozen=True) +class GpuAcquireBatch: + handles: tuple[GpuTensorHandle | None, ...] + bypassed: tuple[int, ...] + deferred: tuple[int, ...] + hits: int + misses: int + admitted: int + + +class GpuTensorCache: + """Fixed-block GPU cache with TinyLFU admission and GDSF eviction.""" + + def __init__(self, pool: GpuResourcePool, allow_eviction: bool) -> None: + self.pool = pool + self.allow_eviction = allow_eviction + self.entries: OrderedDict[tuple[object, ...], _GpuCacheEntry] = OrderedDict() + self.lock = threading.Lock() + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.loaded_bytes = 0 + self.admission_rejections = 0 + self.inflation = 0.0 + self.sketch = TinyLfuSketch() + + def _find_arena(self, payload_bytes: int) -> GpuArena | None: + candidates = [ + arena + for arena in self.pool.arenas + if arena.allocator.largest_free_extent + >= arena.allocator.allocation_bytes(payload_bytes) + ] + if not candidates: + return None + return max(candidates, key=lambda arena: arena.allocator.free_bytes) + + def _victim(self, arena: GpuArena | None = None) -> _GpuCacheEntry | None: + candidates = [ + entry + for entry in self.entries.values() + if not entry.references + and not entry.pinned + and (arena is None or entry.arena is arena) + ] + return min(candidates, key=lambda entry: entry.priority, default=None) + + def _evict(self, entry: _GpuCacheEntry) -> None: + current = self.entries.pop(entry.key, None) + if current is not entry: + raise RuntimeError("GPU eviction directory is inconsistent") + entry.arena.allocator.release(entry.blocks) + self.inflation = max(self.inflation, entry.priority) + self.evictions += 1 + + @staticmethod + def _entry_cost(entry: _GpuCacheEntry) -> int: + return len(entry.blocks) + + def _priority(self, key: tuple[object, ...], blocks: int) -> float: + return self.inflation + self.sketch.estimate(key) / max(1, blocks) + + def _allocate( + self, + key: tuple[object, ...], + payload_bytes: int, + enforce_admission: bool, + ) -> tuple[GpuArena, tuple[int, ...], int] | None: + arena = self._find_arena(payload_bytes) + capable = [ + candidate + for candidate in self.pool.arenas + if candidate.allocator.capacity + >= candidate.allocator.allocation_bytes(payload_bytes) + ] + if not capable: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one tensor exceeds every configured GPU block arena", + ) + if arena is None and self.allow_eviction: + for candidate in sorted( + capable, key=lambda item: item.allocator.free_bytes, reverse=True + ): + required_blocks = candidate.allocator.blocks_for(payload_bytes) + required_bytes = candidate.allocator.allocation_bytes(payload_bytes) + while candidate.allocator.largest_free_extent < required_bytes: + victim = self._victim(candidate) + if victim is None: + break + candidate_priority = self._priority(key, required_blocks) + if enforce_admission and candidate_priority <= victim.priority: + self.admission_rejections += 1 + return None + self._evict(victim) + if candidate.allocator.largest_free_extent >= required_bytes: + arena = candidate + break + if arena is None: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "configured GPU tensor arenas have insufficient free blocks", + ) + blocks = arena.allocator.allocate(payload_bytes) + assert blocks is not None + return arena, blocks, len(blocks) * arena.allocator.block_bytes + + def acquire( + self, + key: tuple[object, ...], + rows: int, + dimension: int, + dtype: int, + loader: Callable[[], bytes], + *, + pin: bool = False, + ) -> tuple[GpuTensorHandle, bool]: + with self.lock: + frequency = self.sketch.increment(key) + cached = self.entries.get(key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + + payload = loader() + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + + with self.lock: + cached = self.entries.get(key) + if cached is not None: + frequency = self.sketch.increment(key) + cached.references += 1 + cached.pinned = cached.pinned or pin + cached.priority = self.inflation + frequency / self._entry_cost(cached) + self.entries.move_to_end(key) + self.hits += 1 + return GpuTensorHandle(cached), True + allocated = self._allocate(key, len(payload), enforce_admission=False) + assert allocated is not None + arena, blocks, allocated_bytes = allocated + try: + arena.copy_from_host(blocks, payload) + except Exception: + arena.allocator.release(blocks) + raise + entry = _GpuCacheEntry( + key, + arena, + blocks, + allocated_bytes, + len(payload), + rows, + dimension, + dtype, + references=1, + pinned=pin, + priority=self._priority(key, len(blocks)), + ) + self.entries[key] = entry + self.misses += 1 + self.loaded_bytes += len(payload) + return GpuTensorHandle(entry), False + + def acquire_many( + self, + loads: list[GpuTensorLoad], + *, + enforce_admission: bool = True, + record_access: bool = True, + count_stats: bool = True, + ) -> GpuAcquireBatch: + """Acquire a request working set and upload all new slabs in batches. + + ``deferred`` items could not be allocated while earlier handles in the + same request are referenced. The caller scores/releases the returned + handles and retries those items. ``bypassed`` items lost TinyLFU + admission and should be scored through the bounded streaming engine. + """ + + handles: list[GpuTensorHandle | None] = [None] * len(loads) + bypassed: list[int] = [] + deferred: list[int] = [] + new_entries: list[tuple[int, _GpuCacheEntry, bytes]] = [] + hits = 0 + misses = 0 + admitted = 0 + with self.lock: + try: + for index, load in enumerate(loads): + expected = protocol.checked_tensor_bytes( + load.rows, load.dimension, load.dtype + ) + if len(load.payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match its shape", + ) + frequency = ( + self.sketch.increment(load.key) + if record_access + else max(1, self.sketch.estimate(load.key)) + ) + cached = self.entries.get(load.key) + if cached is not None: + cached.references += 1 + cached.pinned = cached.pinned or load.pin + cached.priority = self.inflation + frequency / self._entry_cost( + cached + ) + self.entries.move_to_end(load.key) + handles[index] = GpuTensorHandle(cached) + hits += 1 + continue + misses += 1 + try: + allocation = self._allocate( + load.key, + len(load.payload), + enforce_admission and not load.pin, + ) + except protocol.SidecarError as error: + if "one tensor exceeds" in str(error): + bypassed.append(index) + continue + deferred.append(index) + continue + if allocation is None: + bypassed.append(index) + continue + arena, blocks, allocated_bytes = allocation + entry = _GpuCacheEntry( + load.key, + arena, + blocks, + allocated_bytes, + len(load.payload), + load.rows, + load.dimension, + load.dtype, + references=1, + pinned=load.pin, + priority=self._priority(load.key, len(blocks)), + ) + self.entries[load.key] = entry + handles[index] = GpuTensorHandle(entry) + new_entries.append((index, entry, load.payload)) + admitted += 1 + + by_arena: dict[ + int, tuple[GpuArena, list[tuple[tuple[int, ...], bytes]]] + ] = {} + for _, entry, payload in new_entries: + bucket = by_arena.setdefault(id(entry.arena), (entry.arena, [])) + bucket[1].append((entry.blocks, payload)) + for arena, items in by_arena.values(): + arena.copy_many_from_host(items) + if count_stats: + self.hits += hits + self.misses += misses + self.loaded_bytes += sum(len(payload) for _, _, payload in new_entries) + except Exception: + new_ids = {id(entry) for _, entry, _ in new_entries} + for handle in handles: + if handle is None: + continue + entry = handle.entry + if id(entry) in new_ids: + if self.entries.pop(entry.key, None) is entry: + entry.arena.allocator.release(entry.blocks) + else: + entry.references -= 1 + raise + return GpuAcquireBatch( + tuple(handles), + tuple(bypassed), + tuple(deferred), + hits, + misses, + admitted, + ) + + def probe_many( + self, keys: list[tuple[object, ...]] + ) -> tuple[tuple[GpuTensorHandle | None, ...], tuple[int, ...]]: + """Acquire GPU hits without resolving the corresponding host payloads.""" + + handles: list[GpuTensorHandle | None] = [None] * len(keys) + misses = [] + with self.lock: + for index, key in enumerate(keys): + frequency = self.sketch.increment(key) + entry = self.entries.get(key) + if entry is None: + misses.append(index) + self.misses += 1 + continue + entry.references += 1 + entry.priority = self.inflation + frequency / self._entry_cost(entry) + self.entries.move_to_end(key) + handles[index] = GpuTensorHandle(entry) + self.hits += 1 + return tuple(handles), tuple(misses) + + def release(self, handle: GpuTensorHandle) -> None: + with self.lock: + entry = self.entries.get(handle.entry.key) + if entry is not handle.entry or entry.references <= 0: + raise RuntimeError("GPU tensor handle was released more than once") + entry.references -= 1 + + def status(self) -> dict[str, object]: + with self.lock: + allocated_bytes = sum( + entry.allocated_bytes for entry in self.entries.values() + ) + payload_bytes = sum(entry.payload_bytes for entry in self.entries.values()) + return { + "allocator": "segregated-page-runs", + "entries": len(self.entries), + "pinned_entries": sum(entry.pinned for entry in self.entries.values()), + "active_references": sum( + entry.references for entry in self.entries.values() + ), + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "admission_rejections": self.admission_rejections, + "policy": "tinylfu-gdsf", + "gdsf_inflation": self.inflation, + "loaded_bytes": self.loaded_bytes, + "allocated_bytes": allocated_bytes, + "payload_bytes": payload_bytes, + "internal_waste_bytes": allocated_bytes - payload_bytes, + "arenas": self.pool.status(), + } diff --git a/services/tilemaxsim_shard.py b/services/tilemaxsim_shard.py new file mode 100644 index 00000000..453d58c9 --- /dev/null +++ b/services/tilemaxsim_shard.py @@ -0,0 +1,350 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Immutable, content-addressed TileMaxSim tensor shard format. + +The data files contain canonical tensor bytes and alignment padding only. A +generation index maps tensor SHA-256 digests to byte ranges. Data files are +published under their own SHA-256, so a writer never mutates a visible shard. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Iterable + + +FORMAT = "vectorchord.tilemaxsim.shards" +VERSION = 1 +INDEX_NAME = "tilemaxsim-shards-v1.json" +SHARD_DIRECTORY = "shards" +DEFAULT_ALIGNMENT = 4096 +DEFAULT_SHARD_BYTES = 2 * 1024**3 + + +def align_up(value: int, alignment: int) -> int: + if alignment <= 0 or alignment & (alignment - 1): + raise ValueError("alignment must be a positive power of two") + return (value + alignment - 1) // alignment * alignment + + +@dataclass(frozen=True) +class ShardEntry: + digest: str + shard: str + offset: int + length: int + rows: int + dimension: int + dtype: str + + +@dataclass(frozen=True) +class ShardFile: + name: str + size: int + checksum: str + + +@dataclass(frozen=True) +class ShardIndex: + alignment: int + shards: dict[str, ShardFile] + entries: dict[str, ShardEntry] + + +def _digest_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +class ImmutableShardWriter: + """Deterministically publish bounded immutable shards and one atomic index.""" + + def __init__( + self, + root: Path, + target_bytes: int = DEFAULT_SHARD_BYTES, + alignment: int = DEFAULT_ALIGNMENT, + fsync: bool = True, + ) -> None: + if target_bytes <= 0: + raise ValueError("target shard bytes must be positive") + if target_bytes < alignment: + raise ValueError("target shard bytes must be at least one alignment unit") + align_up(0, alignment) + self.root = root + self.shard_root = root / SHARD_DIRECTORY + self.target_bytes = target_bytes + self.alignment = alignment + self.fsync = fsync + self.entries: dict[str, ShardEntry] = {} + self.shards: list[ShardFile] = [] + self._stream: BinaryIO | None = None + self._temporary: Path | None = None + self._hasher: hashlib._Hash | None = None + self._offset = 0 + self._pending: list[tuple[str, int, int, int, int, str]] = [] + self.root.mkdir(parents=True, exist_ok=True) + self.shard_root.mkdir(parents=True, exist_ok=True) + + def _open(self) -> None: + descriptor, name = tempfile.mkstemp( + prefix=".tilemaxsim-shard-", suffix=".tmp", dir=self.shard_root + ) + self._stream = os.fdopen(descriptor, "wb", closefd=True) + self._temporary = Path(name) + self._hasher = hashlib.sha256() + self._offset = 0 + self._pending = [] + + def _write(self, payload: bytes | memoryview) -> None: + assert self._stream is not None and self._hasher is not None + self._stream.write(payload) + self._hasher.update(payload) + self._offset += len(payload) + + def add( + self, + digest: str, + payload: bytes | memoryview, + rows: int, + dimension: int, + dtype: str, + ) -> None: + if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("invalid tensor SHA-256 digest") + if digest in self.entries or any(item[0] == digest for item in self._pending): + return + if not payload: + raise ValueError("tensor payload must not be empty") + if rows <= 0 or dimension <= 0 or dtype not in ("float16", "float32"): + raise ValueError("invalid tensor metadata") + padded = align_up(len(payload), self.alignment) + if self._stream is not None and self._offset and self._offset + padded > self.target_bytes: + self._finish_shard() + if self._stream is None: + self._open() + offset = self._offset + self._write(payload) + padding = padded - len(payload) + if padding: + self._write(bytes(min(padding, self.alignment))) + remaining = padding - min(padding, self.alignment) + while remaining: + chunk = min(remaining, self.alignment) + self._write(bytes(chunk)) + remaining -= chunk + self._pending.append((digest, offset, len(payload), rows, dimension, dtype)) + + def _finish_shard(self) -> None: + if self._stream is None: + return + assert self._temporary is not None and self._hasher is not None + stream = self._stream + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + stream.close() + digest = self._hasher.hexdigest() + name = f"sha256-{digest}.vts" + destination = self.shard_root / name + size = self._offset + if destination.exists(): + if destination.stat().st_size != size or _digest_file(destination) != digest: + raise ValueError(f"existing immutable shard is corrupt: {destination}") + self._temporary.unlink() + else: + os.replace(self._temporary, destination) + if self.fsync: + descriptor = os.open(self.shard_root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + relative = f"{SHARD_DIRECTORY}/{name}" + shard = ShardFile(relative, size, f"sha256:{digest}") + self.shards.append(shard) + for tensor_digest, offset, length, rows, dimension, dtype in self._pending: + self.entries[tensor_digest] = ShardEntry( + tensor_digest, + relative, + offset, + length, + rows, + dimension, + dtype, + ) + self._stream = None + self._temporary = None + self._hasher = None + self._offset = 0 + self._pending = [] + + def finish(self) -> Path: + self._finish_shard() + if not self.entries: + raise ValueError("cannot publish an empty shard generation") + document = { + "format": FORMAT, + "version": VERSION, + "alignment": self.alignment, + "shards": [ + {"name": shard.name, "bytes": shard.size, "checksum": shard.checksum} + for shard in self.shards + ], + "entries": [ + { + "digest": entry.digest, + "shard": entry.shard, + "offset": entry.offset, + "length": entry.length, + "rows": entry.rows, + "dimension": entry.dimension, + "dtype": entry.dtype, + } + for entry in self.entries.values() + ], + } + descriptor, temporary = tempfile.mkstemp( + prefix=f".{INDEX_NAME}.", suffix=".tmp", dir=self.root, text=True + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + json.dump(document, stream, separators=(",", ":"), sort_keys=True) + stream.write("\n") + stream.flush() + if self.fsync: + os.fsync(stream.fileno()) + destination = self.root / INDEX_NAME + os.replace(temporary, destination) + if self.fsync: + directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory) + finally: + os.close(directory) + return destination + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + + def close(self) -> None: + if self._stream is not None: + self._stream.close() + if self._temporary is not None: + try: + self._temporary.unlink() + except FileNotFoundError: + pass + self._stream = None + self._temporary = None + + +def parse_index(document: object) -> ShardIndex: + if not isinstance(document, dict): + raise ValueError("shard index must be a JSON object") + if document.get("format") != FORMAT or document.get("version") != VERSION: + raise ValueError("unsupported TileMaxSim shard index") + alignment = document.get("alignment") + if not isinstance(alignment, int): + raise ValueError("shard index has invalid alignment") + align_up(0, alignment) + raw_shards = document.get("shards") + raw_entries = document.get("entries") + if not isinstance(raw_shards, list) or not isinstance(raw_entries, list): + raise ValueError("shard index has invalid arrays") + shards: dict[str, ShardFile] = {} + for raw in raw_shards: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid shard") + name, size, checksum = raw.get("name"), raw.get("bytes"), raw.get("checksum") + if ( + not isinstance(name, str) + or not name.startswith(f"{SHARD_DIRECTORY}/sha256-") + or "/" in name[len(SHARD_DIRECTORY) + 1 :] + or not isinstance(size, int) + or size <= 0 + or not isinstance(checksum, str) + or not checksum.startswith("sha256:") + or name != f"{SHARD_DIRECTORY}/sha256-{checksum.removeprefix('sha256:')}.vts" + or name in shards + ): + raise ValueError("shard index contains invalid shard metadata") + shards[name] = ShardFile(name, size, checksum) + entries: dict[str, ShardEntry] = {} + intervals: dict[str, list[tuple[int, int]]] = {name: [] for name in shards} + for raw in raw_entries: + if not isinstance(raw, dict): + raise ValueError("shard index contains an invalid tensor entry") + digest = raw.get("digest") + shard = raw.get("shard") + offset = raw.get("offset") + length = raw.get("length") + rows = raw.get("rows") + dimension = raw.get("dimension") + dtype = raw.get("dtype") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(c not in "0123456789abcdef" for c in digest) + or not isinstance(shard, str) + or shard not in shards + or not isinstance(offset, int) + or offset < 0 + or offset % alignment + or not isinstance(length, int) + or length <= 0 + or not isinstance(rows, int) + or rows <= 0 + or not isinstance(dimension, int) + or dimension <= 0 + or dtype not in ("float16", "float32") + or digest in entries + ): + raise ValueError("shard index contains invalid tensor metadata") + scalar_bytes = 2 if dtype == "float16" else 4 + if length != rows * dimension * scalar_bytes: + raise ValueError("shard tensor length disagrees with its shape") + end = offset + length + if end > shards[shard].size: + raise ValueError("shard tensor range is outside its data file") + intervals[shard].append((offset, align_up(end, alignment))) + entries[digest] = ShardEntry( + digest, shard, offset, length, rows, dimension, dtype + ) + if not entries: + raise ValueError("shard index contains no tensor entries") + for shard, ranges in intervals.items(): + previous_end = 0 + for start, end in sorted(ranges): + if start < previous_end: + raise ValueError(f"overlapping tensor ranges in {shard}") + previous_end = end + return ShardIndex(alignment, shards, entries) + + +def load_index(path: Path) -> ShardIndex: + with path.open(encoding="utf-8") as stream: + return parse_index(json.load(stream)) + + +def iter_index_entries(index: ShardIndex) -> Iterable[ShardEntry]: + return index.entries.values() diff --git a/services/tilemaxsim_triton.py b/services/tilemaxsim_triton.py new file mode 100644 index 00000000..8e9afce4 --- /dev/null +++ b/services/tilemaxsim_triton.py @@ -0,0 +1,124 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# Copyright (c) 2026 Hu Xinjing + +"""Fused ragged FP16 TileMaxSim over a process-owned GPU tensor arena.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ragged_tilemaxsim_fp16_kernel( + query, + documents, + document_offsets, + document_rows, + scores, + query_rows, + dimension: tl.constexpr, + max_document_rows: tl.constexpr, + block_query: tl.constexpr, + block_document: tl.constexpr, + block_dimension: tl.constexpr, +): + document_index = tl.program_id(0).to(tl.int64) + query_block_index = tl.program_id(1).to(tl.int64) + query_indices = (query_block_index * block_query + tl.arange(0, block_query)).to( + tl.int64 + ) + query_mask = query_indices < query_rows + document_base = tl.load(document_offsets + document_index).to(tl.int64) + valid_document_rows = tl.load(document_rows + document_index) + running_max = tl.full([block_query], value=float("-inf"), dtype=tl.float32) + + for document_start in range(0, max_document_rows, block_document): + document_indices = (document_start + tl.arange(0, block_document)).to(tl.int64) + document_mask = document_indices < valid_document_rows + similarities = tl.zeros([block_query, block_document], dtype=tl.float32) + for dimension_start in range(0, dimension, block_dimension): + dimension_indices = (dimension_start + tl.arange(0, block_dimension)).to( + tl.int64 + ) + dimension_mask = dimension_indices < dimension + query_pointers = ( + query + query_indices[:, None] * dimension + dimension_indices[None, :] + ) + query_tile = tl.load( + query_pointers, + mask=query_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + document_pointers = ( + documents + + document_base + + document_indices[:, None] * dimension + + dimension_indices[None, :] + ) + document_tile = tl.load( + document_pointers, + mask=document_mask[:, None] & dimension_mask[None, :], + other=0.0, + ) + similarities += tl.dot(query_tile, tl.trans(document_tile)) + similarities = tl.where(document_mask[None, :], similarities, float("-inf")) + running_max = tl.maximum(running_max, tl.max(similarities, axis=1)) + + running_max = tl.where(query_mask, running_max, 0.0) + tl.atomic_add(scores + document_index, tl.sum(running_max, axis=0)) + + +def ragged_tilemaxsim_fp16( + query: torch.Tensor, + document_arena: torch.Tensor, + document_offsets: torch.Tensor, + document_rows: torch.Tensor, + maximum_document_rows: int, +) -> torch.Tensor: + if query.device.type != "cuda" or document_arena.device != query.device: + raise ValueError("query and document arena must be on the same CUDA device") + if query.dtype != torch.float16 or document_arena.dtype != torch.float16: + raise ValueError("ragged TileMaxSim currently requires FP16 tensors") + if query.ndim != 2 or document_arena.ndim != 1: + raise ValueError("invalid query or document arena shape") + if document_offsets.dtype != torch.int64 or document_rows.dtype != torch.int32: + raise ValueError("invalid ragged TileMaxSim metadata dtype") + if document_offsets.shape != document_rows.shape: + raise ValueError("document offsets and rows must have the same shape") + count = document_offsets.numel() + if count == 0: + return torch.empty(0, dtype=torch.float32, device=query.device) + if maximum_document_rows <= 0: + raise ValueError("maximum document rows must be positive") + block_query = 32 + block_document = 32 + block_dimension = 128 + padded_document_rows = ( + (maximum_document_rows + block_document - 1) // block_document * block_document + ) + query_blocks = (query.shape[0] + block_query - 1) // block_query + scores = torch.zeros(count, dtype=torch.float32, device=query.device) + with torch.cuda.device(query.device): + _ragged_tilemaxsim_fp16_kernel[(count, query_blocks)]( + query, + document_arena, + document_offsets, + document_rows, + scores, + query.shape[0], + dimension=query.shape[1], + max_document_rows=padded_document_rows, + block_query=block_query, + block_document=block_document, + block_dimension=block_dimension, + ) + return scores diff --git a/services/tilemaxsimd/Cargo.lock b/services/tilemaxsimd/Cargo.lock new file mode 100644 index 00000000..c9c5a177 --- /dev/null +++ b/services/tilemaxsimd/Cargo.lock @@ -0,0 +1,416 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cc" +version = "1.2.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.185" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tilemaxsimd" +version = "0.1.0" +dependencies = [ + "anyhow", + "cc", + "clap", + "hex", + "libc", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "typenum" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.3+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/services/tilemaxsimd/Cargo.toml b/services/tilemaxsimd/Cargo.toml new file mode 100644 index 00000000..30c5d82a --- /dev/null +++ b/services/tilemaxsimd/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "tilemaxsimd" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +anyhow = "1.0" +clap = { version = "4.5", features = ["derive"] } +hex = "0.4" +libc = "0.2" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sha2 = "0.10" + +[build-dependencies] +cc = { version = "1.2", features = ["parallel"] } + +[workspace] diff --git a/services/tilemaxsimd/build.rs b/services/tilemaxsimd/build.rs new file mode 100644 index 00000000..a8cd9fff --- /dev/null +++ b/services/tilemaxsimd/build.rs @@ -0,0 +1,21 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +fn main() { + println!("cargo:rerun-if-changed=native/tilemaxsim_cuda.cu"); + cc::Build::new() + .cuda(true) + .flag("-O3") + .flag("-lineinfo") + .file("native/tilemaxsim_cuda.cu") + .compile("tilemaxsim_cuda"); + println!("cargo:rustc-link-lib=cudart"); + println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64"); +} diff --git a/services/tilemaxsimd/native/tilemaxsim_cuda.cu b/services/tilemaxsimd/native/tilemaxsim_cuda.cu new file mode 100644 index 00000000..6ebf470c --- /dev/null +++ b/services/tilemaxsimd/native/tilemaxsim_cuda.cu @@ -0,0 +1,353 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +struct VctmGpu { + int device; + unsigned char *allocation; + size_t total_bytes; + size_t tensor_bytes; + size_t workspace_bytes; + unsigned char *host_staging; + size_t host_staging_bytes; + cudaStream_t upload_stream; + cudaStream_t compute_stream; +}; + +static int fail(char *error, size_t capacity, const char *message) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s", message); + } + return 1; +} + +static int cuda_fail(char *error, size_t capacity, const char *operation, + cudaError_t status) { + if (error != nullptr && capacity != 0) { + std::snprintf(error, capacity, "%s: %s", operation, + cudaGetErrorString(status)); + } + return 1; +} + +extern "C" int vctm_gpu_create(int device, size_t total_bytes, + size_t workspace_bytes, VctmGpu **output, + char *error, size_t error_capacity) { + if (output == nullptr || total_bytes == 0 || workspace_bytes == 0 || + workspace_bytes >= total_bytes) { + return fail(error, error_capacity, "invalid GPU arena configuration"); + } + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + size_t free_bytes = 0; + size_t device_bytes = 0; + status = cudaMemGetInfo(&free_bytes, &device_bytes); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemGetInfo", status); + } + if (free_bytes < total_bytes) { + return fail(error, error_capacity, + "configured GPU memory is not currently available"); + } + auto *gpu = new VctmGpu{}; + gpu->device = device; + gpu->total_bytes = total_bytes; + gpu->tensor_bytes = ((total_bytes - workspace_bytes) / 256) * 256; + gpu->workspace_bytes = total_bytes - gpu->tensor_bytes; + gpu->host_staging_bytes = + std::min(gpu->tensor_bytes, static_cast(64) * 1024 * 1024); + status = cudaMalloc(reinterpret_cast(&gpu->allocation), total_bytes); + if (status != cudaSuccess) { + delete gpu; + return cuda_fail(error, error_capacity, "cudaMalloc", status); + } + if ((status = cudaStreamCreateWithFlags(&gpu->upload_stream, + cudaStreamNonBlocking)) != cudaSuccess || + (status = cudaStreamCreateWithFlags(&gpu->compute_stream, + cudaStreamNonBlocking)) != cudaSuccess) { + if (gpu->upload_stream != nullptr) cudaStreamDestroy(gpu->upload_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaStreamCreate", status); + } + status = cudaHostAlloc(reinterpret_cast(&gpu->host_staging), + gpu->host_staging_bytes, cudaHostAllocPortable); + if (status != cudaSuccess) { + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFree(gpu->allocation); + delete gpu; + return cuda_fail(error, error_capacity, "cudaHostAlloc", status); + } + std::memset(gpu->host_staging, 0, gpu->host_staging_bytes); + *output = gpu; + return 0; +} + +extern "C" void vctm_gpu_destroy(VctmGpu *gpu) { + if (gpu == nullptr) return; + cudaSetDevice(gpu->device); + cudaStreamDestroy(gpu->upload_stream); + cudaStreamDestroy(gpu->compute_stream); + cudaFreeHost(gpu->host_staging); + cudaFree(gpu->allocation); + delete gpu; +} + +extern "C" size_t vctm_gpu_tensor_bytes(const VctmGpu *gpu) { + return gpu == nullptr ? 0 : gpu->tensor_bytes; +} + +extern "C" int vctm_gpu_upload_batch( + VctmGpu *gpu, const uint64_t *offsets, + const unsigned char *const *payloads, const size_t *lengths, size_t count, + char *error, size_t error_capacity) { + if (gpu == nullptr || offsets == nullptr || payloads == nullptr || + lengths == nullptr) { + return fail(error, error_capacity, "invalid upload batch"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + for (size_t i = 0; i < count; ++i) { + if (offsets[i] > gpu->tensor_bytes || + lengths[i] > gpu->tensor_bytes - offsets[i]) { + return fail(error, error_capacity, "upload is outside the tensor arena"); + } + } + size_t item = 0; + size_t item_offset = 0; + while (item < count) { + size_t staging_offset = 0; + while (item < count && staging_offset < gpu->host_staging_bytes) { + const size_t remaining = lengths[item] - item_offset; + const size_t chunk = + std::min(remaining, gpu->host_staging_bytes - staging_offset); + std::memcpy(gpu->host_staging + staging_offset, + payloads[item] + item_offset, chunk); + status = cudaMemcpyAsync(gpu->allocation + offsets[item] + item_offset, + gpu->host_staging + staging_offset, chunk, + cudaMemcpyHostToDevice, gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaMemcpyAsync(H2D)", status); + } + staging_offset += chunk; + item_offset += chunk; + if (item_offset == lengths[item]) { + item += 1; + item_offset = 0; + } + } + status = cudaStreamSynchronize(gpu->upload_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaStreamSynchronize(upload)", + status); + } + } + return 0; +} + +template +__device__ float scalar_to_float(Scalar value); + +template <> +__device__ float scalar_to_float(half value) { + return __half2float(value); +} + +template <> +__device__ float scalar_to_float(float value) { + return value; +} + +template +__global__ void tilemaxsim_kernel(const Scalar *query, uint32_t query_rows, + uint32_t dimension, + const unsigned char *documents, + const uint64_t *document_offsets, + const uint32_t *document_rows, + size_t task_count, float *maxima) { + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t warps = blockDim.x >> 5; + __shared__ float warp_best[8]; + for (size_t task = blockIdx.x; task < task_count; task += gridDim.x) { + const size_t candidate = task / query_rows; + const uint32_t query_row = static_cast(task % query_rows); + const auto *document = reinterpret_cast( + documents + document_offsets[candidate]); + const Scalar *query_vector = + query + static_cast(query_row) * dimension; + float best = -CUDART_INF_F; + for (uint32_t row = warp; row < document_rows[candidate]; row += warps) { + const Scalar *document_vector = + document + static_cast(row) * dimension; + float dot = 0.0f; + for (uint32_t index = lane; index < dimension; index += 32) { + dot = fmaf(scalar_to_float(query_vector[index]), + scalar_to_float(document_vector[index]), dot); + } + for (int delta = 16; delta != 0; delta >>= 1) { + dot += __shfl_down_sync(0xffffffff, dot, delta); + } + if (lane == 0) best = fmaxf(best, dot); + } + if (lane == 0) warp_best[warp] = best; + __syncthreads(); + if (threadIdx.x == 0) { + float maximum = -CUDART_INF_F; + for (uint32_t index = 0; index < warps; ++index) { + maximum = fmaxf(maximum, warp_best[index]); + } + maxima[task] = maximum; + } + __syncthreads(); + } +} + +__global__ void tilemaxsim_sum_kernel(const float *maxima, uint32_t query_rows, + size_t count, float *scores) { + const size_t candidate = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (candidate >= count) return; + float score = 0.0f; + for (uint32_t query_row = 0; query_row < query_rows; ++query_row) { + score += maxima[candidate * query_rows + query_row]; + } + scores[candidate] = score; +} + +static size_t aligned(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +static bool reserve_aligned(size_t *cursor, size_t bytes, size_t *offset) { + constexpr size_t alignment = 256; + *offset = *cursor; + if (bytes > std::numeric_limits::max() - *cursor) return false; + const size_t end = *cursor + bytes; + if (end > std::numeric_limits::max() - (alignment - 1)) return false; + *cursor = aligned(end, alignment); + return true; +} + +extern "C" int vctm_gpu_score( + VctmGpu *gpu, const unsigned char *query, size_t query_bytes, + uint32_t query_rows, uint32_t dimension, uint8_t dtype, + const uint64_t *document_offsets, const uint32_t *document_rows, + size_t count, float *output, char *error, size_t error_capacity) { + if (gpu == nullptr || query == nullptr || document_offsets == nullptr || + document_rows == nullptr || output == nullptr || query_rows == 0 || + dimension == 0 || count == 0) { + return fail(error, error_capacity, "invalid TileMaxSim score request"); + } + cudaError_t status = cudaSetDevice(gpu->device); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "cudaSetDevice", status); + } + unsigned char *workspace = gpu->allocation + gpu->tensor_bytes; + if (count > std::numeric_limits::max() / query_rows) { + return fail(error, error_capacity, "TileMaxSim maxima count overflow"); + } + const size_t maxima_count = count * query_rows; + if (count > std::numeric_limits::max() / sizeof(uint64_t) || + count > std::numeric_limits::max() / sizeof(uint32_t) || + count > std::numeric_limits::max() / sizeof(float) || + maxima_count > std::numeric_limits::max() / sizeof(float)) { + return fail(error, error_capacity, "TileMaxSim workspace size overflow"); + } + size_t cursor = 0; + size_t query_offset = 0; + size_t offsets_offset = 0; + size_t rows_offset = 0; + size_t maxima_offset = 0; + size_t scores_offset = 0; + if (!reserve_aligned(&cursor, query_bytes, &query_offset) || + !reserve_aligned(&cursor, count * sizeof(uint64_t), &offsets_offset) || + !reserve_aligned(&cursor, count * sizeof(uint32_t), &rows_offset) || + !reserve_aligned(&cursor, maxima_count * sizeof(float), &maxima_offset) || + !reserve_aligned(&cursor, count * sizeof(float), &scores_offset)) { + return fail(error, error_capacity, "TileMaxSim workspace size overflow"); + } + if (cursor > gpu->workspace_bytes) { + return fail(error, error_capacity, + "TileMaxSim request exceeds the configured GPU workspace"); + } + status = cudaMemcpyAsync(workspace + query_offset, query, query_bytes, + cudaMemcpyHostToDevice, gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + offsets_offset, document_offsets, + count * sizeof(uint64_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status == cudaSuccess) + status = cudaMemcpyAsync(workspace + rows_offset, document_rows, + count * sizeof(uint32_t), cudaMemcpyHostToDevice, + gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "CUDA workspace initialization", + status); + } + // A CUDA grid's Y dimension is limited to 65,535 even on modern devices. + // Flatten candidate/query-row work into X and let blocks stride when the + // task count is larger. This covers every protocol-valid query shape without + // constructing an excessive launch grid. + constexpr size_t maximum_kernel_blocks = 65'535; + const size_t kernel_blocks = std::min(maxima_count, maximum_kernel_blocks); + dim3 grid(static_cast(kernel_blocks)); + dim3 block(256); + if (dtype == 2) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + maxima_count, reinterpret_cast(workspace + maxima_offset)); + } else if (dtype == 1) { + tilemaxsim_kernel<<compute_stream>>>( + reinterpret_cast(workspace + query_offset), query_rows, + dimension, gpu->allocation, + reinterpret_cast(workspace + offsets_offset), + reinterpret_cast(workspace + rows_offset), + maxima_count, reinterpret_cast(workspace + maxima_offset)); + } else { + return fail(error, error_capacity, "unsupported tensor dtype"); + } + status = cudaGetLastError(); + if (status == cudaSuccess) { + constexpr unsigned int threads = 256; + const auto blocks = static_cast((count + threads - 1) / threads); + tilemaxsim_sum_kernel<<compute_stream>>>( + reinterpret_cast(workspace + maxima_offset), query_rows, + count, reinterpret_cast(workspace + scores_offset)); + status = cudaGetLastError(); + } + if (status == cudaSuccess) + status = cudaMemcpyAsync(output, workspace + scores_offset, + count * sizeof(float), cudaMemcpyDeviceToHost, + gpu->compute_stream); + if (status == cudaSuccess) status = cudaStreamSynchronize(gpu->compute_stream); + if (status != cudaSuccess) { + return cuda_fail(error, error_capacity, "TileMaxSim CUDA execution", status); + } + return 0; +} diff --git a/services/tilemaxsimd/src/bin/tilemaxsimctl.rs b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs new file mode 100644 index 00000000..cb8d4140 --- /dev/null +++ b/services/tilemaxsimd/src/bin/tilemaxsimctl.rs @@ -0,0 +1,546 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use anyhow::{Context, Result, anyhow, bail}; +use clap::{Parser, Subcommand}; +use sha2::{Digest, Sha256}; +use std::collections::HashSet; +use std::fs::{self, File, OpenOptions}; +use std::io::{BufRead, BufReader, Read, Write}; +use std::os::fd::AsRawFd; +use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; +use std::os::unix::net::UnixStream; +use std::path::{Path, PathBuf}; +use std::thread; +use std::time::{Duration, Instant, SystemTime}; + +#[derive(Parser)] +#[command(about = "Probe the native TileMaxSim daemon readiness socket")] +struct Args { + #[command(subcommand)] + command: Option, + #[arg(long, default_value = "/run/vectorchord/tilemaxsim-status.sock")] + socket: PathBuf, + #[arg(long, default_value_t = 500)] + io_timeout_ms: u64, + #[arg(long, default_value_t = 0)] + wait_timeout_ms: u64, + #[arg(long, default_value = "ready", value_parser = ["ready", "live"])] + probe: String, +} + +#[derive(Subcommand)] +enum Command { + /// Publish one canonical tensor from stdin into the immutable object store. + PublishObject { + #[arg(long)] + root: PathBuf, + #[arg(long)] + rows: u32, + #[arg(long)] + dimension: u32, + #[arg(long, value_parser = ["float16", "float32"])] + dtype: String, + #[arg(long)] + expected_sha256: Option, + }, + /// Remove old immutable objects not named by the sha256:// refs on stdin. + GcObjects { + #[arg(long)] + root: PathBuf, + #[arg(long, default_value_t = 86_400)] + grace_seconds: u64, + /// Actually remove eligible objects. The default is a dry run. + #[arg(long)] + delete: bool, + }, +} + +fn main() -> Result<()> { + let args = Args::parse(); + if let Some(Command::PublishObject { + root, + rows, + dimension, + dtype, + expected_sha256, + }) = args.command.as_ref() + { + let expected_bytes = tensor_bytes(*rows, *dimension, dtype)?; + let mut payload = Vec::with_capacity(expected_bytes); + std::io::stdin() + .take(expected_bytes.saturating_add(1) as u64) + .read_to_end(&mut payload)?; + if payload.len() != expected_bytes { + bail!( + "stdin tensor length {} does not match declared length {expected_bytes}", + payload.len() + ); + } + let descriptor = publish_object( + root, + *rows, + *dimension, + dtype, + &payload, + expected_sha256.as_deref(), + )?; + println!("{}", serde_json::to_string(&descriptor)?); + return Ok(()); + } + if let Some(Command::GcObjects { + root, + grace_seconds, + delete, + }) = args.command.as_ref() + { + let live = read_live_refs(BufReader::new(std::io::stdin().lock()))?; + let outcome = gc_objects(root, &live, *grace_seconds, *delete)?; + println!("{}", serde_json::to_string(&outcome)?); + return Ok(()); + } + if args.io_timeout_ms == 0 { + bail!("I/O timeout must be positive"); + } + let io_timeout = Duration::from_millis(args.io_timeout_ms); + let wait_timeout = Duration::from_millis(args.wait_timeout_ms); + let deadline = Instant::now() + .checked_add(wait_timeout) + .ok_or_else(|| anyhow!("readiness deadline overflow"))?; + + loop { + let error = match probe(&args.socket, io_timeout, &args.probe) { + Ok(()) => return Ok(()), + Err(error) => error, + }; + if wait_timeout.is_zero() || Instant::now() >= deadline { + return Err(error); + } + thread::sleep(Duration::from_millis(50)); + } +} + +#[derive(serde::Serialize)] +struct PublishedDescriptor { + tensor_ref: String, + tensor_rows: u32, + tensor_dim: u32, + tensor_dtype: String, + tensor_checksum: String, +} + +#[derive(serde::Serialize)] +struct GcOutcome { + live_refs: usize, + scanned_objects: u64, + eligible_objects: u64, + eligible_bytes: u64, + deleted_objects: u64, + deleted_bytes: u64, + skipped_entries: u64, + dry_run: bool, +} + +struct ObjectStoreLock { + _file: File, +} + +fn prepare_store_root(root: &Path) -> Result<()> { + fs::create_dir_all(root)?; + let metadata = fs::symlink_metadata(root)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + bail!("tensor root must be a real directory"); + } + Ok(()) +} + +fn lock_object_store(root: &Path, exclusive: bool) -> Result { + prepare_store_root(root)?; + let path = root.join(".tilemaxsim-objects.lock"); + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .mode(0o640) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(&path) + .with_context(|| format!("cannot open object-store lock {}", path.display()))?; + let operation = if exclusive { + libc::LOCK_EX + } else { + libc::LOCK_SH + }; + loop { + if unsafe { libc::flock(file.as_raw_fd(), operation) } == 0 { + break; + } + let error = std::io::Error::last_os_error(); + if error.kind() != std::io::ErrorKind::Interrupted { + return Err(error) + .with_context(|| format!("cannot lock tensor object store {}", root.display())); + } + } + Ok(ObjectStoreLock { _file: file }) +} + +fn read_live_refs(reader: impl BufRead) -> Result> { + let mut live = HashSet::new(); + for (index, line) in reader.lines().enumerate() { + let line = line?; + let value = line.trim(); + if value.is_empty() { + continue; + } + let Some(digest) = value.strip_prefix("sha256://") else { + bail!("live ref on line {} is not a sha256:// URI", index + 1); + }; + if !is_lower_hex_digest(digest) { + bail!("live ref on line {} has an invalid digest", index + 1); + } + live.insert(digest.to_owned()); + } + Ok(live) +} + +fn is_lower_hex_digest(value: &str) -> bool { + value.len() == 64 + && value + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} + +fn gc_objects( + root: &Path, + live: &HashSet, + grace_seconds: u64, + delete: bool, +) -> Result { + let _lock = lock_object_store(root, true)?; + let mut outcome = GcOutcome { + live_refs: live.len(), + scanned_objects: 0, + eligible_objects: 0, + eligible_bytes: 0, + deleted_objects: 0, + deleted_bytes: 0, + skipped_entries: 0, + dry_run: !delete, + }; + let objects = root.join("objects"); + if !objects.try_exists()? { + return Ok(outcome); + } + let objects_metadata = fs::symlink_metadata(&objects)?; + if objects_metadata.file_type().is_symlink() || !objects_metadata.is_dir() { + bail!("tensor objects path must be a real directory"); + } + + let now = SystemTime::now(); + let grace = Duration::from_secs(grace_seconds); + for directory in fs::read_dir(&objects)? { + let directory = directory?; + let directory_metadata = fs::symlink_metadata(directory.path())?; + if directory_metadata.file_type().is_symlink() { + bail!("tensor objects path contains a symlink"); + } + let Some(prefix) = directory.file_name().to_str().map(str::to_owned) else { + outcome.skipped_entries += 1; + continue; + }; + if !directory_metadata.is_dir() + || prefix.len() != 2 + || !prefix + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + outcome.skipped_entries += 1; + continue; + } + for object in fs::read_dir(directory.path())? { + let object = object?; + let path = object.path(); + let metadata = fs::symlink_metadata(&path)?; + if metadata.file_type().is_symlink() { + bail!("tensor object path contains a symlink"); + } + let Some(filename) = object.file_name().to_str().map(str::to_owned) else { + outcome.skipped_entries += 1; + continue; + }; + let Some(digest) = filename.strip_suffix(".tensor") else { + outcome.skipped_entries += 1; + continue; + }; + if !metadata.is_file() || !is_lower_hex_digest(digest) || !digest.starts_with(&prefix) { + outcome.skipped_entries += 1; + continue; + } + outcome.scanned_objects += 1; + if live.contains(digest) { + continue; + } + let modified = metadata.modified().with_context(|| { + format!("cannot read tensor object timestamp {}", path.display()) + })?; + if now.duration_since(modified).unwrap_or_default() < grace { + continue; + } + outcome.eligible_objects += 1; + outcome.eligible_bytes = outcome.eligible_bytes.saturating_add(metadata.len()); + if delete { + fs::remove_file(&path) + .with_context(|| format!("cannot remove tensor object {}", path.display()))?; + outcome.deleted_objects += 1; + outcome.deleted_bytes = outcome.deleted_bytes.saturating_add(metadata.len()); + } + } + } + Ok(outcome) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: &str) -> Result { + if rows == 0 || rows > 65_536 || dimension == 0 || dimension > 60_000 { + bail!("invalid tensor shape"); + } + let scalar_bytes = match dtype { + "float16" => 2usize, + "float32" => 4usize, + _ => bail!("unsupported tensor dtype"), + }; + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or_else(|| anyhow!("tensor shape overflow")) +} + +fn publish_object( + root: &Path, + rows: u32, + dimension: u32, + dtype: &str, + payload: &[u8], + expected_sha256: Option<&str>, +) -> Result { + if payload.len() != tensor_bytes(rows, dimension, dtype)? { + bail!("tensor payload length disagrees with its shape"); + } + let digest = hex::encode(Sha256::digest(payload)); + if expected_sha256.is_some_and(|expected| expected != digest) { + bail!("tensor payload checksum disagrees with --expected-sha256"); + } + let _lock = lock_object_store(root, false)?; + let objects = root.join("objects"); + fs::create_dir_all(&objects)?; + let objects_metadata = fs::symlink_metadata(&objects)?; + if objects_metadata.file_type().is_symlink() || !objects_metadata.is_dir() { + bail!("tensor objects path must be a real directory"); + } + let directory = objects.join(&digest[..2]); + fs::create_dir_all(&directory)?; + let directory_metadata = fs::symlink_metadata(&directory)?; + if directory_metadata.file_type().is_symlink() || !directory_metadata.is_dir() { + bail!("tensor object directory must be a real directory"); + } + let destination = directory.join(format!("{digest}.tensor")); + if destination.try_exists()? { + verify_existing_object(&destination, payload.len(), &digest)?; + } else { + let temporary = directory.join(format!(".{digest}.{}.tmp", std::process::id())); + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .mode(0o640) + .open(&temporary) + .with_context(|| format!("cannot create {}", temporary.display()))?; + let publish = (|| -> Result<()> { + file.write_all(payload)?; + file.sync_all()?; + fs::set_permissions(&temporary, fs::Permissions::from_mode(0o440))?; + match fs::hard_link(&temporary, &destination) { + Ok(()) => File::open(&directory)?.sync_all()?, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + verify_existing_object(&destination, payload.len(), &digest)?; + } + Err(error) => return Err(error.into()), + } + Ok(()) + })(); + drop(file); + let remove_result = fs::remove_file(&temporary); + if let Err(error) = remove_result + && error.kind() != std::io::ErrorKind::NotFound + { + return Err(error.into()); + } + publish?; + } + Ok(PublishedDescriptor { + tensor_ref: format!("sha256://{digest}"), + tensor_rows: rows, + tensor_dim: dimension, + tensor_dtype: dtype.to_owned(), + tensor_checksum: format!("sha256:{digest}"), + }) +} + +fn verify_existing_object(path: &Path, expected_bytes: usize, expected_digest: &str) -> Result<()> { + let metadata = fs::symlink_metadata(path)?; + if metadata.file_type().is_symlink() + || !metadata.is_file() + || metadata.len() != expected_bytes as u64 + { + bail!("existing immutable tensor object is invalid"); + } + let mut file = OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path)?; + let mut digest = Sha256::new(); + let mut buffer = [0_u8; 1024 * 1024]; + loop { + let count = file.read(&mut buffer)?; + if count == 0 { + break; + } + digest.update(&buffer[..count]); + } + if hex::encode(digest.finalize()) != expected_digest { + bail!("existing immutable tensor object checksum mismatch"); + } + Ok(()) +} + +fn probe(socket: &PathBuf, timeout: Duration, kind: &str) -> Result<()> { + let mut stream = UnixStream::connect(socket) + .with_context(|| format!("cannot connect to status socket {}", socket.display()))?; + stream.set_read_timeout(Some(timeout))?; + stream.set_write_timeout(Some(timeout))?; + let path = if kind == "live" { "/livez" } else { "/healthz" }; + write!( + stream, + "GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" + )?; + let mut response = Vec::new(); + stream.read_to_end(&mut response)?; + if !is_probe_response(&response, kind) { + bail!("TileMaxSim daemon failed its {kind} probe"); + } + Ok(()) +} + +fn is_probe_response(response: &[u8], kind: &str) -> bool { + let expected_body = if kind == "live" { + b"{\"live\":true}".as_slice() + } else { + b"{\"ready\":true}".as_slice() + }; + response.starts_with(b"HTTP/1.1 200 ") + && response + .windows(b"\r\n\r\n".len()) + .any(|window| window == b"\r\n\r\n") + && response.ends_with(expected_body) +} + +#[cfg(test)] +mod tests { + use super::{gc_objects, is_probe_response, publish_object, read_live_refs}; + use std::collections::HashSet; + use std::fs; + use std::io::Cursor; + + #[test] + fn readiness_requires_success_status_and_true_body() { + assert!(is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}", + "ready" + )); + assert!(!is_probe_response( + b"HTTP/1.1 503 Service Unavailable\r\n\r\n{\"ready\":false}", + "ready" + )); + assert!(!is_probe_response(b"HTTP/1.1 200 OK\r\n\r\n", "ready")); + } + + #[test] + fn liveness_is_distinct_from_readiness() { + assert!(is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 13\r\n\r\n{\"live\":true}", + "live", + )); + assert!(!is_probe_response( + b"HTTP/1.1 200 OK\r\nContent-Length: 14\r\n\r\n{\"ready\":true}", + "live", + )); + } + + #[test] + fn publish_object_is_content_addressed_and_idempotent() { + let root = std::env::temp_dir().join(format!( + "tilemaxsimctl-publish-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let _ = fs::remove_dir_all(&root); + let payload = [0_u8, 60, 0, 0]; + let first = publish_object(&root, 1, 2, "float16", &payload, None).unwrap(); + let second = publish_object(&root, 1, 2, "float16", &payload, None).unwrap(); + assert_eq!(first.tensor_ref, second.tensor_ref); + assert_eq!(first.tensor_checksum, second.tensor_checksum); + assert!( + root.join("objects") + .join(&first.tensor_checksum[7..9]) + .exists() + ); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn gc_requires_valid_refs_and_preserves_live_objects() { + let root = std::env::temp_dir().join(format!( + "tilemaxsimctl-gc-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let _ = fs::remove_dir_all(&root); + let first = publish_object(&root, 1, 2, "float16", &[0, 60, 0, 0], None).unwrap(); + let second = publish_object(&root, 1, 2, "float16", &[0, 56, 0, 0], None).unwrap(); + let live = read_live_refs(Cursor::new(format!("{}\n", first.tensor_ref))).unwrap(); + let dry_run = gc_objects(&root, &live, 0, false).unwrap(); + assert_eq!(dry_run.scanned_objects, 2); + assert_eq!(dry_run.eligible_objects, 1); + assert_eq!(dry_run.deleted_objects, 0); + + let deleted = gc_objects(&root, &live, 0, true).unwrap(); + assert_eq!(deleted.deleted_objects, 1); + let first_digest = &first.tensor_checksum[7..]; + let second_digest = &second.tensor_checksum[7..]; + assert!( + root.join("objects") + .join(&first_digest[..2]) + .join(format!("{first_digest}.tensor")) + .exists() + ); + assert!( + !root + .join("objects") + .join(&second_digest[..2]) + .join(format!("{second_digest}.tensor")) + .exists() + ); + assert!(read_live_refs(Cursor::new("not-a-ref\n")).is_err()); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn empty_live_manifest_is_explicitly_supported_for_an_empty_database() { + let live = read_live_refs(Cursor::new("\n")).unwrap(); + assert_eq!(live, HashSet::new()); + } +} diff --git a/services/tilemaxsimd/src/cache.rs b/services/tilemaxsimd/src/cache.rs new file mode 100644 index 00000000..e53f1b49 --- /dev/null +++ b/services/tilemaxsimd/src/cache.rs @@ -0,0 +1,772 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::hash::{Hash, Hasher}; + +#[derive(Debug)] +pub struct PageRunAllocator { + block_bytes: usize, + block_count: usize, + free_by_start: BTreeMap, + free_by_size: BTreeMap>, + allocated: HashMap, + free_blocks: usize, +} + +impl PageRunAllocator { + pub fn new(capacity: usize, block_bytes: usize) -> Result { + if capacity == 0 || block_bytes == 0 || !block_bytes.is_multiple_of(256) { + return Err("invalid fixed-block arena"); + } + let block_count = capacity / block_bytes; + if block_count == 0 { + return Err("fixed-block arena is too small"); + } + let mut allocator = Self { + block_bytes, + block_count, + free_by_start: BTreeMap::new(), + free_by_size: BTreeMap::new(), + allocated: HashMap::new(), + free_blocks: 0, + }; + allocator.add_free_run(0, block_count)?; + Ok(allocator) + } + + fn add_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if blocks == 0 + || start + .checked_add(blocks) + .is_none_or(|end| end > self.block_count) + { + return Err("free GPU page run is outside the arena"); + } + if self.free_by_start.insert(start, blocks).is_some() { + return Err("duplicate free GPU page run"); + } + self.free_by_size.entry(blocks).or_default().insert(start); + self.free_blocks = self + .free_blocks + .checked_add(blocks) + .ok_or("free GPU page count overflow")?; + Ok(()) + } + + fn remove_free_run(&mut self, start: usize, blocks: usize) -> Result<(), &'static str> { + if self.free_by_start.remove(&start) != Some(blocks) { + return Err("free GPU address index is inconsistent"); + } + let remove_size = { + let starts = self + .free_by_size + .get_mut(&blocks) + .ok_or("free GPU size index is inconsistent")?; + if !starts.remove(&start) { + return Err("free GPU size index is inconsistent"); + } + starts.is_empty() + }; + if remove_size { + self.free_by_size.remove(&blocks); + } + self.free_blocks = self + .free_blocks + .checked_sub(blocks) + .ok_or("free GPU page count underflow")?; + Ok(()) + } + + pub fn capacity(&self) -> usize { + self.block_count * self.block_bytes + } + + pub fn block_bytes(&self) -> usize { + self.block_bytes + } + + pub fn allocation_bytes(&self, payload_bytes: usize) -> Option { + if payload_bytes == 0 { + return None; + } + payload_bytes + .div_ceil(self.block_bytes) + .checked_mul(self.block_bytes) + } + + pub fn largest_free(&self) -> usize { + self.free_by_size + .last_key_value() + .map_or(0, |(blocks, _)| blocks * self.block_bytes) + } + + pub fn allocate(&mut self, payload_bytes: usize) -> Option<(usize, usize)> { + let allocation_bytes = self.allocation_bytes(payload_bytes)?; + let required = allocation_bytes / self.block_bytes; + let (&available, starts) = self.free_by_size.range(required..).next()?; + let start = *starts.first()?; + self.remove_free_run(start, available).ok()?; + if available > required { + self.add_free_run(start + required, available - required) + .ok()?; + } + self.allocated.insert(start, required); + Some((start * self.block_bytes, allocation_bytes)) + } + + pub fn release(&mut self, offset: usize) -> Result<(), &'static str> { + if !offset.is_multiple_of(self.block_bytes) { + return Err("unaligned fixed-block release"); + } + let mut start = offset / self.block_bytes; + let mut blocks = self + .allocated + .remove(&start) + .ok_or("fixed-block page run was released twice")?; + let previous = self + .free_by_start + .range(..start) + .next_back() + .map(|(&previous_start, &previous_blocks)| (previous_start, previous_blocks)); + if let Some((previous_start, previous_blocks)) = previous + && previous_start + previous_blocks == start + { + self.remove_free_run(previous_start, previous_blocks)?; + start = previous_start; + blocks += previous_blocks; + } + let next_start = start + blocks; + if let Some(&next_blocks) = self.free_by_start.get(&next_start) { + self.remove_free_run(next_start, next_blocks)?; + blocks += next_blocks; + } + self.add_free_run(start, blocks) + } + + pub fn free_bytes(&self) -> usize { + self.free_blocks * self.block_bytes + } + + pub fn allocated_bytes(&self) -> usize { + self.allocated + .values() + .map(|blocks| blocks * self.block_bytes) + .sum() + } + + #[cfg(test)] + pub fn validate(&self) -> Result<(), &'static str> { + let address_blocks = self.free_by_start.values().sum::(); + let size_blocks = self + .free_by_size + .iter() + .map(|(blocks, starts)| blocks * starts.len()) + .sum::(); + if address_blocks != self.free_blocks || size_blocks != self.free_blocks { + return Err("free GPU page accounting is inconsistent"); + } + for (&start, &blocks) in &self.free_by_start { + if !self + .free_by_size + .get(&blocks) + .is_some_and(|starts| starts.contains(&start)) + { + return Err("free GPU page indexes disagree"); + } + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct TinyLfu { + width: usize, + tables: Vec>, + samples: usize, +} + +impl TinyLfu { + pub fn new(width: usize) -> Self { + Self { + width, + tables: vec![vec![0; width]; 4], + samples: 0, + } + } + + fn indices(&self, key: &T) -> [usize; 4] { + std::array::from_fn(|row| { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + row.hash(&mut hasher); + key.hash(&mut hasher); + hasher.finish() as usize % self.width + }) + } + + pub fn increment(&mut self, key: &T) -> u16 { + let indices = self.indices(key); + for (row, index) in indices.into_iter().enumerate() { + self.tables[row][index] = self.tables[row][index].saturating_add(1); + } + self.samples += 1; + let estimate = self.estimate(key).max(1); + if self.samples >= self.width * 10 { + for table in &mut self.tables { + for value in table { + *value /= 2; + } + } + self.samples /= 2; + } + estimate + } + + pub fn estimate(&self, key: &T) -> u16 { + self.indices(key) + .into_iter() + .enumerate() + .map(|(row, index)| self.tables[row][index]) + .min() + .unwrap_or(0) + } +} + +#[derive(Clone, Debug)] +pub struct CacheEntry { + pub offset: u64, + pub allocated_bytes: usize, + pub payload_bytes: usize, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, + pub references: usize, + pub pinned: bool, + /// False while an H2D upload is in progress. Loading entries reserve pages + /// but are never visible as cache hits. + pub ready: bool, + owner_tenant: String, + priority: f64, +} + +#[derive(Debug, PartialEq)] +pub enum Admission { + Admitted { offset: u64, allocated_bytes: usize }, + Rejected, + Deferred, +} + +pub struct GpuCache { + allocator: PageRunAllocator, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + tenant_allocated: HashMap, + tenant_reservations: HashMap, + default_tenant_max_bytes: usize, + pinned_max_bytes: usize, + pinned_bytes: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl GpuCache { + #[cfg(test)] + pub fn new(capacity: usize, block_bytes: usize) -> Result { + Self::new_with_limits(capacity, block_bytes, 100, 100, HashMap::new()) + } + + pub fn new_with_limits( + capacity: usize, + block_bytes: usize, + tenant_max_percent: u8, + pinned_max_percent: u8, + tenant_reservations: HashMap, + ) -> Result { + if tenant_max_percent == 0 + || tenant_max_percent > 100 + || pinned_max_percent > 100 + || tenant_reservations.values().sum::() > capacity + { + return Err("invalid GPU tenant cache policy"); + } + Ok(Self { + allocator: PageRunAllocator::new(capacity, block_bytes)?, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + tenant_allocated: HashMap::new(), + tenant_reservations, + default_tenant_max_bytes: capacity * tenant_max_percent as usize / 100, + pinned_max_bytes: capacity * pinned_max_percent as usize / 100, + pinned_bytes: 0, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + }) + } + + pub fn capacity(&self) -> usize { + self.allocator.capacity() + } + + pub fn block_bytes(&self) -> usize { + self.allocator.block_bytes() + } + + pub fn entry_count(&self) -> usize { + self.entries.len() + } + + pub fn pinned_entries(&self) -> usize { + self.entries.values().filter(|entry| entry.pinned).count() + } + + pub fn tenant_count(&self) -> usize { + self.tenant_allocated.len() + } + + pub fn pinned_bytes(&self) -> usize { + self.pinned_bytes + } + + pub fn free_bytes(&self) -> usize { + self.allocator.free_bytes() + } + + pub fn largest_free_extent(&self) -> usize { + self.allocator.largest_free() + } + + pub fn allocated_bytes(&self) -> usize { + self.allocator.allocated_bytes() + } + + pub fn payload_bytes(&self) -> usize { + self.entries.values().map(|entry| entry.payload_bytes).sum() + } + + pub fn get(&mut self, key: &str) -> Option { + let frequency = self.sketch.increment(&key); + let Some(entry) = self.entries.get_mut(key).filter(|entry| entry.ready) else { + self.misses += 1; + return None; + }; + entry.references += 1; + entry.priority = self.inflation + f64::from(frequency) / entry.allocated_bytes as f64; + self.hits += 1; + Some(entry.clone()) + } + + pub fn acquire_existing(&mut self, key: &str) -> Option { + let entry = self.entries.get_mut(key).filter(|entry| entry.ready)?; + entry.references += 1; + Some(entry.clone()) + } + + pub fn record_access_miss(&mut self, key: &str) { + self.sketch.increment(&key); + self.misses += 1; + } + + pub fn contains(&self, key: &str) -> bool { + self.entries.get(key).is_some_and(|entry| entry.ready) + } + + fn victim(&self) -> Option<(&String, &CacheEntry)> { + self.entries + .iter() + .filter(|(_, entry)| self.can_evict(entry)) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + } + + fn tenant_victim(&self, tenant: &str) -> Option<(&String, &CacheEntry)> { + self.entries + .iter() + .filter(|(_, entry)| { + entry.owner_tenant == tenant && entry.references == 0 && !entry.pinned + }) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + } + + fn can_evict(&self, entry: &CacheEntry) -> bool { + if entry.references != 0 || entry.pinned { + return false; + } + let usage = self + .tenant_allocated + .get(&entry.owner_tenant) + .copied() + .unwrap_or(0); + let reservation = self + .tenant_reservations + .get(&entry.owner_tenant) + .copied() + .unwrap_or(0); + usage.saturating_sub(entry.allocated_bytes) >= reservation + } + + fn remove_entry(&mut self, key: &str) -> Result { + let entry = self + .entries + .remove(key) + .ok_or("GPU cache entry disappeared")?; + self.allocator.release(entry.offset as usize)?; + if entry.pinned { + self.pinned_bytes = self.pinned_bytes.saturating_sub(entry.allocated_bytes); + } + let remove_tenant = if let Some(usage) = self.tenant_allocated.get_mut(&entry.owner_tenant) + { + *usage = usage.saturating_sub(entry.allocated_bytes); + *usage == 0 + } else { + false + }; + if remove_tenant { + self.tenant_allocated.remove(&entry.owner_tenant); + } + Ok(entry) + } + + fn evict_one(&mut self) -> bool { + let Some(key) = self.victim().map(|(key, _)| key.clone()) else { + return false; + }; + let entry = self.remove_entry(&key).expect("victim disappeared"); + self.inflation = self.inflation.max(entry.priority); + self.evictions += 1; + true + } + + #[allow(clippy::too_many_arguments)] + #[cfg(test)] + pub fn admit( + &mut self, + key: String, + payload_bytes: usize, + rows: u32, + dimension: u32, + dtype: u8, + pinned: bool, + force: bool, + ) -> Admission { + self.admit_for_tenant( + "__default__", + key, + payload_bytes, + rows, + dimension, + dtype, + pinned, + force, + ) + } + + #[allow(clippy::too_many_arguments)] + pub fn admit_for_tenant( + &mut self, + tenant: &str, + key: String, + payload_bytes: usize, + rows: u32, + dimension: u32, + dtype: u8, + pinned: bool, + force: bool, + ) -> Admission { + let Some(allocation_bytes) = self.allocator.allocation_bytes(payload_bytes) else { + return Admission::Deferred; + }; + if allocation_bytes > self.allocator.capacity() { + return Admission::Deferred; + } + if pinned && self.pinned_bytes.saturating_add(allocation_bytes) > self.pinned_max_bytes { + return Admission::Deferred; + } + if !pinned { + let tenant_max_bytes = self + .tenant_reservations + .get(tenant) + .copied() + .unwrap_or(0) + .max(self.default_tenant_max_bytes); + while self + .tenant_allocated + .get(tenant) + .copied() + .unwrap_or(0) + .saturating_add(allocation_bytes) + > tenant_max_bytes + { + let Some((victim_key, victim_priority)) = self + .tenant_victim(tenant) + .map(|(victim_key, victim)| (victim_key.clone(), victim.priority)) + else { + return Admission::Deferred; + }; + let candidate = self.inflation + + f64::from(self.sketch.estimate(&key).max(1)) / allocation_bytes as f64; + if !force && candidate <= victim_priority { + self.admission_rejections += 1; + return Admission::Rejected; + } + let entry = self + .remove_entry(&victim_key) + .expect("tenant victim disappeared"); + self.inflation = self.inflation.max(entry.priority); + self.evictions += 1; + } + } + while self.allocator.largest_free() < allocation_bytes { + let Some((_, victim)) = self.victim() else { + return Admission::Deferred; + }; + let candidate = self.inflation + + f64::from(self.sketch.estimate(&key).max(1)) / allocation_bytes as f64; + if !force && !pinned && candidate <= victim.priority { + self.admission_rejections += 1; + return Admission::Rejected; + } + if !self.evict_one() { + return Admission::Deferred; + } + } + let Some((offset, allocated_bytes)) = self.allocator.allocate(payload_bytes) else { + return Admission::Deferred; + }; + let priority = + self.inflation + f64::from(self.sketch.estimate(&key).max(1)) / allocated_bytes as f64; + self.entries.insert( + key, + CacheEntry { + offset: offset as u64, + allocated_bytes, + payload_bytes, + rows, + dimension, + dtype, + references: 1, + pinned, + ready: false, + owner_tenant: tenant.to_owned(), + priority, + }, + ); + *self.tenant_allocated.entry(tenant.to_owned()).or_default() += allocated_bytes; + if pinned { + self.pinned_bytes += allocated_bytes; + } + Admission::Admitted { + offset: offset as u64, + allocated_bytes, + } + } + + pub fn release(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .get_mut(key) + .ok_or("GPU cache entry disappeared")?; + if entry.references == 0 { + return Err("GPU cache entry was released twice"); + } + entry.references -= 1; + Ok(()) + } + + pub fn mark_ready(&mut self, key: &str) -> Result<(), &'static str> { + let entry = self + .entries + .get_mut(key) + .ok_or("GPU cache entry disappeared before upload completed")?; + entry.ready = true; + Ok(()) + } + + #[cfg(test)] + pub fn entry(&self, key: &str) -> Option<&CacheEntry> { + self.entries.get(key) + } + + pub fn remove(&mut self, key: &str) -> Result<(), &'static str> { + self.remove_entry(key).map(|_| ()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn next_test_random(state: &mut u64) -> u64 { + // Deterministic xorshift64 sequence for allocator churn tests. This is + // test data, not a pointer, address, secret, or production RNG. + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + + #[test] + fn page_runs_use_exact_blocks_and_coalesce() { + let mut allocator = PageRunAllocator::new(8 * 256, 256).unwrap(); + let first = allocator.allocate(300).unwrap(); + let second = allocator.allocate(300).unwrap(); + let third = allocator.allocate(700).unwrap(); + assert_eq!(first, (0, 512)); + assert_eq!(second, (512, 512)); + assert_eq!(third, (1024, 768)); + assert_eq!(allocator.free_bytes(), 256); + allocator.release(second.0).unwrap(); + allocator.release(first.0).unwrap(); + assert_eq!(allocator.largest_free(), 1024); + allocator.release(third.0).unwrap(); + assert_eq!(allocator.largest_free(), 2048); + allocator.validate().unwrap(); + } + + #[test] + fn page_run_indexes_stay_consistent_under_churn() { + let mut allocator = PageRunAllocator::new(64 * 256, 256).unwrap(); + let mut live = Vec::new(); + let mut state = 0x5eed_u64; + for _ in 0..20_000 { + let random = next_test_random(&mut state); + if !live.is_empty() && random & 3 == 0 { + let index = random as usize % live.len(); + let (offset, _) = live.swap_remove(index); + allocator.release(offset).unwrap(); + } else { + let payload = ((random >> 16) as usize % (12 * 256)) + 1; + if let Some(allocation) = allocator.allocate(payload) { + live.push(allocation); + } + } + allocator.validate().unwrap(); + assert_eq!( + allocator.free_bytes() + allocator.allocated_bytes(), + allocator.capacity() + ); + } + } + + #[test] + fn tinylfu_protects_hot_entry() { + let mut cache = GpuCache::new(256, 256).unwrap(); + cache.record_access_miss("hot"); + assert!(matches!( + cache.admit("hot".into(), 64, 1, 32, 2, false, false), + Admission::Admitted { .. } + )); + assert!( + cache.get("hot").is_none(), + "loading pages must not be visible" + ); + cache.mark_ready("hot").unwrap(); + cache.release("hot").unwrap(); + for _ in 0..3 { + cache.get("hot").unwrap(); + cache.release("hot").unwrap(); + } + cache.record_access_miss("cold"); + assert_eq!( + cache.admit("cold".into(), 64, 1, 32, 2, false, false), + Admission::Rejected + ); + assert!(cache.entry("hot").is_some()); + } + + #[test] + fn failed_loading_entry_can_be_rolled_back_without_leaking_pages() { + let mut cache = GpuCache::new(4 * 256, 256).unwrap(); + assert!(matches!( + cache.admit("loading".into(), 300, 1, 150, 2, false, true), + Admission::Admitted { .. } + )); + assert!(!cache.contains("loading")); + assert_eq!(cache.allocated_bytes(), 512); + cache.remove("loading").unwrap(); + assert_eq!(cache.allocated_bytes(), 0); + assert_eq!(cache.free_bytes(), cache.capacity()); + } + + #[test] + fn tenant_limit_prevents_one_tenant_from_owning_the_entire_arena() { + let mut cache = GpuCache::new_with_limits(4 * 256, 256, 50, 100, HashMap::new()).unwrap(); + for key in ["a-1", "a-2"] { + assert!(matches!( + cache.admit_for_tenant("tenant-a", key.into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + cache.mark_ready(key).unwrap(); + cache.release(key).unwrap(); + } + assert!(matches!( + cache.admit_for_tenant("tenant-b", "b-1".into(), 200, 1, 100, 2, false, true,), + Admission::Admitted { .. } + )); + assert_eq!(cache.tenant_count(), 2); + } + + #[test] + fn pinned_budget_fails_closed_before_consuming_pages() { + let mut cache = GpuCache::new_with_limits(4 * 256, 256, 100, 25, HashMap::new()).unwrap(); + assert!(matches!( + cache.admit_for_tenant( + "__resident__", + "resident".into(), + 200, + 1, + 100, + 2, + true, + true, + ), + Admission::Admitted { .. } + )); + assert_eq!(cache.pinned_bytes(), 256); + assert_eq!( + cache.admit_for_tenant( + "__resident__", + "too-much".into(), + 200, + 1, + 100, + 2, + true, + true, + ), + Admission::Deferred + ); + assert_eq!(cache.allocated_bytes(), 256); + } + + #[test] + fn tenant_can_recycle_its_own_reserved_pages() { + let mut reservations = HashMap::new(); + reservations.insert("tenant-a".to_owned(), 2 * 256); + let mut cache = GpuCache::new_with_limits(2 * 256, 256, 50, 100, reservations).unwrap(); + for key in ["a-1", "a-2"] { + assert!(matches!( + cache.admit_for_tenant("tenant-a", key.into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + cache.mark_ready(key).unwrap(); + cache.release(key).unwrap(); + } + assert!(matches!( + cache.admit_for_tenant("tenant-a", "a-3".into(), 200, 1, 100, 2, false, true), + Admission::Admitted { .. } + )); + assert_eq!(cache.entry_count(), 2); + } +} diff --git a/services/tilemaxsimd/src/engine.rs b/services/tilemaxsimd/src/engine.rs new file mode 100644 index 00000000..99480882 --- /dev/null +++ b/services/tilemaxsimd/src/engine.rs @@ -0,0 +1,1163 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use crate::cache::{Admission, GpuCache}; +use crate::gpu::Gpu; +use crate::protocol::{Descriptor, Request}; +use crate::shard::{HostCacheStatus, ShardStore, cache_key}; +use anyhow::{Result, anyhow, bail}; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex, mpsc}; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug)] +pub struct IoPipelineConfig { + pub overlap: bool, + pub batch_bytes: usize, +} + +impl IoPipelineConfig { + pub fn validate(self) -> Result { + if self.batch_bytes == 0 { + bail!("I/O pipeline batch size must be positive"); + } + Ok(self) + } +} + +#[derive(Clone, Debug, Default)] +pub struct IoPipelineStatus { + pub overlap_enabled: bool, + pub batch_bytes: usize, + pub resolve_batches: u64, + pub resolve_bytes: u64, + pub resolve_microseconds: u64, + pub upload_microseconds: u64, + pub compute_microseconds: u64, + pub overlap_cycles: u64, + pub overlap_microseconds: u64, +} + +struct MissingTensor { + candidate_index: usize, + descriptor: Descriptor, + key: String, + payload: Arc<[u8]>, +} + +struct UnresolvedBatch { + indices: Vec, + descriptors: Vec, + payload_bytes: usize, +} + +struct ResolvedBatch { + tensors: VecDeque, + payload_bytes: usize, + elapsed: Duration, +} + +struct PreparedBatch { + chunks: Vec>, + uploads: Vec)>>, +} + +impl PreparedBatch { + fn empty(device_count: usize) -> Self { + Self { + chunks: (0..device_count).map(|_| Vec::new()).collect(), + uploads: (0..device_count).map(|_| Vec::new()).collect(), + } + } + + fn is_empty(&self) -> bool { + self.chunks.iter().all(Vec::is_empty) + } +} + +struct ResidentTensor { + candidate_index: usize, + device: usize, + key: String, + offset: u64, + rows: u32, + transient: bool, + newly_admitted: bool, +} + +struct DeviceState { + gpu: Arc, + cache: GpuCache, + h2d_batches: u64, + h2d_bytes: u64, +} + +pub struct Engine { + devices: Vec, + store: Arc>, + next_device: usize, + io_pipeline: IoPipelineConfig, + io_status: IoPipelineStatus, +} + +#[derive(Clone, Debug, Default)] +pub struct DeviceStatus { + pub slot: usize, + pub device: i32, + pub capacity_bytes: usize, + pub block_bytes: usize, + pub free_bytes: usize, + pub largest_free_extent_bytes: usize, + pub allocated_bytes: usize, + pub payload_bytes: usize, + pub internal_waste_bytes: usize, + pub entries: usize, + pub pinned_entries: usize, + pub pinned_bytes: usize, + pub tenants: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, + pub h2d_batches: u64, + pub h2d_bytes: u64, +} + +#[derive(Clone, Debug, Default)] +pub struct EngineStatus { + pub devices: Vec, + pub host: HostCacheStatus, + pub batch_read_calls: u64, + pub batch_read_bytes: u64, + pub io_pipeline: IoPipelineStatus, +} + +impl Engine { + pub fn new( + gpus: Vec, + block_bytes: usize, + store: ShardStore, + tenant_cache_max_percent: u8, + pinned_cache_max_percent: u8, + tenant_reservations: &HashMap, + io_pipeline: IoPipelineConfig, + ) -> Result { + if gpus.is_empty() { + bail!("at least one GPU is required"); + } + let io_pipeline = io_pipeline.validate()?; + let devices = gpus + .into_iter() + .map(|gpu| { + let cache = GpuCache::new_with_limits( + gpu.tensor_bytes(), + block_bytes, + tenant_cache_max_percent, + pinned_cache_max_percent, + tenant_reservations.clone(), + ) + .map_err(|message| anyhow!(message))?; + Ok(DeviceState { + gpu: Arc::new(gpu), + cache, + h2d_batches: 0, + h2d_bytes: 0, + }) + }) + .collect::>>()?; + Ok(Self { + devices, + store: Arc::new(Mutex::new(store)), + next_device: 0, + io_pipeline, + io_status: IoPipelineStatus { + overlap_enabled: io_pipeline.overlap, + batch_bytes: io_pipeline.batch_bytes, + ..IoPipelineStatus::default() + }, + }) + } + + pub fn reload_shards(&mut self) -> Result<()> { + self.store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .reload() + } + + /// Return the H2D payload expected for a descriptor at the instant the + /// scheduler builds its next cooperative quantum. Host-cache residency does + /// not make this zero: every L0 miss still consumes upload bandwidth. + pub fn gpu_miss_bytes(&self, descriptor: &Descriptor) -> u64 { + let key = cache_key(descriptor); + if self + .devices + .iter() + .any(|device| device.cache.contains(&key)) + { + return 0; + } + descriptor_payload_bytes(descriptor) + .ok() + .and_then(|bytes| u64::try_from(bytes).ok()) + .unwrap_or(u64::MAX) + } + + pub fn prewarm(&mut self, descriptors: &[Descriptor], batch_size: usize) -> Result<()> { + if batch_size == 0 { + bail!("resident prewarm batch size must be positive"); + } + // A content-addressed manifest may legitimately reference the same tensor + // from more than one logical candidate. Upload each cache key once. Without + // this normalization, two duplicates in one batch both try to admit the + // same not-yet-ready entry and the second insert replaces the first arena + // allocation. + let descriptors = unique_descriptors(descriptors); + for batch in descriptors.chunks(batch_size) { + let payloads = self + .store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .resolve_many(batch, "__resident__")?; + let mut uploads = (0..self.devices.len()) + .map(|_| Vec::<(u64, &[u8])>::new()) + .collect::>(); + let mut acquired = Vec::<(usize, String, bool)>::new(); + for (descriptor, payload) in batch.iter().zip(&payloads) { + let key = cache_key(descriptor); + if let Some((device, _)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&key) + .map(|entry| (index, entry)) + }) + { + acquired.push((device, key, false)); + continue; + } + self.devices[self.next_device] + .cache + .record_access_miss(&key); + let mut admission = None; + for step in 0..self.devices.len() { + let device = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device].cache.admit_for_tenant( + "__resident__", + key.clone(), + payload.len(), + descriptor.rows, + descriptor.dimension, + descriptor.dtype, + true, + true, + ) + { + admission = Some((device, offset)); + self.next_device = (device + 1) % self.devices.len(); + break; + } + } + let Some((device, offset)) = admission else { + bail!("resident manifest exceeds the configured Rust GPU block caches"); + }; + uploads[device].push((offset, payload.as_ref())); + acquired.push((device, key, true)); + } + let mut upload_succeeded = vec![true; self.devices.len()]; + let mut upload_error = None; + for (device, items) in uploads.iter().enumerate() { + if items.is_empty() { + continue; + } + match self.devices[device].gpu.upload_batch(items) { + Ok(()) => { + self.devices[device].h2d_batches += 1; + self.devices[device].h2d_bytes += items + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum::(); + } + Err(error) => { + upload_succeeded[device] = false; + upload_error.get_or_insert(error); + } + } + } + for (device, key, newly_admitted) in acquired { + if newly_admitted && upload_succeeded[device] { + self.devices[device] + .cache + .mark_ready(&key) + .map_err(|message| anyhow!(message))?; + } + if newly_admitted && !upload_succeeded[device] { + self.devices[device] + .cache + .remove(&key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[device] + .cache + .release(&key) + .map_err(|message| anyhow!(message))?; + } + } + if let Some(error) = upload_error { + return Err(error); + } + } + Ok(()) + } + + pub fn score(&mut self, request: &Request) -> Result> { + if request.candidates.is_empty() { + return Ok(Vec::new()); + } + let mut scores = vec![None; request.candidates.len()]; + let mut hit_chunks = (0..self.devices.len()) + .map(|_| Vec::::new()) + .collect::>(); + let mut missing_descriptors = Vec::new(); + let mut missing_indices = Vec::new(); + let mut first_candidate_by_key = HashMap::::new(); + let mut duplicate_candidates = Vec::<(usize, usize)>::new(); + for (index, descriptor) in request.candidates.iter().enumerate() { + let key = cache_key(descriptor); + if let Some(first_index) = first_candidate_by_key.get(&key) { + duplicate_candidates.push((index, *first_index)); + continue; + } + first_candidate_by_key.insert(key.clone(), index); + let hit_device = self + .devices + .iter() + .position(|device| device.cache.contains(&key)); + if let Some(device_index) = hit_device { + let entry = self.devices[device_index] + .cache + .get(&key) + .expect("cache hit disappeared"); + validate_entry(descriptor, &entry)?; + hit_chunks[device_index].push(ResidentTensor { + candidate_index: index, + device: device_index, + key, + offset: entry.offset, + rows: entry.rows, + transient: false, + newly_admitted: false, + }); + } else { + // Record one request-level miss on the device that will get the + // first admission opportunity. Other devices are not polluted. + self.devices[self.next_device] + .cache + .record_access_miss(&key); + missing_indices.push(index); + missing_descriptors.push(descriptor.clone()); + } + } + + let unresolved = split_unresolved_batches( + missing_indices, + missing_descriptors, + self.io_pipeline.batch_bytes, + )?; + let hits = PreparedBatch { + chunks: hit_chunks, + uploads: (0..self.devices.len()).map(|_| Vec::new()).collect(), + }; + if self.io_pipeline.overlap { + let stage_before = self.pipeline_stage_microseconds(); + let pipeline_started = Instant::now(); + self.score_overlapped(request, hits, unresolved, &mut scores)?; + let stage_delta = self + .pipeline_stage_microseconds() + .saturating_sub(stage_before); + let hidden = + stage_delta.saturating_sub(duration_microseconds(pipeline_started.elapsed())); + self.io_status.overlap_microseconds = + self.io_status.overlap_microseconds.saturating_add(hidden); + } else { + self.score_serial(request, hits, unresolved, &mut scores)?; + } + + // Equal content-addressed tensors have equal TileMaxSim scores. Preserve + // every logical candidate id while avoiding duplicate cache acquisitions, + // uploads, and kernel work inside one request. + for (duplicate_index, first_index) in duplicate_candidates { + scores[duplicate_index] = scores[first_index]; + } + + request + .candidates + .iter() + .enumerate() + .map(|(index, descriptor)| { + scores[index] + .map(|score| (descriptor.candidate_id, score)) + .ok_or_else(|| anyhow!("missing native TileMaxSim result")) + }) + .collect() + } + + fn score_serial( + &mut self, + request: &Request, + mut current: PreparedBatch, + mut unresolved: VecDeque, + scores: &mut [Option], + ) -> Result<()> { + let mut pending = VecDeque::new(); + loop { + if !current.is_empty() { + let result = score_chunks(&self.gpus(), request, ¤t.chunks); + let cleanup = self.release_chunks(¤t.chunks); + let (computed, elapsed) = result?; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(elapsed)); + apply_scores(scores, computed); + cleanup?; + } + if pending.is_empty() && unresolved.is_empty() { + return Ok(()); + } + current = + self.fill_next_batch(&request.tenant, &mut unresolved, &mut pending, false)?; + } + } + + fn score_overlapped( + &mut self, + request: &Request, + current: PreparedBatch, + unresolved: VecDeque, + scores: &mut [Option], + ) -> Result<()> { + let (sender, receiver) = mpsc::sync_channel::>(1); + let store = Arc::clone(&self.store); + let tenant = request.tenant.clone(); + std::thread::scope(|scope| { + let resolver = scope.spawn(move || { + for batch in unresolved { + let result = resolve_batch(&store, batch, &tenant); + let failed = result.is_err(); + if sender.send(result).is_err() || failed { + break; + } + } + }); + let result = self.score_overlapped_inner(request, current, &receiver, scores); + // An early CUDA or cache error must unblock a resolver waiting on the + // bounded one-batch channel before the scoped worker can join. + drop(receiver); + let resolver_result = resolver + .join() + .map_err(|_| anyhow!("tensor resolver coordinator panicked")); + result.and(resolver_result) + }) + } + + fn score_overlapped_inner( + &mut self, + request: &Request, + mut current: PreparedBatch, + receiver: &mpsc::Receiver>, + scores: &mut [Option], + ) -> Result<()> { + let mut pending = VecDeque::new(); + let mut resolver_done = false; + loop { + if current.is_empty() { + current = self.fill_next_resolved_batch( + &request.tenant, + receiver, + &mut pending, + &mut resolver_done, + false, + )?; + if current.is_empty() && resolver_done { + return Ok(()); + } + continue; + } + + if pending.is_empty() && resolver_done { + let result = score_chunks(&self.gpus(), request, ¤t.chunks); + let cleanup = self.release_chunks(¤t.chunks); + let (computed, elapsed) = result?; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(elapsed)); + apply_scores(scores, computed); + cleanup?; + return Ok(()); + } + + let gpus = self.gpus(); + let (score_result, next_result) = std::thread::scope(|scope| { + let score_worker = scope.spawn(|| score_chunks(&gpus, request, ¤t.chunks)); + let next = self.fill_next_resolved_batch( + &request.tenant, + receiver, + &mut pending, + &mut resolver_done, + true, + ); + let score = score_worker + .join() + .map_err(|_| anyhow!("GPU score coordinator panicked")) + .and_then(|result| result); + (score, next) + }); + let cleanup_current = self.release_chunks(¤t.chunks); + + let (computed, compute_elapsed) = match score_result { + Ok(result) => result, + Err(error) => { + if let Ok(next) = next_result { + self.release_chunks(&next.chunks)?; + } + cleanup_current?; + return Err(error); + } + }; + self.io_status.compute_microseconds = self + .io_status + .compute_microseconds + .saturating_add(duration_microseconds(compute_elapsed)); + apply_scores(scores, computed); + cleanup_current?; + + let next = next_result?; + self.io_status.overlap_cycles = self.io_status.overlap_cycles.saturating_add(1); + current = next; + } + } + + fn fill_next_resolved_batch( + &mut self, + tenant: &str, + receiver: &mpsc::Receiver>, + pending: &mut VecDeque, + resolver_done: &mut bool, + allow_deferred: bool, + ) -> Result { + if pending.is_empty() && !*resolver_done { + match receiver.recv() { + Ok(result) => { + let resolved = result?; + self.record_resolution(&resolved); + *pending = resolved.tensors; + } + Err(_) => *resolver_done = true, + } + } + if pending.is_empty() { + return Ok(PreparedBatch::empty(self.devices.len())); + } + let mut prepared = self.prepare_pending(pending, tenant, allow_deferred)?; + if prepared.is_empty() { + return Ok(prepared); + } + let upload_elapsed = self.upload_prepared(&mut prepared)?; + self.io_status.upload_microseconds = self + .io_status + .upload_microseconds + .saturating_add(duration_microseconds(upload_elapsed)); + Ok(prepared) + } + + fn fill_next_batch( + &mut self, + tenant: &str, + unresolved: &mut VecDeque, + pending: &mut VecDeque, + allow_deferred: bool, + ) -> Result { + if pending.is_empty() + && let Some(batch) = unresolved.pop_front() + { + let resolved = resolve_batch(&self.store, batch, tenant)?; + self.record_resolution(&resolved); + *pending = resolved.tensors; + } + if pending.is_empty() { + return Ok(PreparedBatch::empty(self.devices.len())); + } + let mut prepared = self.prepare_pending(pending, tenant, allow_deferred)?; + if prepared.is_empty() { + return Ok(prepared); + } + let upload_elapsed = self.upload_prepared(&mut prepared)?; + self.io_status.upload_microseconds = self + .io_status + .upload_microseconds + .saturating_add(duration_microseconds(upload_elapsed)); + Ok(prepared) + } + + fn record_resolution(&mut self, resolved: &ResolvedBatch) { + self.io_status.resolve_batches = self.io_status.resolve_batches.saturating_add(1); + self.io_status.resolve_bytes = self + .io_status + .resolve_bytes + .saturating_add(u64::try_from(resolved.payload_bytes).unwrap_or(u64::MAX)); + self.io_status.resolve_microseconds = self + .io_status + .resolve_microseconds + .saturating_add(duration_microseconds(resolved.elapsed)); + } + + fn pipeline_stage_microseconds(&self) -> u64 { + self.io_status + .resolve_microseconds + .saturating_add(self.io_status.upload_microseconds) + .saturating_add(self.io_status.compute_microseconds) + } + + fn prepare_pending( + &mut self, + pending: &mut VecDeque, + tenant: &str, + allow_deferred: bool, + ) -> Result { + let mut prepared = PreparedBatch::empty(self.devices.len()); + let mut consumed = 0; + for tensor in pending.iter() { + if let Some((device_index, entry)) = + self.devices + .iter_mut() + .enumerate() + .find_map(|(index, device)| { + device + .cache + .acquire_existing(&tensor.key) + .map(|entry| (index, entry)) + }) + { + prepared.chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset: entry.offset, + rows: entry.rows, + transient: false, + newly_admitted: false, + }); + consumed += 1; + continue; + } + + let mut admitted = None; + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + let admission = self.devices[device_index].cache.admit_for_tenant( + tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + false, + ); + if let Admission::Admitted { offset, .. } = admission { + admitted = Some((device_index, offset, false)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + + if admitted.is_none() && prepared.is_empty() { + // TinyLFU may reject a cold item, but every requested tensor must + // still be scored. A transient page run is removed after this + // chunk completes. During overlap it may be temporarily deferred + // because the current compute batch still owns all freeable pages. + for step in 0..self.devices.len() { + let device_index = (self.next_device + step) % self.devices.len(); + if let Admission::Admitted { offset, .. } = + self.devices[device_index].cache.admit_for_tenant( + tenant, + tensor.key.clone(), + tensor.payload.len(), + tensor.descriptor.rows, + tensor.descriptor.dimension, + tensor.descriptor.dtype, + false, + true, + ) + { + admitted = Some((device_index, offset, true)); + self.next_device = (device_index + 1) % self.devices.len(); + break; + } + } + } + + let Some((device_index, offset, transient)) = admitted else { + if !prepared.is_empty() || allow_deferred { + break; + } + bail!("one tensor cannot be scheduled in any configured Rust GPU block cache"); + }; + prepared.uploads[device_index].push((offset, Arc::clone(&tensor.payload))); + prepared.chunks[device_index].push(ResidentTensor { + candidate_index: tensor.candidate_index, + device: device_index, + key: tensor.key.clone(), + offset, + rows: tensor.descriptor.rows, + transient, + newly_admitted: true, + }); + consumed += 1; + } + + if consumed == 0 { + if allow_deferred { + return Ok(prepared); + } + bail!("Rust multi-GPU scheduler made no progress"); + } + pending.drain(..consumed); + Ok(prepared) + } + + fn upload_prepared(&mut self, prepared: &mut PreparedBatch) -> Result { + let started = Instant::now(); + let gpus = self.gpus(); + let results = std::thread::scope(|scope| { + let mut workers = Vec::new(); + for (device_index, (gpu, upload)) in gpus.iter().zip(&prepared.uploads).enumerate() { + if upload.is_empty() { + continue; + } + workers.push(( + device_index, + scope.spawn(move || -> Result<()> { + let items = upload + .iter() + .map(|(offset, payload)| (*offset, payload.as_ref())) + .collect::>(); + gpu.upload_batch(&items) + }), + )); + } + workers + .into_iter() + .map(|(device, worker)| { + let result = worker + .join() + .map_err(|_| anyhow!("GPU upload worker panicked")) + .and_then(|result| result); + (device, result) + }) + .collect::>() + }); + + let mut succeeded = vec![true; self.devices.len()]; + let mut first_error = None; + for (device, result) in results { + if let Err(error) = result { + succeeded[device] = false; + first_error.get_or_insert(error); + continue; + } + self.devices[device].h2d_batches = self.devices[device].h2d_batches.saturating_add(1); + self.devices[device].h2d_bytes = self.devices[device].h2d_bytes.saturating_add( + prepared.uploads[device] + .iter() + .map(|(_, payload)| payload.len() as u64) + .sum(), + ); + } + for (device, chunk) in prepared.chunks.iter().enumerate() { + if !succeeded[device] { + continue; + } + for tensor in chunk.iter().filter(|tensor| tensor.newly_admitted) { + if let Err(message) = self.devices[device].cache.mark_ready(&tensor.key) { + succeeded[device] = false; + first_error.get_or_insert_with(|| anyhow!(message)); + break; + } + } + } + prepared.uploads.iter_mut().for_each(Vec::clear); + if let Some(error) = first_error { + self.cleanup_after_upload_failure(&prepared.chunks, &succeeded)?; + Err(error) + } else { + Ok(started.elapsed()) + } + } + + fn gpus(&self) -> Vec> { + self.devices + .iter() + .map(|device| Arc::clone(&device.gpu)) + .collect() + } + + fn release_chunks(&mut self, chunks: &[Vec]) -> Result<()> { + for chunk in chunks { + for tensor in chunk { + if tensor.transient { + self.devices[tensor.device] + .cache + .remove(&tensor.key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[tensor.device] + .cache + .release(&tensor.key) + .map_err(|message| anyhow!(message))?; + } + } + } + Ok(()) + } + + fn cleanup_after_upload_failure( + &mut self, + chunks: &[Vec], + upload_succeeded: &[bool], + ) -> Result<()> { + for chunk in chunks { + for tensor in chunk { + if tensor.newly_admitted && (!upload_succeeded[tensor.device] || tensor.transient) { + self.devices[tensor.device] + .cache + .remove(&tensor.key) + .map_err(|message| anyhow!(message))?; + } else { + self.devices[tensor.device] + .cache + .release(&tensor.key) + .map_err(|message| anyhow!(message))?; + } + } + } + Ok(()) + } + + pub fn status_snapshot(&self) -> EngineStatus { + let devices = self + .devices + .iter() + .enumerate() + .map(|(slot, device)| DeviceStatus { + slot, + device: device.gpu.device(), + capacity_bytes: device.cache.capacity(), + block_bytes: device.cache.block_bytes(), + free_bytes: device.cache.free_bytes(), + largest_free_extent_bytes: device.cache.largest_free_extent(), + allocated_bytes: device.cache.allocated_bytes(), + payload_bytes: device.cache.payload_bytes(), + internal_waste_bytes: device + .cache + .allocated_bytes() + .saturating_sub(device.cache.payload_bytes()), + entries: device.cache.entry_count(), + pinned_entries: device.cache.pinned_entries(), + pinned_bytes: device.cache.pinned_bytes(), + tenants: device.cache.tenant_count(), + hits: device.cache.hits, + misses: device.cache.misses, + evictions: device.cache.evictions, + admission_rejections: device.cache.admission_rejections, + h2d_batches: device.h2d_batches, + h2d_bytes: device.h2d_bytes, + }) + .collect(); + let store = self.store.lock().unwrap_or_else(|error| error.into_inner()); + EngineStatus { + devices, + host: store.host_status(), + batch_read_calls: store.batch_read_calls, + batch_read_bytes: store.batch_read_bytes, + io_pipeline: self.io_status.clone(), + } + } + + pub fn status_json(&self) -> serde_json::Value { + let status = self.status_snapshot(); + let devices = status + .devices + .iter() + .map(|device| { + serde_json::json!({ + "index": device.slot, + "device": device.device, + "gpu_allocator": "segregated-page-runs", + "gpu_tensor_bytes": device.capacity_bytes, + "gpu_block_bytes": device.block_bytes, + "gpu_free_bytes": device.free_bytes, + "gpu_largest_free_extent_bytes": device.largest_free_extent_bytes, + "gpu_allocated_bytes": device.allocated_bytes, + "gpu_payload_bytes": device.payload_bytes, + "gpu_internal_waste_bytes": device.internal_waste_bytes, + "gpu_entries": device.entries, + "gpu_pinned_entries": device.pinned_entries, + "gpu_pinned_bytes": device.pinned_bytes, + "gpu_tenant_count": device.tenants, + "gpu_hits": device.hits, + "gpu_misses": device.misses, + "gpu_evictions": device.evictions, + "gpu_admission_rejections": device.admission_rejections, + "h2d_batches": device.h2d_batches, + "h2d_bytes": device.h2d_bytes, + }) + }) + .collect::>(); + serde_json::json!({ + "devices": devices, + "host_capacity_bytes": status.host.capacity_bytes, + "host_used_bytes": status.host.used_bytes, + "host_entries": status.host.entries, + "host_tenant_count": status.host.tenants, + "host_hits": status.host.hits, + "host_misses": status.host.misses, + "host_evictions": status.host.evictions, + "host_admission_rejections": status.host.admission_rejections, + "batch_read_calls": status.batch_read_calls, + "batch_read_bytes": status.batch_read_bytes, + "io_pipeline": { + "mode": if status.io_pipeline.overlap_enabled { "overlap" } else { "serial" }, + "batch_bytes": status.io_pipeline.batch_bytes, + "resolve_batches": status.io_pipeline.resolve_batches, + "resolve_bytes": status.io_pipeline.resolve_bytes, + "resolve_microseconds": status.io_pipeline.resolve_microseconds, + "upload_microseconds": status.io_pipeline.upload_microseconds, + "compute_microseconds": status.io_pipeline.compute_microseconds, + "overlap_cycles": status.io_pipeline.overlap_cycles, + "overlap_microseconds": status.io_pipeline.overlap_microseconds, + }, + }) + } +} + +fn resolve_batch( + store: &Arc>, + batch: UnresolvedBatch, + tenant: &str, +) -> Result { + let started = Instant::now(); + let payloads = store + .lock() + .map_err(|_| anyhow!("immutable shard store lock is poisoned"))? + .resolve_many(&batch.descriptors, tenant)?; + let elapsed = started.elapsed(); + let tensors = batch + .indices + .into_iter() + .zip(batch.descriptors) + .zip(payloads) + .map(|((candidate_index, descriptor), payload)| MissingTensor { + candidate_index, + key: cache_key(&descriptor), + descriptor, + payload, + }) + .collect(); + Ok(ResolvedBatch { + tensors, + payload_bytes: batch.payload_bytes, + elapsed, + }) +} + +fn split_unresolved_batches( + indices: Vec, + descriptors: Vec, + maximum_bytes: usize, +) -> Result> { + if indices.len() != descriptors.len() { + bail!("internal missing tensor metadata disagrees"); + } + let mut output = VecDeque::new(); + let mut batch_indices = Vec::new(); + let mut batch_descriptors = Vec::new(); + let mut batch_bytes = 0_usize; + for (index, descriptor) in indices.into_iter().zip(descriptors) { + let payload_bytes = descriptor_payload_bytes(&descriptor)?; + if !batch_descriptors.is_empty() + && batch_bytes.saturating_add(payload_bytes) > maximum_bytes + { + output.push_back(UnresolvedBatch { + indices: std::mem::take(&mut batch_indices), + descriptors: std::mem::take(&mut batch_descriptors), + payload_bytes: batch_bytes, + }); + batch_bytes = 0; + } + batch_indices.push(index); + batch_descriptors.push(descriptor); + batch_bytes = batch_bytes + .checked_add(payload_bytes) + .ok_or_else(|| anyhow!("I/O pipeline batch byte count overflow"))?; + } + if !batch_descriptors.is_empty() { + output.push_back(UnresolvedBatch { + indices: batch_indices, + descriptors: batch_descriptors, + payload_bytes: batch_bytes, + }); + } + Ok(output) +} + +fn descriptor_payload_bytes(descriptor: &Descriptor) -> Result { + let scalar_bytes = match descriptor.dtype { + 1 => 4_usize, + 2 => 2_usize, + _ => bail!("unsupported tensor dtype"), + }; + (descriptor.rows as usize) + .checked_mul(descriptor.dimension as usize) + .and_then(|value| value.checked_mul(scalar_bytes)) + .ok_or_else(|| anyhow!("tensor payload byte count overflow")) +} + +fn score_chunks( + gpus: &[Arc], + request: &Request, + chunks: &[Vec], +) -> Result<(Vec<(usize, f32)>, Duration)> { + let started = Instant::now(); + let completed = std::thread::scope(|scope| -> Result>> { + let mut workers = Vec::new(); + for (gpu, chunk) in gpus.iter().zip(chunks) { + if chunk.is_empty() { + continue; + } + workers.push(scope.spawn(move || -> Result> { + let offsets = chunk.iter().map(|item| item.offset).collect::>(); + let rows = chunk.iter().map(|item| item.rows).collect::>(); + let computed = gpu.score( + &request.query, + request.query_rows, + request.dimension, + request.dtype, + &offsets, + &rows, + )?; + Ok(chunk + .iter() + .zip(computed) + .map(|(tensor, score)| (tensor.candidate_index, score)) + .collect()) + })); + } + let mut completed = Vec::new(); + for worker in workers { + completed.push( + worker + .join() + .map_err(|_| anyhow!("GPU worker panicked"))??, + ); + } + Ok(completed) + })?; + Ok((completed.into_iter().flatten().collect(), started.elapsed())) +} + +fn apply_scores(scores: &mut [Option], computed: Vec<(usize, f32)>) { + for (candidate_index, score) in computed { + scores[candidate_index] = Some(score); + } +} + +fn duration_microseconds(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + +fn unique_descriptors(descriptors: &[Descriptor]) -> Vec { + let mut seen = HashSet::with_capacity(descriptors.len()); + descriptors + .iter() + .filter(|descriptor| seen.insert(cache_key(descriptor))) + .cloned() + .collect() +} + +fn validate_entry(descriptor: &Descriptor, entry: &crate::cache::CacheEntry) -> Result<()> { + let scalar_bytes = if descriptor.dtype == 1 { 4 } else { 2 }; + let expected_bytes = descriptor.rows as usize * descriptor.dimension as usize * scalar_bytes; + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.payload_bytes != expected_bytes + || entry.allocated_bytes < expected_bytes + { + bail!("GPU cache metadata disagrees with the tensor descriptor"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn descriptor(candidate_id: u32, digest: &str, rows: u32) -> Descriptor { + Descriptor { + candidate_id, + contract: "colqwen35@test".to_owned(), + digest: digest.to_owned(), + rows, + dimension: 320, + dtype: 2, + } + } + + #[test] + fn resident_prewarm_deduplicates_content_references() { + let descriptors = vec![ + descriptor(1, "a", 100), + descriptor(2, "a", 100), + descriptor(3, "b", 120), + // Shape is part of the cache identity and must not be collapsed. + descriptor(4, "a", 101), + ]; + + let unique = unique_descriptors(&descriptors); + assert_eq!(unique.len(), 3); + assert_eq!(unique[0].candidate_id, 1); + assert_eq!(unique[1].candidate_id, 3); + assert_eq!(unique[2].candidate_id, 4); + } + + #[test] + fn io_batches_are_payload_bounded_and_keep_one_oversized_tensor() { + let descriptors = vec![ + descriptor(1, "a", 1), + descriptor(2, "b", 1), + descriptor(3, "c", 3), + ]; + // One row is 640 bytes. The first two fit exactly; the final tensor is + // larger than the target but must remain a schedulable one-item batch. + let batches = split_unresolved_batches(vec![0, 1, 2], descriptors, 1_280).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].indices, vec![0, 1]); + assert_eq!(batches[0].payload_bytes, 1_280); + assert_eq!(batches[1].indices, vec![2]); + assert_eq!(batches[1].payload_bytes, 1_920); + } +} diff --git a/services/tilemaxsimd/src/gpu.rs b/services/tilemaxsimd/src/gpu.rs new file mode 100644 index 00000000..2059b161 --- /dev/null +++ b/services/tilemaxsimd/src/gpu.rs @@ -0,0 +1,183 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use anyhow::{Result, anyhow, bail}; +use std::ffi::{CStr, c_char, c_int, c_uchar, c_void}; +use std::ptr::NonNull; + +#[repr(C)] +struct NativeGpu(c_void); + +unsafe extern "C" { + fn vctm_gpu_create( + device: c_int, + total_bytes: usize, + workspace_bytes: usize, + output: *mut *mut NativeGpu, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_destroy(gpu: *mut NativeGpu); + fn vctm_gpu_tensor_bytes(gpu: *const NativeGpu) -> usize; + fn vctm_gpu_upload_batch( + gpu: *mut NativeGpu, + offsets: *const u64, + payloads: *const *const c_uchar, + lengths: *const usize, + count: usize, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; + fn vctm_gpu_score( + gpu: *mut NativeGpu, + query: *const c_uchar, + query_bytes: usize, + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: *const u64, + document_rows: *const u32, + count: usize, + output: *mut f32, + error: *mut c_char, + error_capacity: usize, + ) -> c_int; +} + +pub struct Gpu { + native: NonNull, + device: i32, + tensor_bytes: usize, +} + +// SAFETY: `Gpu` uniquely owns the native handle. The engine permits at most one +// upload and one score call at a time. The native implementation assigns those +// operations distinct CUDA streams and disjoint host/device workspaces; tensor +// cache references prevent an upload from overwriting pages used by scoring. +// Destruction happens only after every scoped worker has joined. +unsafe impl Send for Gpu {} +unsafe impl Sync for Gpu {} + +impl Gpu { + pub fn create(device: i32, total_bytes: usize, workspace_bytes: usize) -> Result { + let mut native = std::ptr::null_mut(); + let mut error = [0_i8; 512]; + // SAFETY: the C API writes one opaque pointer and a bounded error string. + let status = unsafe { + vctm_gpu_create( + device, + total_bytes, + workspace_bytes, + &mut native, + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + let native = NonNull::new(native).ok_or_else(|| anyhow!("CUDA returned a null arena"))?; + // SAFETY: `native` is live until Drop. + let tensor_bytes = unsafe { vctm_gpu_tensor_bytes(native.as_ptr()) }; + Ok(Self { + native, + device, + tensor_bytes, + }) + } + + pub fn device(&self) -> i32 { + self.device + } + + pub fn tensor_bytes(&self) -> usize { + self.tensor_bytes + } + + pub fn upload_batch(&self, items: &[(u64, &[u8])]) -> Result<()> { + if items.is_empty() { + return Ok(()); + } + let offsets = items.iter().map(|item| item.0).collect::>(); + let payloads = items.iter().map(|item| item.1.as_ptr()).collect::>(); + let lengths = items.iter().map(|item| item.1.len()).collect::>(); + let mut error = [0_i8; 512]; + // SAFETY: all slices remain alive for the synchronous native batch call. + let status = unsafe { + vctm_gpu_upload_batch( + self.native.as_ptr(), + offsets.as_ptr(), + payloads.as_ptr(), + lengths.as_ptr(), + items.len(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + Ok(()) + } + + pub fn score( + &self, + query: &[u8], + query_rows: u32, + dimension: u32, + dtype: u8, + document_offsets: &[u64], + document_rows: &[u32], + ) -> Result> { + if document_offsets.len() != document_rows.len() || document_offsets.is_empty() { + bail!("invalid native document metadata"); + } + let mut output = vec![0.0_f32; document_offsets.len()]; + let mut error = [0_i8; 512]; + // SAFETY: the native call is synchronous and receives valid slice pointers. + let status = unsafe { + vctm_gpu_score( + self.native.as_ptr(), + query.as_ptr(), + query.len(), + query_rows, + dimension, + dtype, + document_offsets.as_ptr(), + document_rows.as_ptr(), + document_offsets.len(), + output.as_mut_ptr(), + error.as_mut_ptr(), + error.len(), + ) + }; + if status != 0 { + bail!(native_error(&error)); + } + if output.iter().any(|score| !score.is_finite()) { + bail!("native TileMaxSim returned a non-finite score"); + } + Ok(output) + } +} + +impl Drop for Gpu { + fn drop(&mut self) { + // SAFETY: this is the unique owned native pointer. + unsafe { vctm_gpu_destroy(self.native.as_ptr()) }; + } +} + +fn native_error(buffer: &[c_char]) -> String { + // SAFETY: the native helper always NUL-terminates a nonempty error buffer. + unsafe { CStr::from_ptr(buffer.as_ptr()) } + .to_string_lossy() + .into_owned() +} diff --git a/services/tilemaxsimd/src/main.rs b/services/tilemaxsimd/src/main.rs new file mode 100644 index 00000000..b35beda8 --- /dev/null +++ b/services/tilemaxsimd/src/main.rs @@ -0,0 +1,2495 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +mod cache; +mod engine; +mod gpu; +mod protocol; +mod scheduler; +mod shard; + +use anyhow::{Context, Result, anyhow, bail}; +use clap::Parser; +use engine::{Engine, EngineStatus, IoPipelineConfig}; +use gpu::Gpu; +use protocol::{HEADER_BYTES, VERSION_EXTERNAL, VERSION_SCHEDULED_EXTERNAL}; +use scheduler::{RequestQueue, Scheduled, SchedulerPolicy}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use shard::ShardStore; +use std::collections::HashMap; +use std::fmt::Write as FmtWrite; +use std::fs::{self, OpenOptions}; +use std::io::{BufRead, Read, Write}; +use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::os::fd::AsRawFd; +use std::os::unix::fs::{FileTypeExt, PermissionsExt}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, mpsc}; +use std::thread; +use std::time::{Duration, Instant}; + +const GIB: usize = 1024 * 1024 * 1024; +static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false); +static RELOAD_REQUESTED: AtomicBool = AtomicBool::new(false); + +extern "C" fn handle_signal(signal: libc::c_int) { + if signal == libc::SIGHUP { + RELOAD_REQUESTED.store(true, Ordering::Relaxed); + } else { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); + } +} + +fn install_signal_handlers() -> Result<()> { + for signal in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP] { + if unsafe { libc::signal(signal, handle_signal as *const () as libc::sighandler_t) } + == libc::SIG_ERR + { + return Err(std::io::Error::last_os_error().into()); + } + } + Ok(()) +} + +#[derive(Parser)] +#[command(about = "Native Rust/CUDA TileMaxSim shard, cache, and scheduling daemon")] +struct Args { + #[arg(long)] + socket: PathBuf, + /// Optional TCP scoring listener for a database running in another pod. + #[arg(long)] + listen: Option, + #[arg(long, required = true, value_delimiter = ',', value_parser = parse_gpu_memory)] + gpu_memory_gb: Vec, + #[arg(long, default_value = "2", value_parser = parse_gb)] + gpu_workspace_gb: usize, + #[arg(long, default_value = "8", value_parser = parse_gb)] + host_cache_gb: usize, + /// Overlap tensor resolution and H2D upload with the current CUDA batch. + #[arg(long, default_value = "overlap", value_parser = ["serial", "overlap"])] + io_pipeline: String, + /// Maximum logical tensor payload resolved by one pipeline stage, in GB. + #[arg(long, default_value = "0.05", value_parser = parse_gb)] + io_batch_gb: usize, + #[arg(long, default_value_t = 80, value_parser = clap::value_parser!(u8).range(1..=100))] + host_tenant_cache_max_percent: u8, + #[arg(long = "contract-root", required = true, value_parser = parse_contract_root)] + contract_roots: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 32)] + gpu_block_kib: usize, + #[arg(long, default_value_t = 64 * 1024 * 1024)] + max_request_bytes: usize, + #[arg(long, default_value = "1", value_parser = parse_gb)] + max_inflight_request_gb: usize, + #[arg(long, default_value = "600", value_parser = parse_mode)] + socket_mode: u32, + #[arg(long)] + status_socket: Option, + #[arg(long)] + status_listen: Option, + #[arg(long, default_value = "600", value_parser = parse_mode)] + status_socket_mode: u32, + #[arg(long)] + once: bool, + #[arg(long)] + verify_full_shards: bool, + #[arg(long, default_value = "lru", value_parser = ["lru", "resident"])] + gpu_cache_mode: String, + #[arg(long = "resident-manifest", value_parser = parse_contract_root)] + resident_manifests: Vec<(String, PathBuf)>, + #[arg(long, default_value_t = 256)] + prewarm_batch_size: usize, + #[arg(long, default_value_t = 80, value_parser = clap::value_parser!(u8).range(1..=100))] + tenant_cache_max_percent: u8, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u8).range(0..=100))] + pinned_cache_max_percent: u8, + #[arg(long = "tenant-cache-reservation", value_parser = parse_tenant_memory)] + tenant_cache_reservations: Vec<(String, usize)>, + #[arg(long, default_value_t = 256)] + max_connections: usize, + #[arg(long, default_value_t = 128)] + max_queued_requests: usize, + #[arg(long, default_value_t = 16)] + max_tenant_queued_requests: usize, + #[arg(long, default_value_t = 5000)] + socket_io_timeout_ms: u64, + #[arg(long, default_value_t = 8000)] + request_timeout_ms: u64, + #[arg(long, default_value = "fair-priority", value_parser = parse_scheduler_policy)] + scheduler_policy: SchedulerPolicy, + #[arg(long, default_value_t = 1000)] + priority_aging_ms: u64, + #[arg(long, default_value_t = 200)] + priority_band: i32, + #[arg(long, default_value_t = 2)] + scheduler_batch_window_ms: u64, + #[arg(long, default_value_t = 1024)] + scheduler_quantum_candidates: usize, + #[arg(long, default_value_t = 250_000)] + scheduler_quantum_tokens: u64, + #[arg(long, default_value_t = 4_000_000_000)] + scheduler_quantum_fmas: u64, + /// Maximum L0-miss payload admitted to one cooperative quantum, in GB. + #[arg(long, default_value = "1", value_parser = parse_gb)] + scheduler_quantum_io_gb: usize, + #[arg(long = "tenant-weight", value_parser = parse_tenant_weight)] + tenant_weights: Vec<(String, f64)>, +} + +#[derive(Clone)] +struct GpuMemory { + device: i32, + bytes: usize, +} + +fn parse_gpu_memory(value: &str) -> Result { + let (device, gb) = value + .strip_prefix("cuda:") + .unwrap_or(value) + .split_once('=') + .ok_or_else(|| "GPU memory must be GPU=GB, for example 1=20".to_owned())?; + let device = device + .parse::() + .map_err(|_| "GPU index must be a nonnegative integer".to_owned())?; + if device < 0 { + return Err("GPU index must be nonnegative".to_owned()); + } + let bytes = parse_gb(gb)?; + Ok(GpuMemory { device, bytes }) +} + +fn parse_gb(value: &str) -> Result { + if value.is_empty() + || !value + .bytes() + .all(|byte| byte.is_ascii_digit() || byte == b'.') + { + return Err("memory size must be a positive number of GB".to_owned()); + } + let gb = value + .parse::() + .map_err(|_| "memory size must be a positive number of GB".to_owned())?; + if !gb.is_finite() || gb <= 0.0 || gb * GIB as f64 > usize::MAX as f64 { + return Err("memory size must be a positive number of GB".to_owned()); + } + Ok((gb * GIB as f64) as usize) +} + +fn kib_to_bytes(kib: usize) -> Result { + kib.checked_mul(1024).ok_or("GPU block size overflow") +} + +fn parse_contract_root(value: &str) -> Result<(String, PathBuf), String> { + let (contract, path) = value + .split_once('=') + .ok_or_else(|| "contract root must be MODEL_CONTRACT_ID=/absolute/path".to_owned())?; + let path = PathBuf::from(path); + if contract.is_empty() || !path.is_absolute() { + return Err("contract root must contain a nonempty ID and absolute path".to_owned()); + } + Ok((contract.to_owned(), path)) +} + +fn parse_mode(value: &str) -> Result { + let mode = u32::from_str_radix(value, 8).map_err(|_| "invalid octal socket mode".to_owned())?; + if mode > 0o777 { + return Err("socket mode must be between 000 and 777".to_owned()); + } + Ok(mode) +} + +fn parse_scheduler_policy(value: &str) -> Result { + SchedulerPolicy::parse(value) +} + +fn parse_tenant_weight(value: &str) -> Result<(String, f64), String> { + let (tenant, weight) = value + .split_once('=') + .ok_or_else(|| "tenant weight must be TENANT=WEIGHT".to_owned())?; + if tenant.is_empty() + || tenant.len() > 256 + || tenant.chars().any(|character| character.is_control()) + { + return Err("tenant weight has an invalid tenant name".to_owned()); + } + let weight = weight + .parse::() + .map_err(|_| "tenant weight must be a finite positive number".to_owned())?; + if !weight.is_finite() || weight <= 0.0 || weight > 1000.0 { + return Err("tenant weight must be between 0 and 1000".to_owned()); + } + Ok((tenant.to_owned(), weight)) +} + +fn parse_tenant_memory(value: &str) -> Result<(String, usize), String> { + let (tenant, gb) = value + .split_once('=') + .ok_or_else(|| "tenant cache reservation must be TENANT=GB".to_owned())?; + if tenant.is_empty() + || tenant.len() > 256 + || tenant.chars().any(|character| character.is_control()) + { + return Err("tenant cache reservation has an invalid tenant name".to_owned()); + } + Ok((tenant.to_owned(), parse_gb(gb)?)) +} + +fn main() -> Result<()> { + let args = Args::parse(); + if args.max_connections == 0 + || args.max_queued_requests == 0 + || args.max_tenant_queued_requests == 0 + || args.socket_io_timeout_ms == 0 + || args.request_timeout_ms == 0 + || args.priority_aging_ms == 0 + || args.scheduler_quantum_candidates == 0 + || args.scheduler_quantum_tokens == 0 + || args.scheduler_quantum_fmas == 0 + { + bail!("connection, queue, timeout, and priority-aging limits must be positive"); + } + let mut seen_devices = std::collections::HashSet::new(); + for specification in &args.gpu_memory_gb { + if args.gpu_workspace_gb >= specification.bytes { + bail!("every configured GPU allocation must exceed its workspace"); + } + if !seen_devices.insert(specification.device) { + bail!("each CUDA device may be configured only once"); + } + } + let block_bytes = kib_to_bytes(args.gpu_block_kib).map_err(|message| anyhow!(message))?; + if block_bytes == 0 || block_bytes % 256 != 0 { + bail!("GPU block size must be positive and 256-byte aligned"); + } + let store = ShardStore::open( + &args.contract_roots, + args.host_cache_gb, + args.host_tenant_cache_max_percent, + args.verify_full_shards, + )?; + let gpus = args + .gpu_memory_gb + .iter() + .map(|specification| { + Gpu::create( + specification.device, + specification.bytes, + args.gpu_workspace_gb, + ) + }) + .collect::>>()?; + let tenant_cache_reservations = args.tenant_cache_reservations.iter().cloned().collect(); + let mut engine = Engine::new( + gpus, + block_bytes, + store, + args.tenant_cache_max_percent, + args.pinned_cache_max_percent, + &tenant_cache_reservations, + IoPipelineConfig { + overlap: args.io_pipeline == "overlap", + batch_bytes: args.io_batch_gb, + }, + )?; + if args.gpu_cache_mode == "resident" && args.resident_manifests.is_empty() { + bail!("resident GPU cache mode requires at least one resident manifest"); + } + if args.gpu_cache_mode == "lru" && !args.resident_manifests.is_empty() { + bail!("resident manifests are valid only in resident GPU cache mode"); + } + if args.gpu_cache_mode == "resident" { + let prewarm_started = Instant::now(); + let descriptors = load_resident_manifests(&args.resident_manifests)?; + engine.prewarm(&descriptors, args.prewarm_batch_size)?; + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_prewarm_complete", + "entries": descriptors.len(), + "elapsed_ms": prewarm_started.elapsed().as_secs_f64() * 1000.0, + "cache": engine.status_json(), + }) + ); + } + let _instance_lock = acquire_instance_lock(&args.socket)?; + remove_stale_socket(&args.socket)?; + let listener = UnixListener::bind(&args.socket) + .with_context(|| format!("cannot bind {}", args.socket.display()))?; + fs::set_permissions(&args.socket, fs::Permissions::from_mode(args.socket_mode))?; + listener.set_nonblocking(true)?; + if args.listen.is_some() && args.listen == args.status_listen { + bail!("scoring and status TCP listeners must use different addresses"); + } + let tcp_listener = if let Some(address) = args.listen { + let listener = TcpListener::bind(address) + .with_context(|| format!("cannot bind scoring listener {address}"))?; + listener.set_nonblocking(true)?; + Some(listener) + } else { + None + }; + if args.status_socket.as_ref() == Some(&args.socket) { + bail!("status socket must differ from the TileMaxSim protocol socket"); + } + let reload = Arc::new(AtomicBool::new(false)); + let metrics = Arc::new(RuntimeMetrics::new( + args.max_connections, + args.max_queued_requests, + args.max_tenant_queued_requests, + args.max_inflight_request_gb, + )); + metrics.update_engine(engine.status_snapshot()); + install_signal_handlers()?; + let ready_cache = engine.status_json(); + + let (sender, receiver) = mpsc::sync_channel::(args.max_queued_requests); + let frame_admission = Arc::new(ByteAdmission::new( + args.max_inflight_request_gb, + Arc::clone(&metrics), + )); + let pending_admission = Arc::new(PendingAdmission::new( + args.max_queued_requests, + args.max_tenant_queued_requests, + Arc::clone(&metrics), + )); + let tenant_weights = args.tenant_weights.iter().cloned().collect(); + let scheduler_config = SchedulerConfig { + policy: args.scheduler_policy, + priority_aging: Duration::from_millis(args.priority_aging_ms), + priority_band: args.priority_band, + batch_window: Duration::from_millis(args.scheduler_batch_window_ms), + quantum_candidates: args.scheduler_quantum_candidates, + quantum_tokens: args.scheduler_quantum_tokens, + quantum_fmas: args.scheduler_quantum_fmas, + quantum_io_bytes: args.scheduler_quantum_io_gb, + socket_io_timeout: Duration::from_millis(args.socket_io_timeout_ms), + tenant_weights, + }; + let scheduler_reload = Arc::clone(&reload); + let scheduler_metrics = Arc::clone(&metrics); + let scheduler = thread::Builder::new() + .name("tilemaxsim-scheduler".to_owned()) + .spawn(move || { + run_scheduler( + engine, + receiver, + scheduler_config, + scheduler_reload, + scheduler_metrics, + ) + })?; + let status_server = if let Some(path) = args.status_socket.clone() { + remove_stale_socket(&path)?; + let status_listener = UnixListener::bind(&path) + .with_context(|| format!("cannot bind status socket {}", path.display()))?; + fs::set_permissions(&path, fs::Permissions::from_mode(args.status_socket_mode))?; + status_listener.set_nonblocking(true)?; + let status_metrics = Arc::clone(&metrics); + Some( + thread::Builder::new() + .name("tilemaxsim-status".to_owned()) + .spawn(move || run_status_server(status_listener, path, status_metrics))?, + ) + } else { + None + }; + let status_tcp_server = if let Some(address) = args.status_listen { + let status_listener = TcpListener::bind(address) + .with_context(|| format!("cannot bind status listener {address}"))?; + status_listener.set_nonblocking(true)?; + let status_metrics = Arc::clone(&metrics); + Some( + thread::Builder::new() + .name("tilemaxsim-status-tcp".to_owned()) + .spawn(move || run_tcp_status_server(status_listener, status_metrics))?, + ) + } else { + None + }; + metrics.ready.store(true, Ordering::Release); + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_ready", + "socket": args.socket, + "listen": args.listen, + "devices": args.gpu_memory_gb.iter().map(|item| serde_json::json!({ + "device": item.device, + "allocated_bytes": item.bytes, + })).collect::>(), + "workspace_bytes": args.gpu_workspace_gb, + "scheduler_policy": format!("{:?}", args.scheduler_policy), + "max_connections": args.max_connections, + "max_queued_requests": args.max_queued_requests, + "max_tenant_queued_requests": args.max_tenant_queued_requests, + "max_inflight_request_bytes": args.max_inflight_request_gb, + "scheduler_quantum_fmas": args.scheduler_quantum_fmas, + "scheduler_quantum_io_bytes": args.scheduler_quantum_io_gb, + "io_pipeline": args.io_pipeline, + "io_batch_bytes": args.io_batch_gb, + "status_socket": args.status_socket, + "status_listen": args.status_listen, + "cache": ready_cache, + }) + ); + + let mut readers = Vec::new(); + let mut accepted = 0_usize; + let mut status_failed = false; + let mut scheduler_failed = false; + let mut fatal_error = None; + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + if scheduler.is_finished() { + scheduler_failed = true; + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + break; + } + if status_server + .as_ref() + .is_some_and(thread::JoinHandle::is_finished) + || status_tcp_server + .as_ref() + .is_some_and(thread::JoinHandle::is_finished) + { + status_failed = true; + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + break; + } + if RELOAD_REQUESTED.swap(false, Ordering::AcqRel) { + reload.store(true, Ordering::Release); + } + reap_readers(&mut readers); + let mut connections = Vec::with_capacity(2); + match listener.accept() { + Ok((connection, _)) => connections.push(ClientStream::Unix(connection)), + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) => {} + Err(error) => { + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } + if let Some(tcp_listener) = &tcp_listener { + match tcp_listener.accept() { + Ok((connection, _)) => { + connection.set_nodelay(true).ok(); + connections.push(ClientStream::Tcp(connection)); + } + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) => {} + Err(error) => { + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } + } + if connections.is_empty() { + thread::sleep(Duration::from_millis(5)); + continue; + } + for connection in connections { + if SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + break; + } + if !try_acquire_reader(&metrics.active_connections, args.max_connections) { + // Closing immediately is deliberate: we have not read enough + // bytes to know whether the peer expects a v2 or v3 response. + metrics.rejected_connections.fetch_add(1, Ordering::Relaxed); + drop(connection); + continue; + } + accepted += 1; + let reader_sender = sender.clone(); + let reader_metrics = Arc::clone(&metrics); + let reader_admission = Arc::clone(&pending_admission); + let request_metrics = Arc::clone(&metrics); + let reader_frame_admission = Arc::clone(&frame_admission); + let reader_config = ReaderConfig { + maximum: args.max_request_bytes, + io_timeout: Duration::from_millis(args.socket_io_timeout_ms), + server_timeout: Duration::from_millis(args.request_timeout_ms), + maximum_candidate_fmas: args.scheduler_quantum_fmas, + }; + match thread::Builder::new() + .name("tilemaxsim-reader".to_owned()) + .spawn(move || { + let _permit = ReaderPermit(reader_metrics); + read_and_enqueue( + connection, + &reader_sender, + reader_config, + reader_admission, + request_metrics, + reader_frame_admission, + ); + }) { + Ok(reader) => readers.push(reader), + Err(error) => { + metrics.active_connections.fetch_sub(1, Ordering::Release); + metrics.ready.store(false, Ordering::Release); + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + fatal_error = Some(error.into()); + break; + } + } + if args.once && accepted == 1 { + SHUTDOWN_REQUESTED.store(true, Ordering::Relaxed); + } + } + } + SHUTDOWN_REQUESTED.store(true, Ordering::Release); + drop(listener); + drop(tcp_listener); + for reader in readers { + if reader.join().is_err() { + eprintln!("tilemaxsim reader thread panicked during shutdown"); + } + } + drop(sender); + metrics.ready.store(false, Ordering::Release); + let scheduler_result = scheduler + .join() + .map_err(|_| anyhow!("TileMaxSim scheduler thread panicked")) + .and_then(|result| result); + if status_server.is_some_and(|status_server| status_server.join().is_err()) { + eprintln!("TileMaxSim status thread panicked during shutdown"); + } + if status_tcp_server.is_some_and(|status_server| status_server.join().is_err()) { + eprintln!("TileMaxSim TCP status thread panicked during shutdown"); + } + match fs::remove_file(&args.socket) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + if status_failed { + bail!("TileMaxSim status server exited unexpectedly"); + } + if scheduler_failed { + scheduler_result?; + bail!("TileMaxSim scheduler exited unexpectedly"); + } + scheduler_result?; + if let Some(error) = fatal_error { + return Err(error); + } + Ok(()) +} + +enum ClientStream { + Unix(UnixStream), + Tcp(TcpStream), +} + +impl ClientStream { + fn set_read_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_read_timeout(timeout), + Self::Tcp(stream) => stream.set_read_timeout(timeout), + } + } + + fn set_write_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_write_timeout(timeout), + Self::Tcp(stream) => stream.set_write_timeout(timeout), + } + } +} + +impl Read for ClientStream { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.read(buffer), + Self::Tcp(stream) => stream.read(buffer), + } + } +} + +impl Write for ClientStream { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.write(buffer), + Self::Tcp(stream) => stream.write(buffer), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.flush(), + Self::Tcp(stream) => stream.flush(), + } + } +} + +impl AsRawFd for ClientStream { + fn as_raw_fd(&self) -> std::os::fd::RawFd { + match self { + Self::Unix(stream) => stream.as_raw_fd(), + Self::Tcp(stream) => stream.as_raw_fd(), + } + } +} + +struct Work { + request: protocol::Request, + connection: ClientStream, + accepted_at: Instant, + deadline: Instant, + next_candidate: usize, + results: Vec<(u32, f32)>, + gpu_elapsed: Duration, + _pending_permit: PendingPermit, + _frame_permit: BytePermit, +} + +struct SchedulerConfig { + policy: SchedulerPolicy, + priority_aging: Duration, + priority_band: i32, + batch_window: Duration, + quantum_candidates: usize, + quantum_tokens: u64, + quantum_fmas: u64, + quantum_io_bytes: usize, + socket_io_timeout: Duration, + tenant_weights: std::collections::HashMap, +} + +#[derive(Default)] +struct RuntimeMetrics { + ready: AtomicBool, + max_connections: usize, + max_pending_requests: usize, + max_tenant_pending_requests: usize, + max_inflight_request_bytes: usize, + active_connections: AtomicUsize, + inflight_request_bytes: AtomicUsize, + pending_requests: AtomicUsize, + pending_tenants: AtomicUsize, + scheduler_depth: AtomicUsize, + scheduler_depth_high_water: AtomicUsize, + gpu_active: AtomicUsize, + completed: AtomicU64, + failed: AtomicU64, + timed_out: AtomicU64, + disconnected: AtomicU64, + rejected_connections: AtomicU64, + rejected_frame_bytes: AtomicU64, + rejected_queue_global: AtomicU64, + rejected_tenant: AtomicU64, + frame_read_failures: AtomicU64, + invalid_requests: AtomicU64, + gpu_failures: AtomicU64, + scheduler_failures: AtomicU64, + reload_succeeded: AtomicU64, + reload_failed: AtomicU64, + timeout_before_enqueue: AtomicU64, + timeout_in_queue: AtomicU64, + timeout_before_execution: AtomicU64, + timeout_during_execution: AtomicU64, + gpu_quantums: AtomicU64, + scheduler_requeues: AtomicU64, + admitted_priority_negative: AtomicU64, + admitted_priority_zero: AtomicU64, + admitted_priority_positive: AtomicU64, + candidates_scored: AtomicU64, + document_rows_scored: AtomicU64, + latency_observations: AtomicU64, + total_latency_us: AtomicU64, + gpu_latency_us: AtomicU64, + engine: Mutex, +} + +impl RuntimeMetrics { + fn new( + max_connections: usize, + max_pending_requests: usize, + max_tenant_pending_requests: usize, + max_inflight_request_bytes: usize, + ) -> Self { + Self { + max_connections, + max_pending_requests, + max_tenant_pending_requests, + max_inflight_request_bytes, + ..Self::default() + } + } + + fn update_engine(&self, status: EngineStatus) { + *self + .engine + .lock() + .unwrap_or_else(|error| error.into_inner()) = status; + } + + fn update_scheduler_depth(&self, depth: usize) { + self.scheduler_depth.store(depth, Ordering::Relaxed); + self.scheduler_depth_high_water + .fetch_max(depth, Ordering::Relaxed); + } + + fn observe_latency(&self, total: Duration, gpu: Duration) { + self.latency_observations.fetch_add(1, Ordering::Relaxed); + saturating_atomic_add(&self.total_latency_us, duration_micros(total)); + saturating_atomic_add(&self.gpu_latency_us, duration_micros(gpu)); + } +} + +fn duration_micros(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + +fn saturating_atomic_add(counter: &AtomicU64, value: u64) { + counter + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_add(value)) + }) + .ok(); +} + +struct ReaderPermit(Arc); + +impl Drop for ReaderPermit { + fn drop(&mut self) { + self.0.active_connections.fetch_sub(1, Ordering::Release); + } +} + +struct PendingState { + total: usize, + tenants: HashMap, +} + +struct PendingAdmission { + state: Mutex, + max_total: usize, + max_tenant: usize, + metrics: Arc, +} + +struct ByteAdmission { + used: AtomicUsize, + maximum: usize, + metrics: Arc, +} + +struct BytePermit { + admission: Arc, + bytes: usize, +} + +impl ByteAdmission { + fn new(maximum: usize, metrics: Arc) -> Self { + Self { + used: AtomicUsize::new(0), + maximum, + metrics, + } + } + + fn try_acquire(self: &Arc, bytes: usize) -> Option { + let result = self + .used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current + .checked_add(bytes) + .filter(|next| *next <= self.maximum) + }) + .ok(); + if result.is_none() { + self.metrics + .rejected_frame_bytes + .fetch_add(1, Ordering::Relaxed); + return None; + } + self.metrics + .inflight_request_bytes + .fetch_add(bytes, Ordering::Relaxed); + Some(BytePermit { + admission: Arc::clone(self), + bytes, + }) + } +} + +impl Drop for BytePermit { + fn drop(&mut self) { + self.admission.used.fetch_sub(self.bytes, Ordering::Release); + self.admission + .metrics + .inflight_request_bytes + .fetch_sub(self.bytes, Ordering::Release); + } +} + +#[derive(Debug)] +enum AdmissionRejection { + Global, + Tenant, +} + +struct PendingPermit { + admission: Arc, + tenant: String, +} + +impl PendingAdmission { + fn new(max_total: usize, max_tenant: usize, metrics: Arc) -> Self { + Self { + state: Mutex::new(PendingState { + total: 0, + tenants: HashMap::new(), + }), + max_total, + max_tenant, + metrics, + } + } + + fn try_acquire(self: &Arc, tenant: &str) -> Result { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.total >= self.max_total { + return Err(AdmissionRejection::Global); + } + if state.tenants.get(tenant).copied().unwrap_or(0) >= self.max_tenant { + return Err(AdmissionRejection::Tenant); + } + state.total += 1; + *state.tenants.entry(tenant.to_owned()).or_default() += 1; + self.metrics + .pending_requests + .store(state.total, Ordering::Relaxed); + self.metrics + .pending_tenants + .store(state.tenants.len(), Ordering::Relaxed); + Ok(PendingPermit { + admission: Arc::clone(self), + tenant: tenant.to_owned(), + }) + } +} + +impl Drop for PendingPermit { + fn drop(&mut self) { + let mut state = self + .admission + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + state.total = state.total.saturating_sub(1); + if let Some(count) = state.tenants.get_mut(&self.tenant) { + *count = count.saturating_sub(1); + if *count == 0 { + state.tenants.remove(&self.tenant); + } + } + self.admission + .metrics + .pending_requests + .store(state.total, Ordering::Relaxed); + self.admission + .metrics + .pending_tenants + .store(state.tenants.len(), Ordering::Relaxed); + } +} + +fn try_acquire_reader(counter: &AtomicUsize, maximum: usize) -> bool { + counter + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current < maximum).then_some(current + 1) + }) + .is_ok() +} + +fn reap_readers(readers: &mut Vec>) { + let mut index = 0; + while index < readers.len() { + if readers[index].is_finished() { + let reader = readers.swap_remove(index); + if reader.join().is_err() { + eprintln!("tilemaxsim reader thread panicked"); + } + } else { + index += 1; + } + } +} + +#[derive(Clone, Copy)] +struct ReaderConfig { + maximum: usize, + io_timeout: Duration, + server_timeout: Duration, + maximum_candidate_fmas: u64, +} + +fn read_and_enqueue( + mut connection: ClientStream, + sender: &mpsc::SyncSender, + config: ReaderConfig, + pending_admission: Arc, + metrics: Arc, + frame_admission: Arc, +) { + let accepted_at = Instant::now(); + if let Err(error) = connection.set_read_timeout(Some(config.io_timeout)) { + eprintln!("cannot configure TileMaxSim socket read timeout: {error}"); + return; + } + if let Err(error) = connection.set_write_timeout(Some(config.io_timeout)) { + eprintln!("cannot configure TileMaxSim socket write timeout: {error}"); + return; + } + let (frame, frame_permit) = + match read_request(&mut connection, config.maximum, &frame_admission) { + Ok(frame) => frame, + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.frame_read_failures.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(VERSION_EXTERNAL, 0, 1, &format!("{error:#}")), + ); + return; + } + }; + let version = header_version(&frame); + let request_id = header_request_id(&frame); + let request = match protocol::parse(&frame) { + Ok(request) => request, + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.invalid_requests.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(version, request_id, 1, &format!("{error:#}")), + ); + return; + } + }; + if request.candidates.iter().any(|candidate| { + candidate_fmas(request.query_rows, request.dimension, candidate.rows) + > config.maximum_candidate_fmas + }) { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.invalid_requests.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 1, + "one candidate exceeds the configured CUDA kernel work limit", + ), + ); + return; + } + let client_timeout = if request.timeout_ms == 0 { + config.server_timeout + } else { + Duration::from_millis(u64::from(request.timeout_ms)).min(config.server_timeout) + }; + let deadline = accepted_at + .checked_add(client_timeout) + .unwrap_or(accepted_at); + if deadline <= Instant::now() { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_before_enqueue + .fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 2, + "request deadline expired before enqueue", + ), + ); + return; + } + let pending_permit = match pending_admission.try_acquire(&request.tenant) { + Ok(permit) => permit, + Err(AdmissionRejection::Global) => { + metrics + .rejected_queue_global + .fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure(version, request_id, 2, "TileMaxSim scheduler queue is full"), + ); + return; + } + Err(AdmissionRejection::Tenant) => { + metrics.rejected_tenant.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut connection, + &protocol::failure( + version, + request_id, + 2, + "tenant TileMaxSim queue limit exceeded", + ), + ); + return; + } + }; + match sender.try_send(Work { + request, + connection, + accepted_at, + deadline, + next_candidate: 0, + results: Vec::new(), + gpu_elapsed: Duration::ZERO, + _pending_permit: pending_permit, + _frame_permit: frame_permit, + }) { + Ok(()) => {} + Err(mpsc::TrySendError::Full(mut work)) => { + metrics + .rejected_queue_global + .fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut work.connection, + &protocol::failure( + work.request.protocol_version, + work.request.request_id, + 2, + "TileMaxSim scheduler queue is full", + ), + ); + } + Err(mpsc::TrySendError::Disconnected(mut work)) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.scheduler_failures.fetch_add(1, Ordering::Relaxed); + write_response_nonfatal( + &mut work.connection, + &protocol::failure( + work.request.protocol_version, + work.request.request_id, + 3, + "TileMaxSim scheduler is unavailable", + ), + ); + } + } +} + +fn run_scheduler( + mut engine: Engine, + receiver: mpsc::Receiver, + config: SchedulerConfig, + reload: Arc, + metrics: Arc, +) -> Result<()> { + let mut queue = RequestQueue::new( + config.policy, + config.priority_aging, + config.priority_band, + config.tenant_weights.clone(), + ); + let mut channel_open = true; + while channel_open || queue.len() > 0 { + if reload.swap(false, Ordering::AcqRel) { + match engine.reload_shards() { + Ok(()) => { + metrics.reload_succeeded.fetch_add(1, Ordering::Relaxed); + metrics.update_engine(engine.status_snapshot()); + println!( + "{}", + serde_json::json!({"event": "tilemaxsim_rust_shards_reloaded"}) + ); + } + Err(error) => { + metrics.reload_failed.fetch_add(1, Ordering::Relaxed); + eprintln!("TileMaxSim shard reload rejected: {error:#}"); + } + } + } + + if queue.len() == 0 && channel_open { + match receiver.recv_timeout(Duration::from_millis(50)) { + Ok(work) => enqueue_work(&mut queue, work, &config, &engine, &metrics), + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => channel_open = false, + } + } + if channel_open { + let batching_deadline = Instant::now() + config.batch_window; + loop { + let remaining = batching_deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match receiver.recv_timeout(remaining) { + Ok(work) => enqueue_work(&mut queue, work, &config, &engine, &metrics), + Err(mpsc::RecvTimeoutError::Timeout) => break, + Err(mpsc::RecvTimeoutError::Disconnected) => { + channel_open = false; + break; + } + } + } + } + + let now = Instant::now(); + for mut expired in queue.drain_expired(now) { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics.timeout_in_queue.fetch_add(1, Ordering::Relaxed); + metrics.observe_latency(expired.payload.accepted_at.elapsed(), Duration::ZERO); + write_response_nonfatal( + &mut expired.payload.connection, + &protocol::failure( + expired.payload.request.protocol_version, + expired.payload.request.request_id, + 2, + "request deadline expired in scheduler queue", + ), + ); + } + metrics.update_scheduler_depth(queue.len()); + let Some(scheduled) = queue.pop(Instant::now()) else { + continue; + }; + metrics.update_scheduler_depth(queue.len()); + if peer_disconnected(&scheduled.payload.connection) { + metrics.disconnected.fetch_add(1, Ordering::Relaxed); + metrics.observe_latency(scheduled.payload.accepted_at.elapsed(), Duration::ZERO); + continue; + } + let started = Instant::now(); + let mut work = scheduled.payload; + let request_id = work.request.request_id; + let version = work.request.protocol_version; + let tenant = work.request.tenant.clone(); + let priority = work.request.priority; + let response = if work.deadline <= started { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_before_execution + .fetch_add(1, Ordering::Relaxed); + Some(protocol::failure( + version, + request_id, + 2, + "request deadline expired before execution", + )) + } else { + let end = next_quantum_end(&work, &config, &engine); + let quantum = protocol::Request { + protocol_version: work.request.protocol_version, + request_id: work.request.request_id, + tenant: work.request.tenant.clone(), + priority: work.request.priority, + timeout_ms: work.request.timeout_ms, + query_rows: work.request.query_rows, + dimension: work.request.dimension, + dtype: work.request.dtype, + query: work.request.query.clone(), + candidates: work.request.candidates[work.next_candidate..end].to_vec(), + }; + let quantum_started = Instant::now(); + metrics.gpu_quantums.fetch_add(1, Ordering::Relaxed); + saturating_atomic_add( + &metrics.candidates_scored, + u64::try_from(end - work.next_candidate).unwrap_or(u64::MAX), + ); + saturating_atomic_add( + &metrics.document_rows_scored, + quantum + .candidates + .iter() + .map(|candidate| u64::from(candidate.rows)) + .sum(), + ); + metrics.gpu_active.store(1, Ordering::Relaxed); + let score_result = engine.score(&quantum); + metrics.gpu_active.store(0, Ordering::Relaxed); + metrics.update_engine(engine.status_snapshot()); + match score_result { + Ok(results) => { + work.gpu_elapsed += quantum_started.elapsed(); + work.results.extend(results); + work.next_candidate = end; + if work.deadline <= Instant::now() { + metrics.timed_out.fetch_add(1, Ordering::Relaxed); + metrics + .timeout_during_execution + .fetch_add(1, Ordering::Relaxed); + Some(protocol::failure( + version, + request_id, + 2, + "request deadline expired during GPU execution", + )) + } else if end < work.request.candidates.len() { + metrics.scheduler_requeues.fetch_add(1, Ordering::Relaxed); + let cost = estimated_next_work(&work, &config, &engine); + queue.push(Scheduled::new( + tenant.clone(), + priority, + cost, + work.accepted_at, + work.deadline, + work, + )); + metrics.update_scheduler_depth(queue.len()); + continue; + } else { + metrics.completed.fetch_add(1, Ordering::Relaxed); + Some(protocol::success(version, request_id, &work.results)) + } + } + Err(error) => { + metrics.failed.fetch_add(1, Ordering::Relaxed); + metrics.gpu_failures.fetch_add(1, Ordering::Relaxed); + work.gpu_elapsed += quantum_started.elapsed(); + let diagnostic = format!("{error:#}"); + let failure = protocol::failure(version, request_id, 3, &diagnostic); + if is_fatal_cuda_diagnostic(&diagnostic) { + // CUDA execution/context failures are not ordinary bad + // requests. Stop advertising readiness and terminate + // the scheduler after best-effort notification so the + // service supervisor can recreate the CUDA context. + metrics.ready.store(false, Ordering::Release); + write_response_nonfatal(&mut work.connection, &failure); + metrics.observe_latency(work.accepted_at.elapsed(), work.gpu_elapsed); + return Err(anyhow!("fatal CUDA failure: {diagnostic}")); + } + Some(failure) + } + } + }; + let Some(response) = response else { + continue; + }; + work.connection + .set_write_timeout(Some(config.socket_io_timeout)) + .ok(); + write_response_nonfatal(&mut work.connection, &response); + metrics.observe_latency(work.accepted_at.elapsed(), work.gpu_elapsed); + println!( + "{}", + serde_json::json!({ + "event": "tilemaxsim_rust_request", + "request_id": request_id, + "tenant_hash": tenant_hash(&tenant), + "priority": priority, + "total_ms": work.accepted_at.elapsed().as_secs_f64() * 1000.0, + "gpu_ms": work.gpu_elapsed.as_secs_f64() * 1000.0, + "queue_ms": work.accepted_at.elapsed().saturating_sub(work.gpu_elapsed).as_secs_f64() * 1000.0, + "queue_depth": queue.len(), + "cache": engine.status_json(), + }) + ); + } + Ok(()) +} + +fn enqueue_work( + queue: &mut RequestQueue, + work: Work, + config: &SchedulerConfig, + engine: &Engine, + metrics: &RuntimeMetrics, +) { + match work.request.priority.cmp(&0) { + std::cmp::Ordering::Less => metrics + .admitted_priority_negative + .fetch_add(1, Ordering::Relaxed), + std::cmp::Ordering::Equal => metrics + .admitted_priority_zero + .fetch_add(1, Ordering::Relaxed), + std::cmp::Ordering::Greater => metrics + .admitted_priority_positive + .fetch_add(1, Ordering::Relaxed), + }; + let cost = estimated_next_work(&work, config, engine); + queue.push(Scheduled::new( + work.request.tenant.clone(), + work.request.priority, + cost, + work.accepted_at, + work.deadline, + work, + )); + metrics.update_scheduler_depth(queue.len()); +} + +fn tenant_hash(tenant: &str) -> String { + let digest = Sha256::digest(tenant.as_bytes()); + digest[..8] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn is_fatal_cuda_diagnostic(diagnostic: &str) -> bool { + [ + "cudaSetDevice", + "cudaMemcpy", + "cudaStream", + "CUDA workspace initialization", + "TileMaxSim CUDA execution", + "GPU upload worker panicked", + "GPU worker panicked", + ] + .iter() + .any(|marker| diagnostic.contains(marker)) +} + +fn next_quantum_end(work: &Work, config: &SchedulerConfig, engine: &Engine) -> usize { + let compute_end = quantum_end( + &work.request.candidates, + work.next_candidate, + config.quantum_candidates, + config.quantum_tokens, + config.quantum_fmas, + work.request.query_rows, + work.request.dimension, + ); + io_quantum_end( + &work.request.candidates, + work.next_candidate, + compute_end, + config.quantum_io_bytes, + |candidate| engine.gpu_miss_bytes(candidate), + ) +} + +fn io_quantum_end( + candidates: &[protocol::Descriptor], + start: usize, + compute_end: usize, + maximum_io_bytes: usize, + mut miss_bytes: impl FnMut(&protocol::Descriptor) -> u64, +) -> usize { + let maximum_io_bytes = u64::try_from(maximum_io_bytes).unwrap_or(u64::MAX); + let mut end = start; + let mut bytes = 0_u64; + while end < compute_end { + let candidate_bytes = miss_bytes(&candidates[end]); + if end > start && bytes.saturating_add(candidate_bytes) > maximum_io_bytes { + break; + } + bytes = bytes.saturating_add(candidate_bytes); + end += 1; + } + end.max((start + 1).min(compute_end)) +} + +fn quantum_end( + candidates: &[protocol::Descriptor], + start: usize, + maximum_candidates: usize, + maximum_tokens: u64, + maximum_fmas: u64, + query_rows: u32, + dimension: u32, +) -> usize { + let mut end = start; + let mut tokens = 0_u64; + let mut fmas = 0_u64; + while end < candidates.len() && end - start < maximum_candidates { + let rows = u64::from(candidates[end].rows); + let candidate_fmas = candidate_fmas(query_rows, dimension, candidates[end].rows); + if end > start + && (tokens.saturating_add(rows) > maximum_tokens + || fmas.saturating_add(candidate_fmas) > maximum_fmas) + { + break; + } + tokens = tokens.saturating_add(rows); + fmas = fmas.saturating_add(candidate_fmas); + end += 1; + } + end +} + +fn candidate_fmas(query_rows: u32, dimension: u32, document_rows: u32) -> u64 { + u64::from(query_rows) + .saturating_mul(u64::from(document_rows)) + .saturating_mul(u64::from(dimension)) +} + +fn estimated_next_work(work: &Work, config: &SchedulerConfig, engine: &Engine) -> u64 { + let end = next_quantum_end(work, config, engine); + work.request.candidates[work.next_candidate..end] + .iter() + .map(|candidate| { + candidate_fmas( + work.request.query_rows, + work.request.dimension, + candidate.rows, + ) + }) + .fold(0_u64, u64::saturating_add) + .max(1) +} + +fn peer_disconnected(connection: &ClientStream) -> bool { + let mut byte = 0_u8; + let result = unsafe { + libc::recv( + connection.as_raw_fd(), + (&raw mut byte).cast(), + 1, + libc::MSG_PEEK | libc::MSG_DONTWAIT, + ) + }; + if result == 0 { + return true; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + return !matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ); + } + false +} + +fn write_response_nonfatal(connection: &mut ClientStream, response: &[u8]) { + if let Err(error) = connection.write_all(response) { + eprintln!("TileMaxSim response write failed without stopping daemon: {error}"); + } +} + +fn run_status_server(listener: UnixListener, path: PathBuf, metrics: Arc) { + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + match listener.accept() { + Ok((mut connection, _)) => { + connection + .set_read_timeout(Some(Duration::from_millis(250))) + .ok(); + connection + .set_write_timeout(Some(Duration::from_millis(250))) + .ok(); + handle_status_connection(&mut connection, &metrics); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(20)); + } + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Err(error) => { + eprintln!("TileMaxSim status socket failed: {error}"); + break; + } + } + } + drop(listener); + match fs::remove_file(&path) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => eprintln!("cannot remove TileMaxSim status socket: {error}"), + } +} + +fn run_tcp_status_server(listener: TcpListener, metrics: Arc) { + while !SHUTDOWN_REQUESTED.load(Ordering::Relaxed) { + match listener.accept() { + Ok((mut connection, _)) => { + configure_tcp_status_connection(&connection); + handle_status_connection(&mut connection, &metrics); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(20)); + } + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {} + Err(error) => { + eprintln!("TileMaxSim TCP status listener failed: {error}"); + break; + } + } + } +} + +fn configure_tcp_status_connection(connection: &TcpStream) { + connection + .set_read_timeout(Some(Duration::from_millis(250))) + .ok(); + connection + .set_write_timeout(Some(Duration::from_millis(250))) + .ok(); +} + +fn handle_status_connection(connection: &mut (impl Read + Write), metrics: &RuntimeMetrics) { + let mut request = [0_u8; 1024]; + let Ok(count) = connection.read(&mut request) else { + return; + }; + let request = String::from_utf8_lossy(&request[..count]); + let (status, content_type, body) = if request.starts_with("GET /livez ") { + ( + "200 OK", + "application/json", + serde_json::json!({"live": true}).to_string(), + ) + } else if request.starts_with("GET /healthz ") { + let ready = metrics.ready.load(Ordering::Acquire); + ( + if ready { + "200 OK" + } else { + "503 Service Unavailable" + }, + "application/json", + serde_json::json!({"ready": ready}).to_string(), + ) + } else if request.starts_with("GET /metrics ") { + ( + "200 OK", + "text/plain; version=0.0.4", + render_metrics(metrics), + ) + } else { + ("404 Not Found", "text/plain", "not found\n".to_owned()) + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + if let Err(error) = connection.write_all(response.as_bytes()) { + eprintln!("TileMaxSim status response failed: {error}"); + } +} + +fn render_metrics(metrics: &RuntimeMetrics) -> String { + let mut output = String::with_capacity(8 * 1024); + writeln!( + output, + "# HELP tilemaxsim_ready Whether the daemon is ready to accept work." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_ready gauge").unwrap(); + writeln!( + output, + "tilemaxsim_ready {}", + usize::from(metrics.ready.load(Ordering::Relaxed)) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_connections Active reader connections and configured limit." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_connections gauge").unwrap(); + writeln!( + output, + "tilemaxsim_connections{{kind=\"active\"}} {}", + metrics.active_connections.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_connections{{kind=\"limit\"}} {}", + metrics.max_connections + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_pending_requests Requests admitted but not yet completed." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_pending_requests gauge").unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"current\"}} {}", + metrics.pending_requests.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"global_limit\"}} {}", + metrics.max_pending_requests + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_pending_requests{{kind=\"per_tenant_limit\"}} {}", + metrics.max_tenant_pending_requests + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_pending_tenants Tenants with admitted requests." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_pending_tenants gauge").unwrap(); + writeln!( + output, + "tilemaxsim_pending_tenants {}", + metrics.pending_tenants.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_inflight_request_bytes Encoded request frames retained in memory." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_inflight_request_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_inflight_request_bytes{{kind=\"current\"}} {}", + metrics.inflight_request_bytes.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_inflight_request_bytes{{kind=\"limit\"}} {}", + metrics.max_inflight_request_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_scheduler_queue_depth Requests waiting for a GPU quantum." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_queue_depth gauge").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_queue_depth{{kind=\"current\"}} {}", + metrics.scheduler_depth.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_queue_depth{{kind=\"high_water\"}} {}", + metrics.scheduler_depth_high_water.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_active Whether a CUDA quantum is executing." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_active gauge").unwrap(); + writeln!( + output, + "tilemaxsim_gpu_active {}", + metrics.gpu_active.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_requests_total Terminal request outcomes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_requests_total counter").unwrap(); + for (outcome, value) in [ + ("completed", metrics.completed.load(Ordering::Relaxed)), + ("failed", metrics.failed.load(Ordering::Relaxed)), + ("timeout", metrics.timed_out.load(Ordering::Relaxed)), + ("disconnected", metrics.disconnected.load(Ordering::Relaxed)), + ] { + writeln!( + output, + "tilemaxsim_requests_total{{outcome=\"{outcome}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_admission_rejections_total Bounded admission rejection reasons." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_admission_rejections_total counter" + ) + .unwrap(); + for (reason, value) in [ + ( + "connections", + metrics.rejected_connections.load(Ordering::Relaxed), + ), + ( + "frame_bytes", + metrics.rejected_frame_bytes.load(Ordering::Relaxed), + ), + ( + "queue_global", + metrics.rejected_queue_global.load(Ordering::Relaxed), + ), + ( + "queue_tenant", + metrics.rejected_tenant.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_admission_rejections_total{{reason=\"{reason}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_failures_total Internal failure categories." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_failures_total counter").unwrap(); + for (reason, value) in [ + ( + "frame_io_or_validation", + metrics.frame_read_failures.load(Ordering::Relaxed), + ), + ("request", metrics.invalid_requests.load(Ordering::Relaxed)), + ("gpu", metrics.gpu_failures.load(Ordering::Relaxed)), + ( + "scheduler", + metrics.scheduler_failures.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_failures_total{{reason=\"{reason}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_timeouts_total Request timeout phase." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_timeouts_total counter").unwrap(); + for (phase, value) in [ + ( + "before_enqueue", + metrics.timeout_before_enqueue.load(Ordering::Relaxed), + ), + ("queue", metrics.timeout_in_queue.load(Ordering::Relaxed)), + ( + "before_execution", + metrics.timeout_before_execution.load(Ordering::Relaxed), + ), + ( + "gpu", + metrics.timeout_during_execution.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_timeouts_total{{phase=\"{phase}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_requests_admitted_total Requests admitted by bounded priority class." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_requests_admitted_total counter").unwrap(); + for (priority_class, value) in [ + ( + "negative", + metrics.admitted_priority_negative.load(Ordering::Relaxed), + ), + ( + "zero", + metrics.admitted_priority_zero.load(Ordering::Relaxed), + ), + ( + "positive", + metrics.admitted_priority_positive.load(Ordering::Relaxed), + ), + ] { + writeln!( + output, + "tilemaxsim_requests_admitted_total{{priority_class=\"{priority_class}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_scheduler_quantums_total CUDA scheduling quanta executed." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_quantums_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_quantums_total {}", + metrics.gpu_quantums.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_scheduler_requeues_total Cooperative quantum yields requeued." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_scheduler_requeues_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_scheduler_requeues_total {}", + metrics.scheduler_requeues.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_candidates_scored_total Candidate tensors submitted to CUDA." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_candidates_scored_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_candidates_scored_total {}", + metrics.candidates_scored.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_document_rows_scored_total Document token rows submitted to CUDA." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_document_rows_scored_total counter" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_document_rows_scored_total {}", + metrics.document_rows_scored.load(Ordering::Relaxed) + ) + .unwrap(); + let observations = metrics.latency_observations.load(Ordering::Relaxed); + let total_us = metrics.total_latency_us.load(Ordering::Relaxed); + let gpu_us = metrics.gpu_latency_us.load(Ordering::Relaxed); + writeln!( + output, + "# HELP tilemaxsim_request_duration_seconds Request wall-clock duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_request_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_request_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_request_duration_seconds_sum {}", + total_us as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_duration_seconds CUDA execution duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_gpu_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_gpu_duration_seconds_sum {}", + gpu_us as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_queue_duration_seconds Non-CUDA request duration summary." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_queue_duration_seconds summary").unwrap(); + writeln!( + output, + "tilemaxsim_queue_duration_seconds_count {observations}" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_queue_duration_seconds_sum {}", + total_us.saturating_sub(gpu_us) as f64 / 1_000_000.0 + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_shard_reloads_total Immutable shard metadata reload outcomes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_shard_reloads_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_shard_reloads_total{{outcome=\"success\"}} {}", + metrics.reload_succeeded.load(Ordering::Relaxed) + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_shard_reloads_total{{outcome=\"failed\"}} {}", + metrics.reload_failed.load(Ordering::Relaxed) + ) + .unwrap(); + + let engine = metrics + .engine + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_bytes GPU tensor-cache byte accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_bytes gauge").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_entries GPU tensor-cache entry accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_entries gauge").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_cache_events_total GPU tensor-cache cumulative events." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_cache_events_total counter").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_h2d_batches_total Host-to-device transfer batches." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_h2d_batches_total counter").unwrap(); + writeln!( + output, + "# HELP tilemaxsim_gpu_h2d_bytes_total Host-to-device transfer bytes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_gpu_h2d_bytes_total counter").unwrap(); + for device in &engine.devices { + for (kind, value) in [ + ("capacity", device.capacity_bytes), + ("free", device.free_bytes), + ("largest_free_extent", device.largest_free_extent_bytes), + ("allocated", device.allocated_bytes), + ("payload", device.payload_bytes), + ("internal_waste", device.internal_waste_bytes), + ("pinned", device.pinned_bytes), + ("block", device.block_bytes), + ] { + writeln!( + output, + "tilemaxsim_gpu_cache_bytes{{slot=\"{}\",device=\"{}\",kind=\"{kind}\"}} {value}", + device.slot, device.device + ) + .unwrap(); + } + for (kind, value) in [ + ("total", device.entries), + ("pinned", device.pinned_entries), + ("tenants", device.tenants), + ] { + writeln!( + output, + "tilemaxsim_gpu_cache_entries{{slot=\"{}\",device=\"{}\",kind=\"{kind}\"}} {value}", + device.slot, device.device + ) + .unwrap(); + } + for (event, value) in [ + ("hit", device.hits), + ("miss", device.misses), + ("eviction", device.evictions), + ("admission_rejection", device.admission_rejections), + ] { + writeln!(output, "tilemaxsim_gpu_cache_events_total{{slot=\"{}\",device=\"{}\",event=\"{event}\"}} {value}", device.slot, device.device).unwrap(); + } + writeln!( + output, + "tilemaxsim_gpu_h2d_batches_total{{slot=\"{}\",device=\"{}\"}} {}", + device.slot, device.device, device.h2d_batches + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_gpu_h2d_bytes_total{{slot=\"{}\",device=\"{}\"}} {}", + device.slot, device.device, device.h2d_bytes + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_host_cache_bytes Host tensor-cache byte accounting." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_bytes{{kind=\"capacity\"}} {}", + engine.host.capacity_bytes + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_bytes{{kind=\"used\"}} {}", + engine.host.used_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_host_cache_entries Host tensor-cache entries and active tenants." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_entries gauge").unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_entries{{kind=\"entries\"}} {}", + engine.host.entries + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_host_cache_entries{{kind=\"tenants\"}} {}", + engine.host.tenants + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_host_cache_events_total Host tensor-cache cumulative events." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_host_cache_events_total counter").unwrap(); + for (event, value) in [ + ("hit", engine.host.hits), + ("miss", engine.host.misses), + ("eviction", engine.host.evictions), + ("admission_rejection", engine.host.admission_rejections), + ] { + writeln!( + output, + "tilemaxsim_host_cache_events_total{{event=\"{event}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_storage_read_calls_total Immutable tensor storage read calls." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_storage_read_calls_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_storage_read_calls_total {}", + engine.batch_read_calls + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_storage_read_bytes_total Immutable tensor storage read bytes." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_storage_read_bytes_total counter").unwrap(); + writeln!( + output, + "tilemaxsim_storage_read_bytes_total {}", + engine.batch_read_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_enabled Whether tensor I/O overlaps CUDA scoring." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_io_pipeline_enabled gauge").unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_enabled {}", + usize::from(engine.io_pipeline.overlap_enabled) + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_batch_bytes Configured logical tensor bytes per I/O stage." + ) + .unwrap(); + writeln!(output, "# TYPE tilemaxsim_io_pipeline_batch_bytes gauge").unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_batch_bytes {}", + engine.io_pipeline.batch_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_batches_total Tensor pipeline stage executions." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_batches_total counter" + ) + .unwrap(); + for (stage, value) in [ + ("resolve", engine.io_pipeline.resolve_batches), + ("overlap", engine.io_pipeline.overlap_cycles), + ] { + writeln!( + output, + "tilemaxsim_io_pipeline_batches_total{{stage=\"{stage}\"}} {value}" + ) + .unwrap(); + } + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_resolve_bytes_total Logical tensor payload passed through resolve stages." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_resolve_bytes_total counter" + ) + .unwrap(); + writeln!( + output, + "tilemaxsim_io_pipeline_resolve_bytes_total {}", + engine.io_pipeline.resolve_bytes + ) + .unwrap(); + writeln!( + output, + "# HELP tilemaxsim_io_pipeline_seconds_total Cumulative tensor pipeline stage and hidden time." + ) + .unwrap(); + writeln!( + output, + "# TYPE tilemaxsim_io_pipeline_seconds_total counter" + ) + .unwrap(); + for (stage, microseconds) in [ + ("resolve", engine.io_pipeline.resolve_microseconds), + ("upload", engine.io_pipeline.upload_microseconds), + ("compute", engine.io_pipeline.compute_microseconds), + ("overlap", engine.io_pipeline.overlap_microseconds), + ] { + writeln!( + output, + "tilemaxsim_io_pipeline_seconds_total{{stage=\"{stage}\"}} {}", + microseconds as f64 / 1_000_000.0 + ) + .unwrap(); + } + output +} + +fn read_request( + connection: &mut ClientStream, + maximum: usize, + frame_admission: &Arc, +) -> Result<(Vec, BytePermit)> { + let mut header = [0_u8; HEADER_BYTES]; + connection.read_exact(&mut header)?; + let body_bytes = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .context("request body does not fit this host")?; + let total = HEADER_BYTES + .checked_add(body_bytes) + .ok_or_else(|| anyhow!("request length overflow"))?; + if total > maximum { + bail!("request exceeds byte limit"); + } + let permit = frame_admission + .try_acquire(total) + .ok_or_else(|| anyhow!("in-flight TileMaxSim request byte budget is exhausted"))?; + let mut frame = Vec::with_capacity(total); + frame.extend_from_slice(&header); + frame.resize(total, 0); + connection.read_exact(&mut frame[HEADER_BYTES..])?; + Ok((frame, permit)) +} + +fn header_request_id(frame: &[u8]) -> u64 { + if frame.len() < HEADER_BYTES { + 0 + } else { + u64::from_le_bytes(frame[8..16].try_into().unwrap()) + } +} + +fn header_version(frame: &[u8]) -> u16 { + if frame.len() < HEADER_BYTES { + VERSION_EXTERNAL + } else { + match u16::from_le_bytes(frame[4..6].try_into().unwrap()) { + VERSION_SCHEDULED_EXTERNAL => VERSION_SCHEDULED_EXTERNAL, + _ => VERSION_EXTERNAL, + } + } +} + +fn acquire_instance_lock(socket: &Path) -> Result { + let lock_path = PathBuf::from(format!("{}.lock", socket.display())); + let mut lock = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(&lock_path) + .with_context(|| format!("cannot open instance lock {}", lock_path.display()))?; + if unsafe { libc::flock(lock.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } != 0 { + bail!( + "another TileMaxSim daemon holds the instance lock {}", + lock_path.display() + ); + } + lock.set_len(0)?; + writeln!(lock, "{}", std::process::id())?; + lock.flush()?; + Ok(lock) +} + +fn remove_stale_socket(path: &PathBuf) -> Result<()> { + let Ok(metadata) = fs::symlink_metadata(path) else { + return Ok(()); + }; + if !metadata.file_type().is_socket() { + bail!("refusing to remove non-socket path {}", path.display()); + } + if UnixStream::connect(path).is_ok() { + bail!( + "another TileMaxSim daemon is already listening on {}", + path.display() + ); + } + fs::remove_file(path)?; + Ok(()) +} + +#[derive(Deserialize)] +struct ResidentRecord { + tensor_ref: String, + tensor_rows: u32, + tensor_dim: u32, + tensor_dtype: String, + tensor_checksum: String, + canonical_bytes: Option, +} + +fn load_resident_manifests(values: &[(String, PathBuf)]) -> Result> { + let mut descriptors = Vec::new(); + for (contract, path) in values { + let file = fs::File::open(path) + .with_context(|| format!("cannot open resident manifest {}", path.display()))?; + for (line_number, line) in std::io::BufReader::new(file).lines().enumerate() { + let line = line?; + let record: ResidentRecord = serde_json::from_str(&line).with_context(|| { + format!( + "invalid resident manifest {}:{}", + path.display(), + line_number + 1 + ) + })?; + let digest = record + .tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("resident tensor reference is not SHA-256"))?; + if record.tensor_checksum != format!("sha256:{digest}") { + bail!("resident tensor reference and checksum disagree"); + } + let (dtype, scalar_bytes) = match record.tensor_dtype.as_str() { + "float16" => (2, 2), + "float32" => (1, 4), + _ => bail!("resident manifest has an unsupported tensor dtype"), + }; + let expected_bytes = + record.tensor_rows as usize * record.tensor_dim as usize * scalar_bytes; + if record + .canonical_bytes + .is_some_and(|bytes| bytes != expected_bytes) + { + bail!("resident canonical byte length disagrees with its shape"); + } + descriptors.push(protocol::Descriptor { + candidate_id: descriptors.len() as u32, + contract: contract.clone(), + digest: digest.to_owned(), + rows: record.tensor_rows, + dimension: record.tensor_dim, + dtype, + }); + } + } + if descriptors.is_empty() { + bail!("resident manifests contain no tensor descriptors"); + } + Ok(descriptors) +} + +#[cfg(test)] +mod tests { + use super::{ + ByteAdmission, PendingAdmission, RuntimeMetrics, candidate_fmas, io_quantum_end, + is_fatal_cuda_diagnostic, kib_to_bytes, quantum_end, render_metrics, tenant_hash, + }; + use crate::engine::{DeviceStatus, EngineStatus, IoPipelineStatus}; + use crate::protocol::Descriptor; + use crate::shard::HostCacheStatus; + use std::sync::Arc; + + #[test] + fn gpu_block_kib_is_converted_once() { + assert_eq!(kib_to_bytes(32), Ok(32 * 1024)); + assert!(kib_to_bytes(usize::MAX).is_err()); + } + + #[test] + fn scheduler_quantum_always_makes_progress_and_honours_both_limits() { + let descriptor = |rows| Descriptor { + candidate_id: rows, + contract: "model".to_owned(), + digest: "a".repeat(64), + rows, + dimension: 2, + dtype: 2, + }; + let candidates = vec![descriptor(60), descriptor(60), descriptor(60)]; + assert_eq!(quantum_end(&candidates, 0, 8, 100, 1_000, 5, 2), 1); + assert_eq!(quantum_end(&candidates, 0, 2, 1_000, 2_000, 5, 2), 2); + assert_eq!(quantum_end(&candidates, 2, 8, 10, 1_000, 5, 2), 3); + assert_eq!(candidate_fmas(5, 2, 60), 600); + assert_eq!(candidate_fmas(u32::MAX, u32::MAX, u32::MAX), u64::MAX); + } + + #[test] + fn scheduler_quantum_caps_l0_miss_bytes_but_keeps_hot_candidates_free() { + let descriptor = |candidate_id| Descriptor { + candidate_id, + contract: "model".to_owned(), + digest: format!("{candidate_id:064x}"), + rows: 1, + dimension: 2, + dtype: 2, + }; + let candidates = vec![descriptor(1), descriptor(2), descriptor(3)]; + let cold_end = io_quantum_end(&candidates, 0, 3, 8, |_| 5); + assert_eq!(cold_end, 1); + let mixed_end = io_quantum_end(&candidates, 0, 3, 8, |candidate| { + if candidate.candidate_id == 2 { 0 } else { 5 } + }); + assert_eq!(mixed_end, 2); + assert_eq!(io_quantum_end(&candidates, 2, 3, 1, |_| 5), 3); + } + + #[test] + fn pending_admission_is_globally_and_per_tenant_bounded() { + let metrics = Arc::new(RuntimeMetrics::default()); + let admission = Arc::new(PendingAdmission::new(2, 1, metrics)); + let first = admission.try_acquire("a").unwrap(); + assert!(admission.try_acquire("a").is_err()); + let second = admission.try_acquire("b").unwrap(); + assert!(admission.try_acquire("c").is_err()); + drop(first); + assert!(admission.try_acquire("a").is_ok()); + drop(second); + } + + #[test] + fn in_flight_frame_bytes_are_bounded_until_the_permit_drops() { + let metrics = Arc::new(RuntimeMetrics::default()); + let admission = Arc::new(ByteAdmission::new(100, metrics)); + let first = admission.try_acquire(60).unwrap(); + assert!(admission.try_acquire(41).is_none()); + let second = admission.try_acquire(40).unwrap(); + assert!(admission.try_acquire(1).is_none()); + drop(first); + assert!(admission.try_acquire(60).is_some()); + drop(second); + } + + #[test] + fn tenant_labels_are_stable_and_do_not_expose_raw_ids() { + assert_eq!(tenant_hash("tenant-a"), tenant_hash("tenant-a")); + assert_ne!(tenant_hash("tenant-a"), tenant_hash("tenant-b")); + assert!(!tenant_hash("tenant-a").contains("tenant")); + } + + #[test] + fn only_cuda_runtime_failures_force_supervised_restart() { + assert!(is_fatal_cuda_diagnostic( + "TileMaxSim CUDA execution: an illegal memory access was encountered" + )); + assert!(is_fatal_cuda_diagnostic( + "cudaMemcpyAsync(H2D): unspecified launch failure" + )); + assert!(!is_fatal_cuda_diagnostic( + "tensor is missing from immutable shards and objects" + )); + assert!(!is_fatal_cuda_diagnostic( + "TileMaxSim request exceeds the configured GPU workspace" + )); + } + + #[test] + fn prometheus_status_contains_only_bounded_scheduler_counters() { + let metrics = RuntimeMetrics::default(); + metrics + .ready + .store(true, std::sync::atomic::Ordering::Relaxed); + metrics + .completed + .store(7, std::sync::atomic::Ordering::Relaxed); + metrics.update_engine(EngineStatus { + devices: vec![DeviceStatus { + slot: 0, + device: 3, + capacity_bytes: 1_024, + free_bytes: 512, + largest_free_extent_bytes: 256, + ..DeviceStatus::default() + }], + host: HostCacheStatus { + capacity_bytes: 2_048, + used_bytes: 128, + ..HostCacheStatus::default() + }, + batch_read_calls: 4, + batch_read_bytes: 256, + io_pipeline: IoPipelineStatus { + overlap_enabled: true, + batch_bytes: 512, + resolve_batches: 2, + overlap_cycles: 1, + ..IoPipelineStatus::default() + }, + }); + let output = render_metrics(&metrics); + assert!(output.contains("tilemaxsim_ready 1")); + assert!(output.contains("outcome=\"completed\"} 7")); + assert!(output.contains( + "tilemaxsim_gpu_cache_bytes{slot=\"0\",device=\"3\",kind=\"capacity\"} 1024" + )); + assert!(output.contains("tilemaxsim_host_cache_bytes{kind=\"used\"} 128")); + assert!(output.contains("tilemaxsim_storage_read_bytes_total 256")); + assert!(output.contains("tilemaxsim_io_pipeline_enabled 1")); + assert!(output.contains("tilemaxsim_io_pipeline_batch_bytes 512")); + assert!(!output.contains("tenant-a")); + } +} diff --git a/services/tilemaxsimd/src/protocol.rs b/services/tilemaxsimd/src/protocol.rs new file mode 100644 index 00000000..559e00d1 --- /dev/null +++ b/services/tilemaxsimd/src/protocol.rs @@ -0,0 +1,347 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use anyhow::{Result, anyhow, bail}; +use std::collections::HashSet; + +pub const HEADER_BYTES: usize = 24; +pub const VERSION_EXTERNAL: u16 = 2; +pub const VERSION_SCHEDULED_EXTERNAL: u16 = 3; +const MAGIC: &[u8; 4] = b"VCTM"; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; + +#[derive(Clone, Debug)] +pub struct Descriptor { + pub candidate_id: u32, + pub contract: String, + pub digest: String, + pub rows: u32, + pub dimension: u32, + pub dtype: u8, +} + +#[derive(Clone, Debug)] +pub struct Request { + pub protocol_version: u16, + pub request_id: u64, + /// Logical scheduling domain. It is never used for authorization or + /// tensor lookup; the caller must already have applied hard ACL filters. + pub tenant: String, + /// Higher values run first within the scheduler's fairness policy. + pub priority: i32, + /// Client-supplied end-to-end budget. Zero is used only by legacy v2. + pub timeout_ms: u32, + pub query_rows: u32, + pub dimension: u32, + pub dtype: u8, + pub query: Vec, + pub candidates: Vec, +} + +struct Reader<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Reader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn take(&mut self, count: usize) -> Result<&'a [u8]> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| anyhow!("request overflow"))?; + if end > self.bytes.len() { + bail!("truncated request"); + } + let result = &self.bytes[self.offset..end]; + self.offset = end; + Ok(result) + } + + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn i32(&mut self) -> Result { + Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn text(&mut self, count: usize, maximum: usize, name: &str) -> Result { + if count == 0 || count > maximum { + bail!("invalid {name} length"); + } + let value = std::str::from_utf8(self.take(count)?)?; + if value.chars().any(|character| character.is_control()) { + bail!("{name} contains control characters"); + } + Ok(value.to_owned()) + } + + fn finish(self) -> Result<()> { + if self.offset != self.bytes.len() { + bail!("trailing request bytes"); + } + Ok(()) + } +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + if rows == 0 || dimension == 0 || dimension > 60_000 { + bail!("invalid tensor shape"); + } + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape is too large")) +} + +fn validate_finite(payload: &[u8], dtype: u8) -> Result<()> { + if dtype == 1 { + for scalar in payload.chunks_exact(4) { + if !f32::from_bits(u32::from_le_bytes(scalar.try_into().unwrap())).is_finite() { + bail!("query contains a non-finite value"); + } + } + } else { + for scalar in payload.chunks_exact(2) { + let bits = u16::from_le_bytes(scalar.try_into().unwrap()); + if bits & 0x7c00 == 0x7c00 { + bail!("query contains a non-finite value"); + } + } + } + Ok(()) +} + +pub fn parse(frame: &[u8]) -> Result { + if frame.len() < HEADER_BYTES { + bail!("truncated request header"); + } + if &frame[..4] != MAGIC { + bail!("invalid protocol magic"); + } + let version = u16::from_le_bytes(frame[4..6].try_into().unwrap()); + let kind = u16::from_le_bytes(frame[6..8].try_into().unwrap()); + let request_id = u64::from_le_bytes(frame[8..16].try_into().unwrap()); + let body_bytes = u64::from_le_bytes(frame[16..24].try_into().unwrap()); + if !matches!(version, VERSION_EXTERNAL | VERSION_SCHEDULED_EXTERNAL) || kind != REQUEST_KIND { + bail!("Rust daemon requires TileMaxSim external protocol v2 or v3"); + } + if usize::try_from(body_bytes).ok() != Some(frame.len() - HEADER_BYTES) { + bail!("request body length mismatch"); + } + let mut reader = Reader::new(&frame[HEADER_BYTES..]); + let dimension = reader.u32()?; + let query_rows = reader.u32()?; + let candidate_count = reader.u32()?; + let dtype = reader.u8()?; + let scoring = reader.u8()?; + let reserved = reader.u16()?; + let contract_bytes = reader.u32()? as usize; + if scoring != 1 || reserved != 0 { + bail!("unsupported scoring function or reserved bits"); + } + if candidate_count > 65_536 { + bail!("too many candidates"); + } + let (priority, timeout_ms, tenant_bytes) = if version == VERSION_SCHEDULED_EXTERNAL { + (reader.i32()?, reader.u32()?, reader.u32()? as usize) + } else { + (0, 0, 0) + }; + if !(-100..=100).contains(&priority) { + bail!("scheduler priority must be between -100 and 100"); + } + if version == VERSION_SCHEDULED_EXTERNAL && !(1..=600_000).contains(&timeout_ms) { + bail!("scheduler timeout must be between 1 and 600000 milliseconds"); + } + let contract = reader.text(contract_bytes, 512, "model contract")?; + let tenant = if version == VERSION_SCHEDULED_EXTERNAL { + reader.text(tenant_bytes, 256, "scheduler tenant")? + } else { + "__default__".to_owned() + }; + let query_bytes = tensor_bytes(query_rows, dimension, dtype)?; + let query = reader.take(query_bytes)?.to_vec(); + validate_finite(&query, dtype)?; + let mut total_tokens = query_rows as usize; + let mut total_bytes = query_bytes; + let mut candidate_ids = HashSet::new(); + let mut candidates = Vec::with_capacity(candidate_count as usize); + for _ in 0..candidate_count { + let candidate_id = reader.u32()?; + let rows = reader.u32()?; + let reference_bytes = reader.u32()? as usize; + let checksum_bytes = reader.u32()? as usize; + if !candidate_ids.insert(candidate_id) { + bail!("duplicate candidate ID"); + } + let tensor_ref = reader.text(reference_bytes, 4096, "tensor reference")?; + let checksum = reader.text(checksum_bytes, 512, "tensor checksum")?; + let digest = tensor_ref + .strip_prefix("sha256://") + .ok_or_else(|| anyhow!("unsupported tensor reference"))?; + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + || checksum != format!("sha256:{digest}") + { + bail!("invalid content-addressed tensor descriptor"); + } + let bytes = tensor_bytes(rows, dimension, dtype)?; + total_tokens = total_tokens + .checked_add(rows as usize) + .ok_or_else(|| anyhow!("token overflow"))?; + total_bytes = total_bytes + .checked_add(bytes) + .ok_or_else(|| anyhow!("byte overflow"))?; + if total_tokens > 1_000_000 || total_bytes > 1024 * 1024 * 1024 { + bail!("request exceeds tensor limits"); + } + candidates.push(Descriptor { + candidate_id, + contract: contract.clone(), + digest: digest.to_owned(), + rows, + dimension, + dtype, + }); + } + reader.finish()?; + Ok(Request { + protocol_version: version, + request_id, + tenant, + priority, + timeout_ms, + query_rows, + dimension, + dtype, + query, + candidates, + }) +} + +pub fn success(version: u16, request_id: u64, results: &[(u32, f32)]) -> Vec { + let body_bytes = 8 + results.len() * 8; + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&version.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&0_u32.to_le_bytes()); + frame.extend_from_slice(&(results.len() as u32).to_le_bytes()); + for (candidate_id, score) in results { + frame.extend_from_slice(&candidate_id.to_le_bytes()); + frame.extend_from_slice(&score.to_le_bytes()); + } + frame +} + +pub fn failure(version: u16, request_id: u64, status: u32, message: &str) -> Vec { + let message = message.as_bytes(); + let message = &message[..message.len().min(64 * 1024)]; + let body_bytes = 8 + message.len(); + let mut frame = Vec::with_capacity(HEADER_BYTES + body_bytes); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&version.to_le_bytes()); + frame.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + frame.extend_from_slice(&request_id.to_le_bytes()); + frame.extend_from_slice(&(body_bytes as u64).to_le_bytes()); + frame.extend_from_slice(&status.max(1).to_le_bytes()); + frame.extend_from_slice(&(message.len() as u32).to_le_bytes()); + frame.extend_from_slice(message); + frame +} + +#[cfg(test)] +mod tests { + use super::*; + + fn scheduled_frame(priority: i32, timeout_ms: u32, tenant: &str) -> Vec { + let contract = "model@1"; + let digest = "a".repeat(64); + let reference = format!("sha256://{digest}"); + let checksum = format!("sha256:{digest}"); + let mut body = Vec::new(); + body.extend_from_slice(&2_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.push(2); + body.push(1); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&(contract.len() as u32).to_le_bytes()); + body.extend_from_slice(&priority.to_le_bytes()); + body.extend_from_slice(&timeout_ms.to_le_bytes()); + body.extend_from_slice(&(tenant.len() as u32).to_le_bytes()); + body.extend_from_slice(contract.as_bytes()); + body.extend_from_slice(tenant.as_bytes()); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&0_u16.to_le_bytes()); + body.extend_from_slice(&7_u32.to_le_bytes()); + body.extend_from_slice(&1_u32.to_le_bytes()); + body.extend_from_slice(&(reference.len() as u32).to_le_bytes()); + body.extend_from_slice(&(checksum.len() as u32).to_le_bytes()); + body.extend_from_slice(reference.as_bytes()); + body.extend_from_slice(checksum.as_bytes()); + let mut frame = Vec::new(); + frame.extend_from_slice(MAGIC); + frame.extend_from_slice(&VERSION_SCHEDULED_EXTERNAL.to_le_bytes()); + frame.extend_from_slice(&REQUEST_KIND.to_le_bytes()); + frame.extend_from_slice(&42_u64.to_le_bytes()); + frame.extend_from_slice(&(body.len() as u64).to_le_bytes()); + frame.extend_from_slice(&body); + frame + } + + #[test] + fn scheduled_protocol_carries_tenant_priority_and_deadline() { + let request = parse(&scheduled_frame(17, 4_000, "tenant-a")).unwrap(); + assert_eq!(request.protocol_version, VERSION_SCHEDULED_EXTERNAL); + assert_eq!(request.request_id, 42); + assert_eq!(request.tenant, "tenant-a"); + assert_eq!(request.priority, 17); + assert_eq!(request.timeout_ms, 4_000); + assert_eq!(request.candidates[0].candidate_id, 7); + } + + #[test] + fn scheduled_protocol_rejects_priority_outside_the_public_contract() { + assert!(parse(&scheduled_frame(101, 4_000, "tenant-a")).is_err()); + } + + #[test] + fn response_uses_the_request_protocol_version() { + let response = success(VERSION_SCHEDULED_EXTERNAL, 9, &[(1, 0.5)]); + assert_eq!( + u16::from_le_bytes(response[4..6].try_into().unwrap()), + VERSION_SCHEDULED_EXTERNAL + ); + } +} diff --git a/services/tilemaxsimd/src/scheduler.rs b/services/tilemaxsimd/src/scheduler.rs new file mode 100644 index 00000000..81e438a2 --- /dev/null +++ b/services/tilemaxsimd/src/scheduler.rs @@ -0,0 +1,338 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use std::collections::HashMap; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SchedulerPolicy { + Fair, + Priority, + FairPriority, +} + +impl SchedulerPolicy { + pub fn parse(value: &str) -> Result { + match value { + "fair" => Ok(Self::Fair), + "priority" => Ok(Self::Priority), + "fair-priority" => Ok(Self::FairPriority), + _ => Err("scheduler policy must be fair, priority, or fair-priority".to_owned()), + } + } +} + +pub struct Scheduled { + pub tenant: String, + pub priority: i32, + pub cost: u64, + pub enqueued_at: Instant, + pub deadline: Instant, + pub payload: T, + sequence: u64, +} + +impl Scheduled { + pub fn new( + tenant: String, + priority: i32, + cost: u64, + enqueued_at: Instant, + deadline: Instant, + payload: T, + ) -> Self { + Self { + tenant, + priority, + cost: cost.max(1), + enqueued_at, + deadline, + payload, + sequence: 0, + } + } +} + +/// Hierarchical request scheduler inspired by vLLM's priority queue, with a +/// tenant-fair first level that vLLM does not need in its usual single-owner +/// deployment. `fair-priority` admits the highest aged priority band and then +/// selects the tenant with the lowest normalized service time. This preserves +/// explicit urgency while preventing equal-priority noisy neighbours from +/// monopolizing the GPU. +pub struct RequestQueue { + policy: SchedulerPolicy, + aging: Duration, + priority_band: i32, + pending: Vec>, + virtual_runtime: HashMap, + weights: HashMap, + next_sequence: u64, +} + +impl RequestQueue { + pub fn new( + policy: SchedulerPolicy, + aging: Duration, + priority_band: i32, + weights: HashMap, + ) -> Self { + Self { + policy, + aging, + priority_band: priority_band.max(0), + pending: Vec::new(), + virtual_runtime: HashMap::new(), + weights, + next_sequence: 0, + } + } + + pub fn len(&self) -> usize { + self.pending.len() + } + + pub fn push(&mut self, mut item: Scheduled) { + item.sequence = self.next_sequence; + self.next_sequence = self.next_sequence.wrapping_add(1); + if self.pending.is_empty() { + // A new busy period should not inherit unbounded floating-point + // service history from a previous burst. + self.virtual_runtime.clear(); + } + let tenant_is_active = self + .pending + .iter() + .any(|pending| pending.tenant == item.tenant); + let baseline = self + .pending + .iter() + .filter_map(|pending| self.virtual_runtime.get(&pending.tenant).copied()) + .reduce(f64::min) + .unwrap_or(0.0); + let runtime = self + .virtual_runtime + .entry(item.tenant.clone()) + .or_insert(baseline); + if !tenant_is_active { + *runtime = runtime.max(baseline); + } + self.pending.push(item); + } + + pub fn drain_expired(&mut self, now: Instant) -> Vec> { + let mut expired = Vec::new(); + let mut index = 0; + while index < self.pending.len() { + if self.pending[index].deadline <= now { + expired.push(self.pending.swap_remove(index)); + } else { + index += 1; + } + } + expired.sort_by_key(|item| item.sequence); + expired + } + + pub fn pop(&mut self, now: Instant) -> Option> { + if self.pending.is_empty() { + return None; + } + let effective = self + .pending + .iter() + .map(|item| self.effective_priority(item, now)) + .collect::>(); + let highest = effective.iter().copied().max().unwrap_or(0); + let mut best = 0; + for candidate in 1..self.pending.len() { + if self.better(candidate, best, &effective, highest) { + best = candidate; + } + } + let item = self.pending.swap_remove(best); + let weight = self.weights.get(&item.tenant).copied().unwrap_or(1.0); + let runtime = self.virtual_runtime.entry(item.tenant.clone()).or_default(); + *runtime += item.cost as f64 / weight; + Some(item) + } + + fn effective_priority(&self, item: &Scheduled, now: Instant) -> i32 { + if self.aging.is_zero() { + return item.priority; + } + let waited = now.saturating_duration_since(item.enqueued_at).as_millis(); + let steps = waited / self.aging.as_millis().max(1); + item.priority + .saturating_add(i32::try_from(steps.min(200)).unwrap_or(200)) + .min(100) + } + + fn better(&self, candidate: usize, current: usize, effective: &[i32], highest: i32) -> bool { + let left = &self.pending[candidate]; + let right = &self.pending[current]; + match self.policy { + SchedulerPolicy::Fair => self.fair_cmp(left, right), + SchedulerPolicy::Priority => { + effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence) + } + SchedulerPolicy::FairPriority => { + let floor = highest.saturating_sub(self.priority_band); + let left_eligible = effective[candidate] >= floor; + let right_eligible = effective[current] >= floor; + if left_eligible != right_eligible { + return left_eligible; + } + if !left_eligible { + return effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence); + } + if left.tenant == right.tenant { + return effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence); + } + let left_runtime = self + .virtual_runtime + .get(&left.tenant) + .copied() + .unwrap_or(0.0); + let right_runtime = self + .virtual_runtime + .get(&right.tenant) + .copied() + .unwrap_or(0.0); + left_runtime < right_runtime + || (left_runtime == right_runtime + && (effective[candidate] > effective[current] + || (effective[candidate] == effective[current] + && left.sequence < right.sequence))) + } + } + } + + fn fair_cmp(&self, left: &Scheduled, right: &Scheduled) -> bool { + let left_runtime = self + .virtual_runtime + .get(&left.tenant) + .copied() + .unwrap_or(0.0); + let right_runtime = self + .virtual_runtime + .get(&right.tenant) + .copied() + .unwrap_or(0.0); + left_runtime < right_runtime + || (left_runtime == right_runtime && left.sequence < right.sequence) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn item(tenant: &str, priority: i32, cost: u64, now: Instant, id: u32) -> Scheduled { + Scheduled::new( + tenant.to_owned(), + priority, + cost, + now, + now + Duration::from_secs(60), + id, + ) + } + + #[test] + fn strict_priority_uses_higher_priority_then_fcfs() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::Priority, + Duration::from_secs(60), + 0, + HashMap::new(), + ); + queue.push(item("a", 0, 1, now, 1)); + queue.push(item("b", 10, 1, now, 2)); + queue.push(item("c", 10, 1, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 2); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 1); + } + + #[test] + fn fair_priority_alternates_equal_priority_tenants() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(60), + 0, + HashMap::new(), + ); + queue.push(item("noisy", 0, 10, now, 1)); + queue.push(item("noisy", 0, 10, now, 2)); + queue.push(item("quiet", 0, 10, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 1); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } + + #[test] + fn full_priority_band_keeps_priority_without_sacrificing_tenant_fairness() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(60), + 200, + HashMap::new(), + ); + queue.push(item("noisy", 100, 10, now, 1)); + queue.push(item("noisy", 100, 10, now, 2)); + queue.push(item("quiet", -100, 10, now, 3)); + assert_eq!(queue.pop(now).unwrap().payload, 1); + assert_eq!(queue.pop(now).unwrap().payload, 3); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } + + #[test] + fn aging_eventually_promotes_waiting_work() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::Priority, + Duration::from_millis(10), + 0, + HashMap::new(), + ); + queue.push(item("old", 0, 1, now, 1)); + queue.push(item("new", 5, 1, now + Duration::from_millis(60), 2)); + assert_eq!( + queue.pop(now + Duration::from_millis(60)).unwrap().payload, + 1 + ); + } + + #[test] + fn expired_requests_leave_without_consuming_service_credit() { + let now = Instant::now(); + let mut queue = RequestQueue::new( + SchedulerPolicy::FairPriority, + Duration::from_secs(1), + 0, + HashMap::new(), + ); + let mut expired = item("a", 0, 1, now, 1); + expired.deadline = now; + queue.push(expired); + queue.push(item("b", 0, 1, now, 2)); + assert_eq!(queue.drain_expired(now).len(), 1); + assert_eq!(queue.pop(now).unwrap().payload, 2); + } +} diff --git a/services/tilemaxsimd/src/shard.rs b/services/tilemaxsimd/src/shard.rs new file mode 100644 index 00000000..3d15b51f --- /dev/null +++ b/services/tilemaxsimd/src/shard.rs @@ -0,0 +1,617 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use crate::cache::TinyLfu; +use crate::protocol::Descriptor; +use anyhow::{Context, Result, anyhow, bail}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use std::collections::HashMap; +use std::fs::{File, OpenOptions}; +use std::io::Read; +use std::os::unix::fs::{FileExt, OpenOptionsExt}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +const INDEX_NAME: &str = "tilemaxsim-shards-v1.json"; + +#[derive(Deserialize)] +struct RawIndex { + format: String, + version: u32, + alignment: usize, + shards: Vec, + entries: Vec, +} + +#[derive(Deserialize)] +struct RawShard { + name: String, + #[serde(rename = "bytes")] + size: usize, + checksum: String, +} + +#[derive(Deserialize)] +struct RawEntry { + digest: String, + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: String, +} + +struct ShardFile { + file: File, + size: usize, + checksum: String, + verified: bool, +} + +#[derive(Clone)] +struct Entry { + shard: String, + offset: u64, + length: usize, + rows: u32, + dimension: u32, + dtype: u8, +} + +struct ContractStore { + root: PathBuf, + shards: HashMap, + entries: HashMap, +} + +struct HostEntry { + payload: Arc<[u8]>, + priority: f64, + owner_tenant: String, +} + +struct HostCache { + maximum_bytes: usize, + current_bytes: usize, + entries: HashMap, + sketch: TinyLfu, + inflation: f64, + tenant_bytes: HashMap, + default_tenant_max_bytes: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl HostCache { + fn new(maximum_bytes: usize, tenant_max_percent: u8) -> Self { + Self { + maximum_bytes, + current_bytes: 0, + entries: HashMap::new(), + sketch: TinyLfu::new(4096), + inflation: 0.0, + tenant_bytes: HashMap::new(), + default_tenant_max_bytes: maximum_bytes * tenant_max_percent as usize / 100, + hits: 0, + misses: 0, + evictions: 0, + admission_rejections: 0, + } + } + + fn get(&mut self, key: &str) -> Option> { + let frequency = self.sketch.increment(&key); + let entry = self.entries.get_mut(key); + if let Some(entry) = entry { + entry.priority = self.inflation + f64::from(frequency) / entry.payload.len() as f64; + self.hits += 1; + Some(entry.payload.clone()) + } else { + self.misses += 1; + None + } + } + + fn remove(&mut self, key: &str) -> Option { + let entry = self.entries.remove(key)?; + self.current_bytes = self.current_bytes.saturating_sub(entry.payload.len()); + let remove_tenant = if let Some(bytes) = self.tenant_bytes.get_mut(&entry.owner_tenant) { + *bytes = bytes.saturating_sub(entry.payload.len()); + *bytes == 0 + } else { + false + }; + if remove_tenant { + self.tenant_bytes.remove(&entry.owner_tenant); + } + Some(entry) + } + + fn put(&mut self, tenant: &str, key: String, payload: Arc<[u8]>) { + if self.maximum_bytes == 0 || payload.len() > self.maximum_bytes { + return; + } + self.remove(&key); + let frequency = self.sketch.estimate(&key).max(1); + let mut priority = self.inflation + f64::from(frequency) / payload.len() as f64; + while self + .tenant_bytes + .get(tenant) + .copied() + .unwrap_or(0) + .saturating_add(payload.len()) + > self.default_tenant_max_bytes + { + let Some((victim_key, victim_priority)) = self + .entries + .iter() + .filter(|(_, entry)| entry.owner_tenant == tenant) + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + .map(|(key, entry)| (key.clone(), entry.priority)) + else { + return; + }; + if priority <= victim_priority { + self.admission_rejections += 1; + return; + } + let victim = self + .remove(&victim_key) + .expect("host tenant victim disappeared"); + self.inflation = self.inflation.max(victim.priority); + priority = self.inflation + f64::from(frequency) / payload.len() as f64; + self.evictions += 1; + } + while self.current_bytes + payload.len() > self.maximum_bytes { + let Some((victim_key, victim_priority)) = self + .entries + .iter() + .min_by(|left, right| left.1.priority.total_cmp(&right.1.priority)) + .map(|(key, entry)| (key.clone(), entry.priority)) + else { + return; + }; + if priority < victim_priority { + self.admission_rejections += 1; + return; + } + let victim = self.remove(&victim_key).expect("host victim disappeared"); + self.inflation = self.inflation.max(victim.priority); + priority = self.inflation + f64::from(frequency) / payload.len() as f64; + self.evictions += 1; + } + self.current_bytes += payload.len(); + *self.tenant_bytes.entry(tenant.to_owned()).or_default() += payload.len(); + self.entries.insert( + key, + HostEntry { + payload, + priority, + owner_tenant: tenant.to_owned(), + }, + ); + } +} + +pub struct ShardStore { + contracts: HashMap, + roots: Vec<(String, PathBuf)>, + host_cache: HostCache, + pub batch_read_calls: u64, + pub batch_read_bytes: u64, + verify_full_shards: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct HostCacheStatus { + pub capacity_bytes: usize, + pub used_bytes: usize, + pub entries: usize, + pub tenants: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub admission_rejections: u64, +} + +impl ShardStore { + pub fn open( + roots: &[(String, PathBuf)], + host_cache_bytes: usize, + host_tenant_max_percent: u8, + verify_full_shards: bool, + ) -> Result { + let mut contracts = HashMap::new(); + for (contract, root) in roots { + if contract.is_empty() || contracts.contains_key(contract) { + bail!("invalid or duplicate model contract {contract:?}"); + } + contracts.insert(contract.clone(), open_contract(root)?); + } + if contracts.is_empty() { + bail!("at least one immutable shard contract is required"); + } + Ok(Self { + contracts, + roots: roots.to_vec(), + host_cache: HostCache::new(host_cache_bytes, host_tenant_max_percent), + batch_read_calls: 0, + batch_read_bytes: 0, + verify_full_shards, + }) + } + + /// Re-open every immutable contract index and publish the new generation + /// only after all roots validate. Content-addressed host-cache entries stay + /// valid across the atomic metadata swap. + pub fn reload(&mut self) -> Result<()> { + let mut contracts = HashMap::new(); + for (contract, root) in &self.roots { + contracts.insert(contract.clone(), open_contract(root)?); + } + self.contracts = contracts; + Ok(()) + } + + pub fn resolve_many( + &mut self, + descriptors: &[Descriptor], + tenant: &str, + ) -> Result>> { + let mut output = vec![None; descriptors.len()]; + let mut groups = HashMap::<(String, String), Vec<(usize, Entry)>>::new(); + for (index, descriptor) in descriptors.iter().enumerate() { + let key = cache_key(descriptor); + if let Some(payload) = self.host_cache.get(&key) { + output[index] = Some(payload); + continue; + } + let contract = self + .contracts + .get(&descriptor.contract) + .ok_or_else(|| anyhow!("model contract has no immutable shard root"))?; + let Some(entry) = contract.entries.get(&descriptor.digest).cloned() else { + let expected = + tensor_bytes(descriptor.rows, descriptor.dimension, descriptor.dtype)?; + let mut file = open_object(&contract.root, &descriptor.digest) + .with_context(|| "tensor is missing from immutable shards and objects")?; + if file.metadata()?.len() != expected as u64 { + bail!("immutable tensor object length disagrees with its descriptor"); + } + let mut payload = Vec::with_capacity(expected); + file.read_to_end(&mut payload)?; + if hex::encode(Sha256::digest(&payload)) != descriptor.digest { + bail!("immutable tensor object checksum mismatch"); + } + self.batch_read_calls += 1; + self.batch_read_bytes += payload.len() as u64; + let payload: Arc<[u8]> = payload.into(); + output[index] = Some(Arc::clone(&payload)); + self.host_cache.put(tenant, key, payload); + continue; + }; + if entry.rows != descriptor.rows + || entry.dimension != descriptor.dimension + || entry.dtype != descriptor.dtype + || entry.length + != tensor_bytes(descriptor.rows, descriptor.dimension, descriptor.dtype)? + { + bail!("tensor descriptor disagrees with the immutable shard index"); + } + groups + .entry((descriptor.contract.clone(), entry.shard.clone())) + .or_default() + .push((index, entry)); + } + + for ((contract_name, shard_name), mut entries) in groups { + let contract = self.contracts.get_mut(&contract_name).unwrap(); + let shard = contract.shards.get_mut(&shard_name).unwrap(); + if self.verify_full_shards { + verify_shard(shard)?; + } + entries.sort_by_key(|(_, entry)| entry.offset); + let mut ranges = Vec::<(u64, u64, Vec<(usize, Entry)>)>::new(); + let mut cursor = 0; + while cursor < entries.len() { + let start = entries[cursor].1.offset; + let mut end = start + entries[cursor].1.length as u64; + let mut limit = cursor + 1; + while limit < entries.len() { + let candidate = &entries[limit].1; + let candidate_end = candidate.offset + candidate.length as u64; + if candidate.offset.saturating_sub(end) > 64 * 1024 + || candidate_end.saturating_sub(start) > 8 * 1024 * 1024 + { + break; + } + end = end.max(candidate_end); + limit += 1; + } + ranges.push((start, end, entries[cursor..limit].to_vec())); + cursor = limit; + } + let worker_count = ranges.len().min(8); + let ranges_per_worker = ranges.len().div_ceil(worker_count); + let resolved = std::thread::scope(|scope| -> Result> { + let mut workers = Vec::new(); + for ranges in ranges.chunks(ranges_per_worker) { + let file = shard.file.try_clone()?; + workers.push(scope.spawn(move || -> Result)>> { + let mut resolved = Vec::new(); + for (start, end, entries) in ranges { + let mut range = vec![0_u8; (end - start) as usize]; + file.read_exact_at(&mut range, *start)?; + for (output_index, entry) in entries { + let local = (entry.offset - start) as usize; + let payload = range[local..local + entry.length].to_vec(); + let actual = hex::encode(Sha256::digest(&payload)); + if actual != descriptors[*output_index].digest { + bail!("tensor checksum mismatch inside immutable shard"); + } + resolved.push((*output_index, payload)); + } + } + Ok(resolved) + })); + } + let mut resolved = Vec::new(); + for worker in workers { + resolved.extend( + worker + .join() + .map_err(|_| anyhow!("shard reader worker panicked"))??, + ); + } + Ok(resolved) + })?; + self.batch_read_calls += ranges.len() as u64; + self.batch_read_bytes += ranges + .iter() + .map(|(start, end, _)| end - start) + .sum::(); + for (output_index, payload) in resolved { + let payload: Arc<[u8]> = payload.into(); + output[output_index] = Some(Arc::clone(&payload)); + self.host_cache + .put(tenant, cache_key(&descriptors[output_index]), payload); + } + } + output + .into_iter() + .map(|payload| payload.ok_or_else(|| anyhow!("tensor resolution produced no payload"))) + .collect() + } + + pub fn host_status(&self) -> HostCacheStatus { + HostCacheStatus { + capacity_bytes: self.host_cache.maximum_bytes, + used_bytes: self.host_cache.current_bytes, + entries: self.host_cache.entries.len(), + tenants: self.host_cache.tenant_bytes.len(), + hits: self.host_cache.hits, + misses: self.host_cache.misses, + evictions: self.host_cache.evictions, + admission_rejections: self.host_cache.admission_rejections, + } + } +} + +fn open_contract(root: &Path) -> Result { + let index_path = root.join(INDEX_NAME); + let root_metadata = std::fs::symlink_metadata(root) + .with_context(|| format!("cannot inspect immutable tensor root {}", root.display()))?; + if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { + bail!("immutable tensor root must be a real directory"); + } + if !index_path.try_exists()? { + return Ok(ContractStore { + root: root.to_path_buf(), + shards: HashMap::new(), + entries: HashMap::new(), + }); + } + let mut index_file = nofollow(&index_path)?; + let mut document = String::new(); + index_file.read_to_string(&mut document)?; + let raw: RawIndex = serde_json::from_str(&document).context("invalid immutable shard index")?; + if raw.format != "vectorchord.tilemaxsim.shards" || raw.version != 1 { + bail!("unsupported immutable shard index"); + } + if raw.alignment == 0 || !raw.alignment.is_power_of_two() { + bail!("invalid immutable shard alignment"); + } + let mut shards = HashMap::new(); + for shard in raw.shards { + let digest = shard + .checksum + .strip_prefix("sha256:") + .ok_or_else(|| anyhow!("invalid immutable shard checksum"))?; + if shard.name != format!("shards/sha256-{digest}.vts") + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + bail!("invalid immutable shard name"); + } + let path = root.join(&shard.name); + let file = nofollow(&path)?; + if file.metadata()?.len() != shard.size as u64 { + bail!("immutable shard size mismatch"); + } + if shards + .insert( + shard.name, + ShardFile { + file, + size: shard.size, + checksum: digest.to_owned(), + verified: false, + }, + ) + .is_some() + { + bail!("duplicate immutable shard"); + } + } + let mut entries = HashMap::new(); + for entry in raw.entries { + let dtype = match entry.dtype.as_str() { + "float32" => 1, + "float16" => 2, + _ => bail!("invalid shard tensor dtype"), + }; + let shard = shards + .get(&entry.shard) + .ok_or_else(|| anyhow!("shard tensor references an unknown file"))?; + if !(entry.offset as usize).is_multiple_of(raw.alignment) + || entry.length != tensor_bytes(entry.rows, entry.dimension, dtype)? + || entry.offset + entry.length as u64 > shard.size as u64 + { + bail!("invalid shard tensor range"); + } + if entries + .insert( + entry.digest, + Entry { + shard: entry.shard, + offset: entry.offset, + length: entry.length, + rows: entry.rows, + dimension: entry.dimension, + dtype, + }, + ) + .is_some() + { + bail!("duplicate tensor digest in shard index"); + } + } + Ok(ContractStore { + root: root.to_path_buf(), + shards, + entries, + }) +} + +fn object_path(root: &Path, digest: &str) -> Result { + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + { + bail!("invalid tensor object digest"); + } + Ok(root + .join("objects") + .join(&digest[..2]) + .join(format!("{digest}.tensor"))) +} + +fn open_object(root: &Path, digest: &str) -> Result { + let path = object_path(root, digest)?; + let objects = root.join("objects"); + let prefix = objects.join(&digest[..2]); + for directory in [&objects, &prefix] { + let metadata = std::fs::symlink_metadata(directory)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + bail!("immutable tensor object path contains a non-directory component"); + } + } + nofollow(&path) +} + +fn nofollow(path: &Path) -> Result { + OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path) + .with_context(|| format!("cannot open immutable shard path {}", path.display())) +} + +fn verify_shard(shard: &mut ShardFile) -> Result<()> { + if shard.verified { + return Ok(()); + } + let mut digest = Sha256::new(); + let mut offset = 0_u64; + let mut buffer = vec![0_u8; 8 * 1024 * 1024]; + while offset < shard.size as u64 { + let count = buffer.len().min(shard.size - offset as usize); + shard.file.read_exact_at(&mut buffer[..count], offset)?; + digest.update(&buffer[..count]); + offset += count as u64; + } + if hex::encode(digest.finalize()) != shard.checksum { + bail!("immutable shard checksum mismatch"); + } + shard.verified = true; + Ok(()) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: u8) -> Result { + let scalar = match dtype { + 1 => 4, + 2 => 2, + _ => bail!("unsupported tensor dtype"), + }; + (rows as usize) + .checked_mul(dimension as usize) + .and_then(|value| value.checked_mul(scalar)) + .ok_or_else(|| anyhow!("tensor shape overflow")) +} + +pub fn cache_key(descriptor: &Descriptor) -> String { + format!( + "{}:{}:{}:{}:{}", + descriptor.contract, + descriptor.digest, + descriptor.rows, + descriptor.dimension, + descriptor.dtype + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[test] + fn content_addressed_objects_are_visible_without_index_reload() { + let root = + std::env::temp_dir().join(format!("tilemaxsim-object-store-{}", std::process::id())); + let _ = fs::remove_dir_all(&root); + fs::create_dir_all(&root).unwrap(); + let mut store = + ShardStore::open(&[("model@1".to_owned(), root.clone())], 1024, 100, false).unwrap(); + let payload = [0_u8, 60, 0, 0]; + let digest = hex::encode(Sha256::digest(payload)); + let directory = root.join("objects").join(&digest[..2]); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(format!("{digest}.tensor")), payload).unwrap(); + let descriptor = Descriptor { + candidate_id: 1, + contract: "model@1".to_owned(), + digest, + rows: 1, + dimension: 2, + dtype: 2, + }; + + let resolved = store.resolve_many(&[descriptor], "tenant-a").unwrap(); + assert_eq!(&*resolved[0], &payload); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/sql/install/vchord--0.1.0.sql b/sql/install/vchord--0.1.0.sql new file mode 100644 index 00000000..8a8fd495 --- /dev/null +++ b/sql/install/vchord--0.1.0.sql @@ -0,0 +1,164 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:17 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:53 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:29 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:5 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am.rs:29 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:6 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:6 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:1 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:18 +-- finalize +-- List of data types + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.2.0.sql b/sql/install/vchord--0.2.0.sql new file mode 100644 index 00000000..dca70ae1 --- /dev/null +++ b/sql/install/vchord--0.2.0.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:18 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:54 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:30 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:6 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:54 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:30 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:6 +-- vchord::datatype::operators_pgvector_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_pgvector_vector::PgvectorVectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am.rs:33 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:12 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:6 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:1 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:19 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.2.1.sql b/sql/install/vchord--0.2.1.sql new file mode 100644 index 00000000..e5cf547b --- /dev/null +++ b/sql/install/vchord--0.2.1.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:10 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:54 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:30 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:6 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:54 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:30 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:6 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:60 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:9 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:11 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.2.2.sql b/sql/install/vchord--0.2.2.sql new file mode 100644 index 00000000..f408248c --- /dev/null +++ b/sql/install/vchord--0.2.2.sql @@ -0,0 +1,500 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:11 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:54 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:30 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:6 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:54 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:30 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:6 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:57 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:5 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:12 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.3.0.sql b/sql/install/vchord--0.3.0.sql new file mode 100644 index 00000000..7d1088fd --- /dev/null +++ b/sql/install/vchord--0.3.0.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:9 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:79 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:55 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:31 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:7 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:79 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:55 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:31 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:7 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:60 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:5 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:36 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:31 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:46 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:11 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:41 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:10 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/install/vchord--0.4.0.sql b/sql/install/vchord--0.4.0.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.0.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/install/vchord--0.4.1.sql b/sql/install/vchord--0.4.1.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.1.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/install/vchord--0.4.2.sql b/sql/install/vchord--0.4.2.sql new file mode 100644 index 00000000..522e9c8a --- /dev/null +++ b/sql/install/vchord--0.4.2.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:74 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/install/vchord--0.4.3.sql b/sql/install/vchord--0.4.3.sql new file mode 100644 index 00000000..583a2f61 --- /dev/null +++ b/sql/install/vchord--0.4.3.sql @@ -0,0 +1,566 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:22 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:21 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:133 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/am/mod.rs:75 +-- vchord::index::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:23 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); +/* */ + diff --git a/sql/install/vchord--0.5.0.sql b/sql/install/vchord--0.5.0.sql new file mode 100644 index 00000000..173d6807 --- /dev/null +++ b/sql/install/vchord--0.5.0.sql @@ -0,0 +1,752 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:32 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:22 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:77 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:77 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:41 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.5.1.sql b/sql/install/vchord--0.5.1.sql new file mode 100644 index 00000000..dcaf9b98 --- /dev/null +++ b/sql/install/vchord--0.5.1.sql @@ -0,0 +1,752 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:77 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:19 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:77 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:41 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); +/* */ + diff --git a/sql/install/vchord--0.5.2.sql b/sql/install/vchord--0.5.2.sql new file mode 100644 index 00000000..0e7e473e --- /dev/null +++ b/sql/install/vchord--0.5.2.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:46 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:59 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:77 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:78 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:47 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/install/vchord--0.5.3.sql b/sql/install/vchord--0.5.3.sql new file mode 100644 index 00000000..9c428f43 --- /dev/null +++ b/sql/install/vchord--0.5.3.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:46 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array */ + "rhs" halfvec[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:58 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:76 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array */ + "rhs" vector[] /* pgrx::datum::array::Array */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:80 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:47 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/install/vchord--1.0.0.sql b/sql/install/vchord--1.0.0.sql new file mode 100644 index 00000000..e50ca725 --- /dev/null +++ b/sql/install/vchord--1.0.0.sql @@ -0,0 +1,857 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:45 +-- bootstrap +-- List of shell types + +CREATE TYPE scalar8; +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:31 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:20 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:40 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:20 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:30 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:132 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:36 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:21 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:98 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:50 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:74 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input<'_> */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:58 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:76 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:21 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:78 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:81 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:46 +-- finalize +-- List of data types + +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/install/vchord--1.1.0.sql b/sql/install/vchord--1.1.0.sql new file mode 100644 index 00000000..89241be2 --- /dev/null +++ b/sql/install/vchord--1.1.0.sql @@ -0,0 +1,1291 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:40 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:40 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:22 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:22 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/install/vchord--1.1.1.sql b/sql/install/vchord--1.1.1.sql new file mode 100644 index 00000000..c5ecd748 --- /dev/null +++ b/sql/install/vchord--1.1.1.sql @@ -0,0 +1,1323 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:41 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:41 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:72 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:59 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:72 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:59 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:23 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:23 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; +/* */ + diff --git a/sql/install/vchord--1.2.0.sql b/sql/install/vchord--1.2.0.sql new file mode 100644 index 00000000..e6ffb9cc --- /dev/null +++ b/sql/install/vchord--1.2.0.sql @@ -0,0 +1,2912 @@ +/* */ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ +/* */ + +/* */ +-- src/lib.rs:47 +-- bootstrap +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:93 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_operator_maxsim +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ + "rhs" halfvec[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_halfvec::HalfvecInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:41 +-- vchord::datatype::functions_rabitq4::_vchord_halfvec_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:41 +-- vchord::datatype::functions_rabitq8::_vchord_halfvec_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:69 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:45 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_halfvec.rs:21 +-- vchord::datatype::operators_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_halfvec::HalfvecInput<'_> */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:72 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:59 +-- vchord::datatype::functions_rabitq4::_vchord_rabitq4_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:20 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_in +CREATE FUNCTION "_vchord_rabitq4_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:41 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_cosine +CREATE FUNCTION "_vchord_rabitq4_operator_cosine"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:31 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_ip +CREATE FUNCTION "_vchord_rabitq4_operator_ip"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:21 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_l2 +CREATE FUNCTION "_vchord_rabitq4_operator_l2"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:123 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq4.rs:140 +-- vchord::datatype::text_rabitq4::_vchord_rabitq4_out +CREATE FUNCTION "_vchord_rabitq4_out"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:36 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_recv +CREATE FUNCTION "_vchord_rabitq4_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq4.rs:21 +-- vchord::datatype::binary_rabitq4::_vchord_rabitq4_send +CREATE FUNCTION "_vchord_rabitq4_send"( + "vector" rabitq4 /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:99 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq4_sphere_cosine_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:75 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq4_sphere_ip_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq4.rs:51 +-- vchord::datatype::operators_rabitq4::_vchord_rabitq4_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq4_sphere_l2_in"( + "lhs" rabitq4, /* vchord::datatype::memory_rabitq4::Rabitq4Input<'_> */ + "rhs" sphere_rabitq4 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq4.rs:17 +-- vchord::datatype::typmod_rabitq4::_vchord_rabitq4_typmod_in +CREATE FUNCTION "_vchord_rabitq4_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:72 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_halfvec +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:59 +-- vchord::datatype::functions_rabitq8::_vchord_rabitq8_dequantize_to_vector +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:20 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_in +CREATE FUNCTION "_vchord_rabitq8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:41 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_cosine +CREATE FUNCTION "_vchord_rabitq8_operator_cosine"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:31 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_ip +CREATE FUNCTION "_vchord_rabitq8_operator_ip"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:21 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_l2 +CREATE FUNCTION "_vchord_rabitq8_operator_l2"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:123 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_operator_maxsim +/* */ + +/* */ +-- src/datatype/text_rabitq8.rs:140 +-- vchord::datatype::text_rabitq8::_vchord_rabitq8_out +CREATE FUNCTION "_vchord_rabitq8_out"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:36 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_recv +CREATE FUNCTION "_vchord_rabitq8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_rabitq8.rs:21 +-- vchord::datatype::binary_rabitq8::_vchord_rabitq8_send +CREATE FUNCTION "_vchord_rabitq8_send"( + "vector" rabitq8 /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:99 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_cosine_in +CREATE FUNCTION "_vchord_rabitq8_sphere_cosine_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:75 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_ip_in +CREATE FUNCTION "_vchord_rabitq8_sphere_ip_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_rabitq8.rs:51 +-- vchord::datatype::operators_rabitq8::_vchord_rabitq8_sphere_l2_in +CREATE FUNCTION "_vchord_rabitq8_sphere_l2_in"( + "lhs" rabitq8, /* vchord::datatype::memory_rabitq8::Rabitq8Input<'_> */ + "rhs" sphere_rabitq8 /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod_rabitq8.rs:17 +-- vchord::datatype::typmod_rabitq8::_vchord_rabitq8_typmod_in +CREATE FUNCTION "_vchord_rabitq8_typmod_in"( + "list" cstring[] /* pgrx::datum::array::Array<'_, &core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:93 +-- vchord::datatype::operators_vector::_vchord_vector_operator_maxsim +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ + "rhs" vector[] /* pgrx::datum::array::Array<'_, vchord::datatype::memory_vector::VectorInput<'_>> */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_rabitq4.rs:23 +-- vchord::datatype::functions_rabitq4::_vchord_vector_quantize_to_rabitq4 +/* */ + +/* */ +-- src/datatype/functions_rabitq8.rs:23 +-- vchord::datatype::functions_rabitq8::_vchord_vector_quantize_to_rabitq8 +/* */ + +/* */ +-- src/datatype/operators_vector.rs:69 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_cosine_in +CREATE FUNCTION "_vchord_vector_sphere_cosine_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:45 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_ip_in +CREATE FUNCTION "_vchord_vector_sphere_ip_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_vector.rs:21 +-- vchord::datatype::operators_vector::_vchord_vector_sphere_l2_in +CREATE FUNCTION "_vchord_vector_sphere_l2_in"( + "lhs" vector, /* vchord::datatype::memory_vector::VectorInput<'_> */ + "rhs" sphere_vector /* pgrx::heap_tuple::PgHeapTuple<'_, pgrx::pgbox::AllocatedByRust> */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_vector_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/index/vchordg/am/mod.rs:110 +-- vchord::index::vchordg::am::_vchordg_amhandler +/* */ + +/* */ +-- src/index/functions.rs:21 +-- vchord::index::functions::_vchordg_prewarm +/* */ + +/* */ +-- src/index/opclass.rs:35 +-- vchord::index::opclass::_vchordg_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordg_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:40 +-- vchord::index::opclass::_vchordg_support_halfvec_ip_ops +CREATE FUNCTION "_vchordg_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:30 +-- vchord::index::opclass::_vchordg_support_halfvec_l2_ops +CREATE FUNCTION "_vchordg_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:65 +-- vchord::index::opclass::_vchordg_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:70 +-- vchord::index::opclass::_vchordg_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:60 +-- vchord::index::opclass::_vchordg_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:50 +-- vchord::index::opclass::_vchordg_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordg_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:55 +-- vchord::index::opclass::_vchordg_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordg_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:45 +-- vchord::index::opclass::_vchordg_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordg_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:20 +-- vchord::index::opclass::_vchordg_support_vector_cosine_ops +CREATE FUNCTION "_vchordg_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:25 +-- vchord::index::opclass::_vchordg_support_vector_ip_ops +CREATE FUNCTION "_vchordg_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:15 +-- vchord::index::opclass::_vchordg_support_vector_l2_ops +CREATE FUNCTION "_vchordg_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/am/mod.rs:192 +-- vchord::index::vchordrq::am::_vchordrq_amhandler +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/search.rs:45 +-- vchord::index::vchordrq::scanners::maxsim::search::_vchordrq_maxsim_search_external +/* */ + +/* */ +-- src/index/functions.rs:43 +-- vchord::index::functions::_vchordrq_prewarm +/* */ + +/* */ +-- src/index/functions.rs:90 +-- vchord::index::functions::_vchordrq_sampled_values +/* */ + +/* */ +-- src/index/opclass.rs:100 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:95 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:90 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:140 +-- vchord::index::opclass::_vchordrq_support_halfvec_maxsim_ops +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:130 +-- vchord::index::opclass::_vchordrq_support_rabitq4_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:125 +-- vchord::index::opclass::_vchordrq_support_rabitq4_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:120 +-- vchord::index::opclass::_vchordrq_support_rabitq4_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:150 +-- vchord::index::opclass::_vchordrq_support_rabitq4_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq4_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:115 +-- vchord::index::opclass::_vchordrq_support_rabitq8_cosine_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:110 +-- vchord::index::opclass::_vchordrq_support_rabitq8_ip_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:105 +-- vchord::index::opclass::_vchordrq_support_rabitq8_l2_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:145 +-- vchord::index::opclass::_vchordrq_support_rabitq8_maxsim_ops +CREATE FUNCTION "_vchordrq_support_rabitq8_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:85 +-- vchord::index::opclass::_vchordrq_support_vector_cosine_ops +CREATE FUNCTION "_vchordrq_support_vector_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:80 +-- vchord::index::opclass::_vchordrq_support_vector_ip_ops +CREATE FUNCTION "_vchordrq_support_vector_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:75 +-- vchord::index::opclass::_vchordrq_support_vector_l2_ops +CREATE FUNCTION "_vchordrq_support_vector_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_l2_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:135 +-- vchord::index::opclass::_vchordrq_support_vector_maxsim_ops +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/exact.rs:51 +-- vchord::index::vchordrq::scanners::maxsim::exact::_vchordrq_tilemaxsim_rerank_halfvec +/* */ + +/* */ +-- src/index/vchordrq/scanners/maxsim/exact.rs:34 +-- vchord::index::vchordrq::scanners::maxsim::exact::_vchordrq_tilemaxsim_rerank_vector +/* */ + +/* */ +-- src/lib.rs:48 +-- finalize +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); +/* */ + diff --git a/sql/upgrade/vchord--0.1.0--0.2.0.sql b/sql/upgrade/vchord--0.1.0--0.2.0.sql new file mode 100644 index 00000000..4360f27f --- /dev/null +++ b/sql/upgrade/vchord--0.1.0--0.2.0.sql @@ -0,0 +1,343 @@ +-- src/lib.rs:18 +CREATE TYPE scalar8; +CREATE TYPE sphere_halfvec; +CREATE TYPE sphere_scalar8; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:18 +-- vchord::datatype::functions_scalar8::_vchord_halfvec_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:54 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_cosine_in +CREATE FUNCTION "_vchord_halfvec_sphere_cosine_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:30 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_ip_in +CREATE FUNCTION "_vchord_halfvec_sphere_ip_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_pgvector_halfvec.rs:6 +-- vchord::datatype::operators_pgvector_halfvec::_vchord_halfvec_sphere_l2_in +CREATE FUNCTION "_vchord_halfvec_sphere_l2_in"( + "lhs" halfvec, /* vchord::datatype::memory_pgvector_halfvec::PgvectorHalfvecInput */ + "rhs" sphere_halfvec /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_halfvec_sphere_l2_in_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:7 +-- vchord::datatype::text_scalar8::_vchord_scalar8_in +CREATE FUNCTION "_vchord_scalar8_in"( + "input" cstring, /* &core::ffi::c_str::CStr */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:26 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_cosine +CREATE FUNCTION "_vchord_scalar8_operator_cosine"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_cosine_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:6 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_ip +CREATE FUNCTION "_vchord_scalar8_operator_ip"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_ip_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:16 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_operator_l2 +CREATE FUNCTION "_vchord_scalar8_operator_l2"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS real /* f32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_operator_l2_wrapper'; +/* */ + +/* */ +-- src/datatype/text_scalar8.rs:119 +-- vchord::datatype::text_scalar8::_vchord_scalar8_out +CREATE FUNCTION "_vchord_scalar8_out"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_out_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:22 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_recv +CREATE FUNCTION "_vchord_scalar8_recv"( + "internal" internal, /* pgrx::datum::internal::Internal */ + "oid" oid, /* pgrx_pg_sys::submodules::oids::Oid */ + "typmod" INT /* i32 */ +) RETURNS scalar8 /* vchord::datatype::memory_scalar8::Scalar8Output */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_recv_wrapper'; +/* */ + +/* */ +-- src/datatype/binary_scalar8.rs:7 +-- vchord::datatype::binary_scalar8::_vchord_scalar8_send +CREATE FUNCTION "_vchord_scalar8_send"( + "vector" scalar8 /* vchord::datatype::memory_scalar8::Scalar8Input */ +) RETURNS bytea /* alloc::vec::Vec */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_send_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:84 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_cosine_in +CREATE FUNCTION "_vchord_scalar8_sphere_cosine_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_cosine_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:36 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_ip_in +CREATE FUNCTION "_vchord_scalar8_sphere_ip_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_ip_in_wrapper'; +/* */ + +/* */ +-- src/datatype/operators_scalar8.rs:60 +-- vchord::datatype::operators_scalar8::_vchord_scalar8_sphere_l2_in +CREATE FUNCTION "_vchord_scalar8_sphere_l2_in"( + "lhs" scalar8, /* vchord::datatype::memory_scalar8::Scalar8Input */ + "rhs" sphere_scalar8 /* pgrx::heap_tuple::PgHeapTuple */ +) RETURNS bool /* bool */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_scalar8_sphere_l2_in_wrapper'; +-- src/datatype/typmod.rs:45 +-- vchord::datatype::typmod::_vchord_typmod_in_65535 +CREATE FUNCTION "_vchord_typmod_in_65535"( + "list" cstring[] /* pgrx::datum::array::Array<&core::ffi::c_str::CStr> */ +) RETURNS INT /* i32 */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_in_65535_wrapper'; +/* */ + +/* */ +-- src/datatype/typmod.rs:63 +-- vchord::datatype::typmod::_vchord_typmod_out +CREATE FUNCTION "_vchord_typmod_out"( + "typmod" INT /* i32 */ +) RETURNS cstring /* alloc::ffi::c_str::CString */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchord_typmod_out_wrapper'; +/* */ + +/* */ +-- src/datatype/functions_scalar8.rs:8 +-- vchord::datatype::functions_scalar8::_vchord_vector_quantize_to_scalar8 +/* */ + +/* */ +-- src/datatype/operators_pgvector_vector.rs:54 +-- src/datatype/operators_pgvector_vector.rs:30 +-- src/datatype/operators_pgvector_vector.rs:6 +-- src/index/am.rs:33 +-- src/index/functions.rs:12 +-- src/index/opclass.rs:26 +-- vchord::index::opclass::_vchordrq_support_halfvec_cosine_ops +CREATE FUNCTION "_vchordrq_support_halfvec_cosine_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_cosine_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:21 +-- vchord::index::opclass::_vchordrq_support_halfvec_ip_ops +CREATE FUNCTION "_vchordrq_support_halfvec_ip_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_ip_ops_wrapper'; +/* */ + +/* */ +-- src/index/opclass.rs:16 +-- vchord::index::opclass::_vchordrq_support_halfvec_l2_ops +CREATE FUNCTION "_vchordrq_support_halfvec_l2_ops"() RETURNS TEXT /* alloc::string::String */ +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c /* Rust */ +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_l2_ops_wrapper'; +/* */ + +/* */ +-- src/lib.rs:19 +CREATE TYPE scalar8 ( + INPUT = _vchord_scalar8_in, + OUTPUT = _vchord_scalar8_out, + RECEIVE = _vchord_scalar8_recv, + SEND = _vchord_scalar8_send, + TYPMOD_IN = _vchord_typmod_in_65535, + TYPMOD_OUT = _vchord_typmod_out, + STORAGE = EXTERNAL, + INTERNALLENGTH = VARIABLE, + ALIGNMENT = double +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_scalar8 AS ( + center scalar8, + radius REAL +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_scalar8_operator_l2, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_scalar8_operator_ip, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_scalar8_operator_cosine, + LEFTARG = scalar8, + RIGHTARG = scalar8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_scalar8_sphere_l2_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<->> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_scalar8_sphere_ip_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<#>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec, + COMMUTATOR = <<=>> +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_scalar8_sphere_cosine_in, + LEFTARG = scalar8, + RIGHTARG = sphere_scalar8, + COMMUTATOR = <<=>> +); + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(scalar8, real) RETURNS sphere_scalar8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_scalar8(vector) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_scalar8_wrapper'; + +CREATE FUNCTION quantize_to_scalar8(halfvec) RETURNS scalar8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_scalar8_wrapper'; + +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); diff --git a/sql/upgrade/vchord--0.2.0--0.2.1.sql b/sql/upgrade/vchord--0.2.0--0.2.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.2.0--0.2.1.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.2.1--0.2.2.sql b/sql/upgrade/vchord--0.2.1--0.2.2.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.2.1--0.2.2.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.2.2--0.3.0.sql b/sql/upgrade/vchord--0.2.2--0.3.0.sql new file mode 100644 index 00000000..bbbc98bf --- /dev/null +++ b/sql/upgrade/vchord--0.2.2--0.3.0.sql @@ -0,0 +1,51 @@ +CREATE FUNCTION "_vchord_halfvec_operator_maxsim"( + "lhs" halfvec[], + "rhs" halfvec[] +) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchord_halfvec_operator_maxsim_wrapper'; + +CREATE FUNCTION "_vchord_vector_operator_maxsim"( + "lhs" vector[], + "rhs" vector[] +) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchord_vector_operator_maxsim_wrapper'; + +CREATE FUNCTION "_vchordrq_support_halfvec_maxsim_ops"() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_support_halfvec_maxsim_ops_wrapper'; + +CREATE FUNCTION "_vchordrq_support_vector_maxsim_ops"() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_support_vector_maxsim_ops_wrapper'; + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; + +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); diff --git a/sql/upgrade/vchord--0.3.0--0.4.0.sql b/sql/upgrade/vchord--0.3.0--0.4.0.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.3.0--0.4.0.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.4.0--0.4.1.sql b/sql/upgrade/vchord--0.4.0--0.4.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.0--0.4.1.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.4.1--0.4.2.sql b/sql/upgrade/vchord--0.4.1--0.4.2.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.1--0.4.2.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.4.2--0.4.3.sql b/sql/upgrade/vchord--0.4.2--0.4.3.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.4.2--0.4.3.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.4.3--0.5.0.sql b/sql/upgrade/vchord--0.4.3--0.5.0.sql new file mode 100644 index 00000000..ee5bdbcf --- /dev/null +++ b/sql/upgrade/vchord--0.4.3--0.5.0.sql @@ -0,0 +1,151 @@ +-- internal changes + +CREATE FUNCTION _vchordg_support_halfvec_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_halfvec_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_halfvec_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_halfvec_l2_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_vector_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_vector_l2_ops_wrapper'; + +-- List of functions + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + +-- List of access methods + +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; + +-- List of operator families + +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); diff --git a/sql/upgrade/vchord--0.5.0--0.5.1.sql b/sql/upgrade/vchord--0.5.0--0.5.1.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.0--0.5.1.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.5.1--0.5.2.sql b/sql/upgrade/vchord--0.5.1--0.5.2.sql new file mode 100644 index 00000000..a28eefcd --- /dev/null +++ b/sql/upgrade/vchord--0.5.1--0.5.2.sql @@ -0,0 +1,101 @@ +-- List of functions + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; diff --git a/sql/upgrade/vchord--0.5.2--0.5.3.sql b/sql/upgrade/vchord--0.5.2--0.5.3.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.2--0.5.3.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--0.5.3--1.0.0.sql b/sql/upgrade/vchord--0.5.3--1.0.0.sql new file mode 100644 index 00000000..be321fe8 --- /dev/null +++ b/sql/upgrade/vchord--0.5.3--1.0.0.sql @@ -0,0 +1 @@ +/* nothing to do */ diff --git a/sql/upgrade/vchord--1.0.0--1.1.0.sql b/sql/upgrade/vchord--1.0.0--1.1.0.sql new file mode 100644 index 00000000..555c3e6a --- /dev/null +++ b/sql/upgrade/vchord--1.0.0--1.1.0.sql @@ -0,0 +1,405 @@ +-- List of shell types + +CREATE TYPE rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq8; +CREATE TYPE sphere_rabitq4; + +-- List of internal changes + +DROP OPERATOR <<=>>(sphere_scalar8, scalar8); +DROP OPERATOR <<#>>(sphere_scalar8, scalar8); +DROP OPERATOR <<->>(sphere_scalar8, scalar8); +DROP FUNCTION quantize_to_scalar8(halfvec); +DROP FUNCTION quantize_to_scalar8(vector); +DROP FUNCTION sphere(scalar8, real); +DROP OPERATOR <<=>>(scalar8, sphere_scalar8); +DROP OPERATOR <<#>>(scalar8, sphere_scalar8); +DROP OPERATOR <<->>(scalar8, sphere_scalar8); +DROP OPERATOR <=>(scalar8, scalar8); +DROP OPERATOR <#>(scalar8, scalar8); +DROP OPERATOR <->(scalar8, scalar8); +DROP FUNCTION _vchord_scalar8_operator_cosine; +DROP FUNCTION _vchord_scalar8_operator_ip; +DROP FUNCTION _vchord_scalar8_operator_l2; +DROP FUNCTION _vchord_scalar8_sphere_cosine_in; +DROP FUNCTION _vchord_scalar8_sphere_ip_in; +DROP FUNCTION _vchord_scalar8_sphere_l2_in; +DROP TYPE sphere_scalar8; +ALTER TYPE scalar8 SET (TYPMOD_IN = NONE); +DROP FUNCTION _vchord_typmod_in_65535; +ALTER TYPE scalar8 SET (TYPMOD_OUT = NONE); +DROP FUNCTION _vchord_typmod_out; +ALTER TYPE scalar8 SET (RECEIVE = NONE); +DROP FUNCTION _vchord_scalar8_recv; +ALTER TYPE scalar8 SET (SEND = NONE); +DROP FUNCTION _vchord_scalar8_send; +DROP TYPE scalar8 CASCADE; +-- DROP FUNCTION _vchord_scalar8_in; +-- DROP FUNCTION _vchord_scalar8_out; + +CREATE FUNCTION _vchord_rabitq4_in(input cstring, oid oid, typmod INT) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_cosine(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_cosine_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_ip(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_ip_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_l2(lhs rabitq4, rhs rabitq4) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_l2_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_out(vector rabitq4) RETURNS cstring +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_out_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_recv(internal internal, oid oid, typmod INT) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_recv_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_send(vector rabitq4) RETURNS bytea +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_send_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_cosine_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_cosine_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_ip_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_ip_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_sphere_l2_in(lhs rabitq4, rhs sphere_rabitq4) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_sphere_l2_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_typmod_in(list cstring[]) RETURNS INT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_typmod_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_in(input cstring, oid oid, typmod INT) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_cosine(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_cosine_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_ip(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_ip_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_operator_l2(lhs rabitq8, rhs rabitq8) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_l2_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_out(vector rabitq8) RETURNS cstring +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_out_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_recv(internal internal, oid oid, typmod INT) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_recv_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_send(vector rabitq8) RETURNS bytea +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_send_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_cosine_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_cosine_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_ip_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_ip_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_sphere_l2_in(lhs rabitq8, rhs sphere_rabitq8) RETURNS bool +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_sphere_l2_in_wrapper'; + +CREATE FUNCTION _vchord_rabitq8_typmod_in(list cstring[]) RETURNS INT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_typmod_in_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq4_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq4_l2_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_ip_ops_wrapper'; + +CREATE FUNCTION _vchordg_support_rabitq8_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_support_rabitq8_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_ip_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq4_maxsim_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq4_maxsim_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_cosine_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_cosine_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_ip_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_ip_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_l2_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_l2_ops_wrapper'; + +CREATE FUNCTION _vchordrq_support_rabitq8_maxsim_ops() RETURNS TEXT +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_support_rabitq8_maxsim_ops_wrapper'; + +DROP OPERATOR <<=>>(sphere_halfvec, halfvec); +DROP OPERATOR <<#>>(sphere_halfvec, halfvec); +DROP OPERATOR <<->>(sphere_halfvec, halfvec); +DROP OPERATOR <<=>>(sphere_vector, vector); +DROP OPERATOR <<#>>(sphere_vector, vector); +DROP OPERATOR <<->>(sphere_vector, vector); + +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +-- List of operator families + +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; + +-- List of operator classes + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); diff --git a/sql/upgrade/vchord--1.1.0--1.1.1.sql b/sql/upgrade/vchord--1.1.0--1.1.1.sql new file mode 100644 index 00000000..6490df1e --- /dev/null +++ b/sql/upgrade/vchord--1.1.0--1.1.1.sql @@ -0,0 +1,13 @@ +-- List of functions + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; diff --git a/sql/upgrade/vchord--1.1.1--1.2.0.sql b/sql/upgrade/vchord--1.1.1--1.2.0.sql new file mode 100644 index 00000000..b0bbb785 --- /dev/null +++ b/sql/upgrade/vchord--1.1.1--1.2.0.sql @@ -0,0 +1,1587 @@ +-- Preserve the post-1.1.1 composite-type compatibility fix when upgrading. + +CREATE OR REPLACE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; + +CREATE OR REPLACE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; + +CREATE OR REPLACE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; + +CREATE OR REPLACE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); diff --git a/src/algorithm/build.rs b/src/algorithm/build.rs deleted file mode 100644 index 4c7450af..00000000 --- a/src/algorithm/build.rs +++ /dev/null @@ -1,357 +0,0 @@ -use crate::algorithm::rabitq; -use crate::algorithm::tuples::*; -use crate::postgres::BufferWriteGuard; -use crate::postgres::Relation; -use crate::types::RabbitholeIndexingOptions; -use base::distance::DistanceKind; -use base::index::VectorOptions; -use base::scalar::ScalarLike; -use base::search::Pointer; -use common::vec2::Vec2; -use rand::Rng; -use rkyv::ser::serializers::AllocSerializer; -use std::marker::PhantomData; -use std::sync::atomic::AtomicBool; -use std::sync::Arc; - -pub trait HeapRelation { - fn traverse(&self, callback: F) - where - F: FnMut((Pointer, Vec)); -} - -pub fn build( - vector_options: VectorOptions, - rabbithole_options: RabbitholeIndexingOptions, - heap_relation: T, - index_relation: Relation, -) { - let dims = vector_options.dims; - let is_residual = - rabbithole_options.residual_quantization && vector_options.d == DistanceKind::L2; - let samples = { - let mut rand = rand::thread_rng(); - let max_number_of_samples = rabbithole_options.nlist.saturating_mul(256); - let mut samples = Vec::new(); - let mut number_of_samples = 0_u32; - heap_relation.traverse(|(_, vector)| { - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions",); - let vector = rabitq::project(&vector); - if number_of_samples < max_number_of_samples { - samples.extend(vector); - number_of_samples += 1; - } else { - let index = rand.gen_range(0..max_number_of_samples) as usize; - let start = index * dims as usize; - let end = start + dims as usize; - samples[start..end].copy_from_slice(&vector); - } - }); - Vec2::from_vec((number_of_samples as _, dims as _), samples) - }; - let structure = Structure::compute(vector_options.clone(), rabbithole_options.clone(), samples); - let h2_len = structure.h2_len(); - let h1_len = structure.h1_len(); - let mut meta = Tape::create(&index_relation); - assert_eq!(meta.first(), 0); - let mut vectors = Tape::create(&index_relation); - assert_eq!(vectors.first(), 1); - let h2_means = (0..h2_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h2_means(i).clone(), - }) - }) - .collect::>(); - let h1_means = (0..h1_len) - .map(|i| { - vectors.push(&VectorTuple { - payload: None, - vector: structure.h1_means(i).clone(), - }) - }) - .collect::>(); - let h1_firsts = (0..h1_len) - .map(|_| { - let tape = Tape::::create(&index_relation); - tape.first() - }) - .collect::>(); - let h2_firsts = (0..h2_len) - .map(|i| { - let mut tape = Tape::::create(&index_relation); - let mut cache = Vec::new(); - let h2_mean = structure.h2_means(i); - let children = structure.h2_children(i); - for child in children.iter().copied() { - let h1_mean = structure.h1_means(child); - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(h1_mean, h2_mean)) - } else { - rabitq::code(dims, h1_mean) - }; - cache.push((child, code)); - if cache.len() == 32 { - let group = std::mem::take(&mut cache); - let code = std::array::from_fn(|i| group[i].1.clone()); - let packed = rabitq::pack_codes(dims, code); - tape.push(&Height1Tuple { - mask: [true; 32], - mean: std::array::from_fn(|i| h1_means[group[i].0 as usize]), - first: std::array::from_fn(|i| h1_firsts[group[i].0 as usize]), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - } - if !cache.is_empty() { - let group = std::mem::take(&mut cache); - let codes = std::array::from_fn(|i| { - if i < group.len() { - group[i].1.clone() - } else { - rabitq::dummy_code(dims) - } - }); - let packed = rabitq::pack_codes(dims, codes); - tape.push(&Height1Tuple { - mask: std::array::from_fn(|i| i < group.len()), - mean: std::array::from_fn(|i| { - if i < group.len() { - h1_means[group[i].0 as usize] - } else { - Default::default() - } - }), - first: std::array::from_fn(|i| { - if i < group.len() { - h1_firsts[group[i].0 as usize] - } else { - Default::default() - } - }), - dis_u_2: packed.dis_u_2, - factor_ppc: packed.factor_ppc, - factor_ip: packed.factor_ip, - factor_err: packed.factor_err, - t: packed.t, - }); - } - tape.first() - }) - .collect::>(); - meta.push(&MetaTuple { - dims, - is_residual, - vectors_first: vectors.first(), - mean: h2_means[0], - first: h2_firsts[0], - }); - drop(meta); - let mut heads = Vec::new(); - for i in 0..structure.h1_len() { - heads.push(h1_firsts[i as usize]); - } - heap_relation.traverse(|(payload, vector)| { - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = rabitq::project(&vector); - let h0_vector = vectors.push(&VectorTuple { - vector: vector.clone(), - payload: Some(payload.as_u64()), - }); - let h0_payload = payload.as_u64(); - let h2_id = 0_u32; - let h1_id = { - let mut target = (0_u32, f32::INFINITY); - for &i in structure.h2_children(h2_id) { - let dis = f32::reduce_sum_of_d2(&vector, structure.h1_means(i)); - if dis < target.1 { - target = (i, dis); - } - } - target.0 - }; - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, structure.h1_means(h1_id))) - } else { - rabitq::code(dims, &vector) - }; - let mut write = index_relation.write(heads[h1_id as usize]); - let page = write.get_mut(); - if page.len() != 0 { - let flag = put( - page.get_mut(page.len()).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - if flag { - return; - } - } - let tuple = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mask: [false; 32], - mean: [(0, 0); 32], - payload: [0; 32], - dis_u_2: [0.0; 32], - factor_ppc: [0.0; 32], - factor_ip: [0.0; 32], - factor_err: [0.0; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], - }) - .unwrap(); - if let Some(i) = page.alloc(&tuple) { - let flag = put( - page.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } - let mut extend = index_relation.extend(); - heads[h1_id as usize] = extend.id(); - page.get_opaque_mut().next = extend.id(); - let page = extend.get_mut(); - if let Some(i) = page.alloc(&tuple) { - let flag = put( - page.get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } - panic!("a tuple cannot even be fit in a fresh page") - }); -} - -struct Structure { - h2_mean: Vec, - h2_children: Vec, - h1_means: Vec>, - h1_children: Vec>, -} - -impl Structure { - fn compute( - vector_options: VectorOptions, - rabbithole_options: RabbitholeIndexingOptions, - samples: Vec2, - ) -> Self { - let dims = vector_options.dims; - let h1_means = base::parallelism::RayonParallelism::scoped( - rabbithole_options.build_threads as _, - Arc::new(AtomicBool::new(false)), - |parallelism| { - let raw = k_means::k_means( - parallelism, - rabbithole_options.nlist as usize, - samples, - rabbithole_options.spherical_centroids, - 10, - false, - ); - let mut centroids = Vec::new(); - for i in 0..rabbithole_options.nlist { - centroids.push(raw[(i as usize,)].to_vec()); - } - centroids - }, - ) - .expect("k_means panics") - .expect("k_means interrupted"); - let h2_mean = { - let mut centroid = vec![0.0; dims as _]; - for i in 0..rabbithole_options.nlist { - for j in 0..dims { - centroid[j as usize] += h1_means[i as usize][j as usize]; - } - } - for j in 0..dims { - centroid[j as usize] /= rabbithole_options.nlist as f32; - } - centroid - }; - Structure { - h2_mean, - h2_children: (0..rabbithole_options.nlist).collect(), - h1_means, - h1_children: (0..rabbithole_options.nlist).map(|_| Vec::new()).collect(), - } - } - fn h2_len(&self) -> u32 { - 1 - } - fn h2_means(&self, i: u32) -> &Vec { - assert!(i == 0); - &self.h2_mean - } - fn h2_children(&self, i: u32) -> &Vec { - assert!(i == 0); - &self.h2_children - } - fn h1_len(&self) -> u32 { - self.h1_means.len() as _ - } - fn h1_means(&self, i: u32) -> &Vec { - &self.h1_means[i as usize] - } - #[allow(dead_code)] - fn h1_children(&self, i: u32) -> &Vec { - &self.h1_children[i as usize] - } -} - -struct Tape<'a, T> { - relation: &'a Relation, - head: BufferWriteGuard, - first: u32, - _phantom: PhantomData T>, -} - -impl<'a, T> Tape<'a, T> { - fn create(relation: &'a Relation) -> Self { - let head = relation.extend(); - let first = head.id(); - Self { - relation, - head, - first, - _phantom: PhantomData, - } - } - fn first(&self) -> u32 { - self.first - } -} - -impl<'a, T> Tape<'a, T> -where - T: rkyv::Serialize>, -{ - fn push(&mut self, x: &T) -> (u32, u16) { - let bytes = rkyv::to_bytes(x).expect("failed to serialize"); - if let Some(i) = self.head.get_mut().alloc(&bytes) { - (self.head.id(), i) - } else { - let next = self.relation.extend(); - self.head.get_mut().get_opaque_mut().next = next.id(); - self.head = next; - if let Some(i) = self.head.get_mut().alloc(&bytes) { - (self.head.id(), i) - } else { - panic!("tuple is too large to fit in a fresh page") - } - } - } -} diff --git a/src/algorithm/insert.rs b/src/algorithm/insert.rs deleted file mode 100644 index d40da899..00000000 --- a/src/algorithm/insert.rs +++ /dev/null @@ -1,194 +0,0 @@ -use crate::algorithm::rabitq; -use crate::algorithm::tuples::*; -use crate::postgres::Relation; -use base::distance::Distance; -use base::scalar::ScalarLike; -use base::search::Pointer; - -pub fn insert(relation: Relation, payload: Pointer, vector: Vec) { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get() - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = rabitq::project(&vector); - let is_residual = meta_tuple.is_residual; - let h0_vector = { - let tuple = rkyv::to_bytes::<_, 8192>(&VectorTuple { - vector: vector.clone(), - payload: Some(payload.as_u64()), - }) - .unwrap(); - let mut current = meta_tuple.vectors_first; - loop { - let read = relation.read(current); - let flag = 'flag: { - if read.get().freespace() as usize >= tuple.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - if let Some(i) = write.get_mut().alloc(&tuple) { - break (current, i); - } - if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&tuple) { - break (extend.id(), i); - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - } - }; - let h0_payload = payload.as_u64(); - let list = (meta_tuple.first,); - let list = { - let mut result = None::<(Distance, u32, Option>)>; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { - let h1_tuple = h1_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - let vector_guard = relation.read(h1_tuple.mean[j].0); - let vector_tuple = vector_guard - .get() - .get(h1_tuple.mean[j].1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis = Distance::from_f32(f32::reduce_sum_of_d2( - &vector, - &vector_tuple.vector, - )); - if result.is_none() || dis < result.as_ref().unwrap().0 { - result = Some(( - dis, - h1_tuple.first[j], - if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - )); - } - } - } - } - current = h1_guard.get().get_opaque().next; - } - let result = result.unwrap(); - (result.1, result.2) - }; - let code = if is_residual { - rabitq::code(dims, &f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - rabitq::code(dims, &vector) - }; - let dummy = rkyv::to_bytes::<_, 8192>(&Height0Tuple { - mask: [false; 32], - mean: [(0, 0); 32], - payload: [0; 32], - dis_u_2: [0.0f32; 32], - factor_ppc: [0.0f32; 32], - factor_ip: [0.0f32; 32], - factor_err: [0.0f32; 32], - t: vec![0; (dims.div_ceil(4) * 16) as usize], - }) - .unwrap(); - let first = list.0; - assert!(first != u32::MAX); - let mut current = first; - loop { - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.get().len() { - let h0_tuple = read - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if h0_tuple.mask.iter().any(|x| *x) { - break 'flag true; - } - } - if read.get().freespace() as usize >= dummy.len() { - break 'flag true; - } - if read.get().get_opaque().next == u32::MAX { - break 'flag true; - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - for i in 1..=write.get().len() { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - if flag { - return; - } - } - if let Some(i) = write.get_mut().alloc(&dummy) { - let flag = put( - write.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } - if write.get().get_opaque().next == u32::MAX { - let mut extend = relation.extend(); - write.get_mut().get_opaque_mut().next = extend.id(); - if let Some(i) = extend.get_mut().alloc(&dummy) { - let flag = put( - extend.get_mut().get_mut(i).expect("data corruption"), - dims, - &code, - h0_vector, - h0_payload, - ); - assert!(flag, "a put fails even on a fresh tuple"); - return; - } else { - panic!("a tuple cannot even be fit in a fresh page"); - } - } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - } -} diff --git a/src/algorithm/mod.rs b/src/algorithm/mod.rs deleted file mode 100644 index 6fcddd1e..00000000 --- a/src/algorithm/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub mod build; -pub mod insert; -pub mod rabitq; -pub mod scan; -pub mod tuples; -pub mod vacuum; diff --git a/src/algorithm/rabitq.rs b/src/algorithm/rabitq.rs deleted file mode 100644 index 45b512be..00000000 --- a/src/algorithm/rabitq.rs +++ /dev/null @@ -1,204 +0,0 @@ -use base::distance::Distance; -use base::scalar::ScalarLike; -use nalgebra::DMatrix; -use quantization::utils::InfiniteByteChunks; -use std::sync::OnceLock; - -fn random_matrix(n: usize) -> DMatrix { - use rand::{Rng, SeedableRng}; - use rand_chacha::ChaCha12Rng; - use rand_distr::StandardNormal; - let mut rng = ChaCha12Rng::from_seed([7; 32]); - DMatrix::from_fn(n, n, |_, _| rng.sample(StandardNormal)) -} - -#[ignore] -#[test] -fn check_all_matrixs_are_full_rank() { - let parallelism = std::thread::available_parallelism().unwrap().get(); - std::thread::scope(|scope| { - let mut threads = vec![]; - for remainder in 0..parallelism { - threads.push(scope.spawn(move || { - for n in (0..=2000).filter(|x| x % parallelism == remainder) { - let matrix = random_matrix(n); - assert!(matrix.is_invertible()); - } - })); - } - for thread in threads { - thread.join().unwrap(); - } - }); -} - -fn orthogonal_matrix(n: usize) -> Vec> { - use nalgebra::QR; - let matrix = random_matrix(n); - // QR decomposition is unique if the matrix is full rank - let qr = QR::new(matrix); - let q = qr.q(); - let mut projection = Vec::new(); - for row in q.row_iter() { - projection.push(row.iter().copied().collect::>()); - } - projection -} - -static MATRIXS: [OnceLock>>; 1 + 2000] = [const { OnceLock::new() }; 1 + 2000]; - -pub fn project(vector: &[f32]) -> Vec { - use base::scalar::ScalarLike; - let n = vector.len(); - let matrix = MATRIXS[n].get_or_init(|| orthogonal_matrix(n)); - (0..n) - .map(|i| f32::reduce_sum_of_xy(vector, &matrix[i])) - .collect() -} - -#[derive(Debug, Clone)] -pub struct Code { - pub dis_u_2: f32, - pub factor_ppc: f32, - pub factor_ip: f32, - pub factor_err: f32, - pub signs: Vec, -} - -pub fn code(dims: u32, vector: &[f32]) -> Code { - let sum_of_abs_x = f32::reduce_sum_of_abs_x(vector); - let sum_of_x_2 = f32::reduce_sum_of_x2(vector); - let dis_u = sum_of_x_2.sqrt(); - let x0 = sum_of_abs_x / (sum_of_x_2 * (dims as f32)).sqrt(); - let x_x0 = dis_u / x0; - let fac_norm = (dims as f32).sqrt(); - let max_x1 = 1.0f32 / (dims as f32 - 1.0).sqrt(); - let factor_err = 2.0f32 * max_x1 * (x_x0 * x_x0 - dis_u * dis_u).sqrt(); - let factor_ip = -2.0f32 / fac_norm * x_x0; - let cnt_pos = vector - .iter() - .map(|x| x.is_sign_positive() as i32) - .sum::(); - let cnt_neg = vector - .iter() - .map(|x| x.is_sign_negative() as i32) - .sum::(); - let factor_ppc = factor_ip * (cnt_pos - cnt_neg) as f32; - let mut signs = Vec::new(); - for i in 0..dims { - signs.push(vector[i as usize].is_sign_positive() as u8); - } - Code { - dis_u_2: sum_of_x_2, - factor_ppc, - factor_ip, - factor_err, - signs, - } -} - -pub fn dummy_code(dims: u32) -> Code { - Code { - dis_u_2: 0.0, - factor_ppc: 0.0, - factor_ip: 0.0, - factor_err: 0.0, - signs: vec![0; dims as _], - } -} - -pub struct PackedCodes { - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn pack_codes(dims: u32, codes: [Code; 32]) -> PackedCodes { - PackedCodes { - dis_u_2: std::array::from_fn(|i| codes[i].dis_u_2), - factor_ppc: std::array::from_fn(|i| codes[i].factor_ppc), - factor_ip: std::array::from_fn(|i| codes[i].factor_ip), - factor_err: std::array::from_fn(|i| codes[i].factor_err), - t: { - let signs = codes.map(|code| { - InfiniteByteChunks::new(code.signs.into_iter()) - .map(|[b0, b1, b2, b3]| b0 | b1 << 1 | b2 << 2 | b3 << 3) - .take(dims.div_ceil(4) as usize) - .collect::>() - }); - quantization::fast_scan::b4::pack(dims.div_ceil(4), signs).collect() - }, - } -} - -pub fn fscan_preprocess(vector: &[f32]) -> (f32, f32, f32, f32, Vec) { - use quantization::quantize; - let dis_v_2 = f32::reduce_sum_of_x2(vector); - let (k, b, qvector) = quantize::quantize::<15>(vector); - let qvector_sum = if vector.len() <= 4369 { - quantize::reduce_sum_of_x_as_u16(&qvector) as f32 - } else { - quantize::reduce_sum_of_x_as_u32(&qvector) as f32 - }; - (dis_v_2, b, k, qvector_sum, compress(qvector)) -} - -pub fn fscan_process_lowerbound( - dims: u32, - lut: &(f32, f32, f32, f32, Vec), - (dis_u_2, factor_ppc, factor_ip, factor_err, t): ( - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[f32; 32], - &[u8], - ), -) -> [Distance; 32] { - let &(dis_v_2, b, k, qvector_sum, ref s) = lut; - let r = quantization::fast_scan::b4::fast_scan_b4(dims.div_ceil(4), t, s); - std::array::from_fn(|i| { - let rough = dis_u_2[i] - + dis_v_2 - + b * factor_ppc[i] - + ((2.0 * r[i] as f32) - qvector_sum) * factor_ip[i] * k; - let err = factor_err[i] * dis_v_2.sqrt(); - Distance::from_f32(rough - 1.9 * err) - }) -} - -fn compress(mut qvector: Vec) -> Vec { - let dims = qvector.len() as u32; - let width = dims.div_ceil(4); - qvector.resize(qvector.len().next_multiple_of(4), 0); - let mut t = vec![0u8; width as usize * 16]; - for i in 0..width as usize { - unsafe { - // this hint is used to skip bound checks - std::hint::assert_unchecked(4 * i + 3 < qvector.len()); - std::hint::assert_unchecked(16 * i + 15 < t.len()); - } - let t0 = qvector[4 * i + 0]; - let t1 = qvector[4 * i + 1]; - let t2 = qvector[4 * i + 2]; - let t3 = qvector[4 * i + 3]; - t[16 * i + 0b0000] = 0; - t[16 * i + 0b0001] = t0; - t[16 * i + 0b0010] = t1; - t[16 * i + 0b0011] = t1 + t0; - t[16 * i + 0b0100] = t2; - t[16 * i + 0b0101] = t2 + t0; - t[16 * i + 0b0110] = t2 + t1; - t[16 * i + 0b0111] = t2 + t1 + t0; - t[16 * i + 0b1000] = t3; - t[16 * i + 0b1001] = t3 + t0; - t[16 * i + 0b1010] = t3 + t1; - t[16 * i + 0b1011] = t3 + t1 + t0; - t[16 * i + 0b1100] = t3 + t2; - t[16 * i + 0b1101] = t3 + t2 + t0; - t[16 * i + 0b1110] = t3 + t2 + t1; - t[16 * i + 0b1111] = t3 + t2 + t1 + t0; - } - t -} diff --git a/src/algorithm/scan.rs b/src/algorithm/scan.rs deleted file mode 100644 index fc1d3718..00000000 --- a/src/algorithm/scan.rs +++ /dev/null @@ -1,186 +0,0 @@ -use crate::algorithm::rabitq; -use crate::algorithm::tuples::*; -use crate::postgres::Relation; -use base::always_equal::AlwaysEqual; -use base::distance::Distance; -use base::scalar::ScalarLike; -use base::search::Pointer; -use std::cmp::Reverse; -use std::collections::BinaryHeap; - -pub fn scan( - relation: Relation, - vector: Vec, - h1_nprobe: u32, -) -> impl Iterator { - assert!(h1_nprobe >= 1); - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get() - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dims = meta_tuple.dims; - assert_eq!(dims as usize, vector.len(), "invalid vector dimensions"); - let vector = rabitq::project(&vector); - let is_residual = meta_tuple.is_residual; - let default_lut = if !is_residual { - Some(rabitq::fscan_preprocess(&vector)) - } else { - None - }; - let lists: Vec<_> = vec![( - meta_tuple.first, - if is_residual { - let vector_guard = relation.read(meta_tuple.mean.0); - let vector_tuple = vector_guard - .get() - .get(meta_tuple.mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - Some(vector_tuple.vector.to_vec()) - } else { - None - }, - )]; - let lists: Vec<_> = { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { - let h1_tuple = h1_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( - dims, - lut, - ( - &h1_tuple.dis_u_2, - &h1_tuple.factor_ppc, - &h1_tuple.factor_ip, - &h1_tuple.factor_err, - &h1_tuple.t, - ), - ); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h1_tuple.mean[j]), - AlwaysEqual(h1_tuple.first[j]), - )); - } - } - } - current = h1_guard.get().get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _, _)>::new(); - std::iter::from_fn(|| { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(first)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let vector_tuple = vector_guard - .get() - .get(mean.1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); - cache.push(( - Reverse(dis_u), - AlwaysEqual(first), - AlwaysEqual(if is_residual { - Some(vector_tuple.vector.to_vec()) - } else { - None - }), - )); - } - let (_, AlwaysEqual(first), AlwaysEqual(mean)) = cache.pop()?; - Some((first, mean)) - }) - .take(h1_nprobe as usize) - .collect() - }; - { - let mut results = Vec::new(); - for list in lists { - let lut = if is_residual { - &rabitq::fscan_preprocess(&f32::vector_sub(&vector, list.1.as_ref().unwrap())) - } else { - default_lut.as_ref().unwrap() - }; - let mut current = list.0; - while current != u32::MAX { - let h0_guard = relation.read(current); - for i in 1..=h0_guard.get().len() { - let h0_tuple = h0_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let lowerbounds = rabitq::fscan_process_lowerbound( - dims, - lut, - ( - &h0_tuple.dis_u_2, - &h0_tuple.factor_ppc, - &h0_tuple.factor_ip, - &h0_tuple.factor_err, - &h0_tuple.t, - ), - ); - for j in 0..32 { - if h0_tuple.mask[j] { - results.push(( - Reverse(lowerbounds[j]), - AlwaysEqual(h0_tuple.mean[j]), - AlwaysEqual(h0_tuple.payload[j]), - )); - } - } - } - current = h0_guard.get().get_opaque().next; - } - } - let mut heap = BinaryHeap::from(results); - let mut cache = BinaryHeap::<(Reverse, _)>::new(); - std::iter::from_fn(move || { - while !heap.is_empty() && heap.peek().map(|x| x.0) > cache.peek().map(|x| x.0) { - let (_, AlwaysEqual(mean), AlwaysEqual(pay_u)) = heap.pop().unwrap(); - let vector_guard = relation.read(mean.0); - let Some(vector_tuple) = vector_guard.get().get(mean.1) else { - // fails consistency check - continue; - }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if vector_tuple.payload != Some(pay_u) { - // fails consistency check - continue; - } - let dis_u = - Distance::from_f32(f32::reduce_sum_of_d2(&vector, &vector_tuple.vector)); - cache.push((Reverse(dis_u), AlwaysEqual(pay_u))); - } - let (Reverse(dis_u), AlwaysEqual(pay_u)) = cache.pop()?; - Some((dis_u, Pointer::new(pay_u))) - }) - } -} diff --git a/src/algorithm/tuples.rs b/src/algorithm/tuples.rs deleted file mode 100644 index aaf3380c..00000000 --- a/src/algorithm/tuples.rs +++ /dev/null @@ -1,137 +0,0 @@ -use crate::algorithm::rabitq; -use rkyv::{Archive, Deserialize, Serialize}; - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct MetaTuple { - pub dims: u32, - pub is_residual: bool, - pub vectors_first: u32, - // raw vector - pub mean: (u32, u16), - // for meta tuple, it's pointers to next level - pub first: u32, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct VectorTuple { - pub vector: Vec, - // this field is saved only for vacuum - pub payload: Option, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height1Tuple { - pub mask: [bool; 32], - // raw vector - pub mean: [(u32, u16); 32], - // for height 1 tuple, it's pointers to next level - pub first: [u32; 32], - // RaBitQ algoithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -#[derive(Clone, PartialEq, Archive, Serialize, Deserialize)] -#[archive(check_bytes)] -pub struct Height0Tuple { - pub mask: [bool; 32], - // raw vector - pub mean: [(u32, u16); 32], - // for height 0 tuple, it's pointers to heap relation - pub payload: [u64; 32], - // RaBitQ algoithm - pub dis_u_2: [f32; 32], - pub factor_ppc: [f32; 32], - pub factor_ip: [f32; 32], - pub factor_err: [f32; 32], - pub t: Vec, -} - -pub fn put( - bytes: &mut [u8], - dims: u32, - code: &rabitq::Code, - vector: (u32, u16), - payload: u64, -) -> bool { - // todo: use mutable api - let mut x = rkyv::from_bytes::(bytes).expect("data corruption"); - for j in 0..32 { - if !x.mask[j] { - x.mean[j] = vector; - x.payload[j] = payload; - x.mask[j] = true; - x.dis_u_2[j] = code.dis_u_2; - x.factor_ppc[j] = code.factor_ppc; - x.factor_ip[j] = code.factor_ip; - x.factor_err[j] = code.factor_err; - let width = dims.div_ceil(4) as usize; - let table = [ - (0, 0), - (2, 0), - (4, 0), - (6, 0), - (8, 0), - (10, 0), - (12, 0), - (14, 0), - (1, 0), - (3, 0), - (5, 0), - (7, 0), - (9, 0), - (11, 0), - (13, 0), - (15, 0), - (0, 1), - (2, 1), - (4, 1), - (6, 1), - (8, 1), - (10, 1), - (12, 1), - (14, 1), - (1, 1), - (3, 1), - (5, 1), - (7, 1), - (9, 1), - (11, 1), - (13, 1), - (15, 1), - ]; - let pos = table[j].0; - let mask = match table[j].1 { - 0 => 0xf0, - 1 => 0x0f, - _ => unreachable!(), - }; - let shift = match table[j].1 { - 0 => 0, - 1 => 4, - _ => unreachable!(), - }; - let mut buffer = vec![0u8; width]; - for j in 0..width { - let b0 = code.signs.get(4 * j + 0).copied().unwrap_or_default(); - let b1 = code.signs.get(4 * j + 1).copied().unwrap_or_default(); - let b2 = code.signs.get(4 * j + 2).copied().unwrap_or_default(); - let b3 = code.signs.get(4 * j + 3).copied().unwrap_or_default(); - buffer[j] = b0 | b1 << 1 | b2 << 2 | b3 << 3; - } - for j in 0..width { - x.t[16 * j + pos] &= mask; - x.t[16 * j + pos] |= buffer[j] << shift; - } - bytes.copy_from_slice(&rkyv::to_bytes::<_, 8192>(&x).unwrap()); - return true; - } - } - false -} diff --git a/src/algorithm/vacuum.rs b/src/algorithm/vacuum.rs deleted file mode 100644 index 880a4bc6..00000000 --- a/src/algorithm/vacuum.rs +++ /dev/null @@ -1,138 +0,0 @@ -use crate::algorithm::tuples::*; -use crate::postgres::Relation; -use base::search::Pointer; - -pub fn vacuum(relation: Relation, callback: impl Fn(Pointer) -> bool) { - // step 1: vacuum height_0_tuple - { - let h1_firsts = { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get() - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - vec![meta_tuple.first] - }; - let h0_firsts = { - let mut results = Vec::new(); - for first in h1_firsts { - let mut current = first; - while current != u32::MAX { - let h1_guard = relation.read(current); - for i in 1..=h1_guard.get().len() { - let h1_tuple = h1_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if h1_tuple.mask[j] { - results.push(h1_tuple.first[j]); - } - } - } - current = h1_guard.get().get_opaque().next; - } - } - results - }; - for first in h0_firsts { - let mut current = first; - while current != u32::MAX { - let mut h0_guard = relation.write(current); - for i in 1..=h0_guard.get().len() { - let h0_tuple = h0_guard - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - let flag = 'flag: { - for j in 0..32 { - if h0_tuple.mask[j] && callback(Pointer::new(h0_tuple.payload[j])) { - break 'flag true; - } - } - false - }; - if flag { - // todo: use mutable API - let mut temp = h0_guard - .get() - .get(i) - .map(rkyv::from_bytes::) - .expect("data corruption") - .expect("data corruption"); - for j in 0..32 { - if temp.mask[j] && callback(Pointer::new(temp.payload[j])) { - temp.mask[j] = false; - } - } - let temp = rkyv::to_bytes::<_, 8192>(&temp).expect("failed to serialize"); - h0_guard - .get_mut() - .get_mut(i) - .expect("data corruption") - .copy_from_slice(&temp); - } - } - // todo: cross-tuple vacuum so that we can skip a tuple - current = h0_guard.get().get_opaque().next; - } - } - } - // step 2: vacuum vector_tuple - { - let mut current = { - let meta_guard = relation.read(0); - let meta_tuple = meta_guard - .get() - .get(1) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - meta_tuple.vectors_first - }; - while current != u32::MAX { - let read = relation.read(current); - let flag = 'flag: { - for i in 1..=read.get().len() { - let vector_tuple = read - .get() - .get(i) - .map(rkyv::check_archived_root::) - .expect("data corruption") - .expect("data corruption"); - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { - break 'flag true; - } - } - } - false - }; - if flag { - drop(read); - let mut write = relation.write(current); - for i in 1..=write.get().len() { - let Some(vector_tuple) = write.get().get(i) else { - continue; - }; - let vector_tuple = rkyv::check_archived_root::(vector_tuple) - .expect("data corruption"); - if let Some(payload) = vector_tuple.payload.as_ref().copied() { - if callback(Pointer::new(payload)) { - write.get_mut().free(i); - } - } - } - current = write.get().get_opaque().next; - } else { - current = read.get().get_opaque().next; - } - } - } -} diff --git a/src/bin/pgrx_embed.rs b/src/bin/pgrx_embed.rs index 5f5c4d85..3307dd2a 100644 --- a/src/bin/pgrx_embed.rs +++ b/src/bin/pgrx_embed.rs @@ -1 +1,54 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#![allow(unsafe_code)] + ::pgrx::pgrx_embed!(); + +#[macro_export] +macro_rules! schema_generation { + ($($symbol:ident)*; $($import:ident)*) => { + pub fn main() -> Result<(), Box> { + $( + const _: () = { + #[unsafe(no_mangle)] + unsafe extern "C" fn $import() { + panic!("{} is called unexpectedly.", stringify!($import)); + } + }; + )* + + extern crate vchord as _; + + use ::pgrx::pgrx_sql_entity_graph::ControlFile; + use ::pgrx::pgrx_sql_entity_graph::PgrxSql; + use ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity; + + let p = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/vchord.control")); + let control_file = ControlFile::try_from(p)?; + + unsafe extern "Rust" { + $(safe fn $symbol() -> SqlGraphEntity;)* + } + + let mut e = vec![SqlGraphEntity::ExtensionRoot(control_file)]; + $(e.push($symbol());)* + + let pgrx_sql = PgrxSql::build(e.into_iter(), "vchord".to_string(), false)?; + pgrx_sql.write(&mut std::io::stdout())?; + + Ok(()) + } + }; +} diff --git a/src/datatype/binary_rabitq4.rs b/src/datatype/binary_rabitq4.rs new file mode 100644 index 00000000..b623c62d --- /dev/null +++ b/src/datatype/binary_rabitq4.rs @@ -0,0 +1,96 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::datum::Internal; +use pgrx::pg_sys::Oid; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_send(vector: Rabitq4Input<'_>) -> Vec { + let vector = vector.as_borrowed(); + let mut stream = Vec::::new(); + stream.extend(vector.dim().to_be_bytes()); + stream.extend(vector.sum_of_x2().to_be_bytes()); + stream.extend(vector.norm_of_lattice().to_be_bytes()); + stream.extend(vector.sum_of_code().to_be_bytes()); + stream.extend(vector.sum_of_abs_x().to_be_bytes()); + for &c in vector.packed_code() { + stream.extend(c.to_be_bytes()); + } + stream +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq4Output { + let _ = (oid, typmod); + let buf = unsafe { internal.get_mut::().unwrap() }; + + let dim = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + u32::from_be_bytes(raw) + }; + let sum_of_x2 = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let norm_of_lattice = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_code = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_abs_x = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let packed_code = { + let mut result = Vec::new(); + for _ in 0..dim.div_ceil(2) { + result.push({ + assert!(buf.cursor < i32::MAX - 1 && buf.cursor + 1 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 1]>().read() }; + buf.cursor += 1; + u8::from_be_bytes(raw) + }); + } + result + }; + + if let Some(x) = Rabitq4Borrowed::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq4Output::new(x) + } else { + pgrx::error!("detect data corruption"); + } +} diff --git a/src/datatype/binary_rabitq8.rs b/src/datatype/binary_rabitq8.rs new file mode 100644 index 00000000..023fed84 --- /dev/null +++ b/src/datatype/binary_rabitq8.rs @@ -0,0 +1,96 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use pgrx::datum::Internal; +use pgrx::pg_sys::Oid; +use vector::VectorBorrowed; +use vector::rabitq8::Rabitq8Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_send(vector: Rabitq8Input<'_>) -> Vec { + let vector = vector.as_borrowed(); + let mut stream = Vec::::new(); + stream.extend(vector.dim().to_be_bytes()); + stream.extend(vector.sum_of_x2().to_be_bytes()); + stream.extend(vector.norm_of_lattice().to_be_bytes()); + stream.extend(vector.sum_of_code().to_be_bytes()); + stream.extend(vector.sum_of_abs_x().to_be_bytes()); + for &c in vector.packed_code() { + stream.extend(c.to_be_bytes()); + } + stream +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_recv(mut internal: Internal, oid: Oid, typmod: i32) -> Rabitq8Output { + let _ = (oid, typmod); + let buf = unsafe { internal.get_mut::().unwrap() }; + + let dim = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + u32::from_be_bytes(raw) + }; + let sum_of_x2 = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let norm_of_lattice = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_code = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let sum_of_abs_x = { + assert!(buf.cursor < i32::MAX - 4 && buf.cursor + 4 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 4]>().read() }; + buf.cursor += 4; + f32::from_be_bytes(raw) + }; + let packed_code = { + let mut result = Vec::new(); + for _ in 0..dim.div_ceil(1) { + result.push({ + assert!(buf.cursor < i32::MAX - 1 && buf.cursor + 1 <= buf.len); + let raw = unsafe { buf.data.add(buf.cursor as _).cast::<[u8; 1]>().read() }; + buf.cursor += 1; + u8::from_be_bytes(raw) + }); + } + result + }; + + if let Some(x) = Rabitq8Borrowed::new_checked( + dim, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq8Output::new(x) + } else { + pgrx::error!("detect data corruption"); + } +} diff --git a/src/datatype/functions_rabitq4.rs b/src/datatype/functions_rabitq4.rs new file mode 100644 index 00000000..eb40e559 --- /dev/null +++ b/src/datatype/functions_rabitq4.rs @@ -0,0 +1,84 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use simd::{Floating, f16}; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; +use vector::vect::VectBorrowed; + +#[pgrx::pg_extern(sql = "")] +fn _vchord_vector_quantize_to_rabitq4(vector: VectorInput) -> Rabitq4Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = vector.slice().to_vec(); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Output::new(Rabitq4Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_halfvec_quantize_to_rabitq4(vector: HalfvecInput) -> Rabitq4Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = f16::vector_to_f32(vector.slice()); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Output::new(Rabitq4Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_dequantize_to_vector(vector: Rabitq4Input) -> VectorOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + VectorOutput::new(VectBorrowed::new(&result)) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_dequantize_to_halfvec(vector: Rabitq4Input) -> HalfvecOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + let result = f16::vector_from_f32(&result); + HalfvecOutput::new(VectBorrowed::new(&result)) +} diff --git a/src/datatype/functions_rabitq8.rs b/src/datatype/functions_rabitq8.rs new file mode 100644 index 00000000..b5c56cad --- /dev/null +++ b/src/datatype/functions_rabitq8.rs @@ -0,0 +1,84 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use simd::{Floating, f16}; +use vector::VectorBorrowed; +use vector::rabitq8::Rabitq8Borrowed; +use vector::vect::VectBorrowed; + +#[pgrx::pg_extern(sql = "")] +fn _vchord_vector_quantize_to_rabitq8(vector: VectorInput) -> Rabitq8Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = vector.slice().to_vec(); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); + Rabitq8Output::new(Rabitq8Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_halfvec_quantize_to_rabitq8(vector: HalfvecInput) -> Rabitq8Output { + let vector = vector.as_borrowed(); + let dim = vector.dim(); + let mut vector = f16::vector_to_f32(vector.slice()); + rabitq::rotate::rotate_inplace(&mut vector); + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); + Rabitq8Output::new(Rabitq8Borrowed::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + &elements, + )) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_dequantize_to_vector(vector: Rabitq8Input) -> VectorOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + VectorOutput::new(VectBorrowed::new(&result)) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_dequantize_to_halfvec(vector: Rabitq8Input) -> HalfvecOutput { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + let result = f16::vector_from_f32(&result); + HalfvecOutput::new(VectBorrowed::new(&result)) +} diff --git a/src/datatype/memory_halfvec.rs b/src/datatype/memory_halfvec.rs new file mode 100644 index 00000000..ca185bdd --- /dev/null +++ b/src/datatype/memory_halfvec.rs @@ -0,0 +1,259 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; +use simd::f16; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[repr(C)] +struct HalfvecHeader { + varlena: u32, + dim: u16, + unused: u16, + elements: [f16; 0], +} + +impl HalfvecHeader { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + size_of::() + size_of::() * len + } + unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f16> { + unsafe { + let this = this.as_ptr(); + VectBorrowed::new(std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dim).read() as usize, + )) + } + } +} + +pub struct HalfvecInput<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl HalfvecInput<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + HalfvecInput(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { + unsafe { HalfvecHeader::as_borrowed(self.0) } + } +} + +impl Drop for HalfvecInput<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct HalfvecOutput(NonNull); + +impl HalfvecOutput { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(HalfvecHeader::size_of(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Self(q) + } + pub fn new(vector: VectBorrowed<'_, f16>) -> Self { + unsafe { + let slice = vector.slice(); + let size = HalfvecHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut HalfvecHeader; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dim).write(vector.dim() as _); + (&raw mut (*ptr).unused).write(0); + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + (&raw mut (*ptr).elements).cast(), + slice.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f16> { + unsafe { HalfvecHeader::as_borrowed(self.0) } + } + fn into_raw(self) -> *mut HalfvecHeader { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for HalfvecOutput { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for HalfvecInput<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for HalfvecOutput { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for HalfvecOutput { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl<'a> pgrx::datum::UnboxDatum for HalfvecInput<'a> { + type As<'src> + = HalfvecInput<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for HalfvecOutput { + type As<'src> = HalfvecOutput; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for HalfvecInput<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +unsafe impl SqlTranslatable for HalfvecOutput { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("halfvec"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("halfvec")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for HalfvecInput<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxRet + +unsafe impl pgrx::callconv::BoxRet for HalfvecOutput { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/memory_rabitq4.rs b/src/datatype/memory_rabitq4.rs new file mode 100644 index 00000000..2f1e9433 --- /dev/null +++ b/src/datatype/memory_rabitq4.rs @@ -0,0 +1,276 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[repr(C)] +struct Rabitq4Header { + varlena: u32, + dim: u16, + unused: u16, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + elements: [u8; 0], +} + +impl Rabitq4Header { + fn size_of_by_dim(dim: usize) -> usize { + Self::size_of_by_len(dim.div_ceil(2)) + } + fn size_of_by_len(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + size_of::() + size_of::() * len + } + unsafe fn as_borrowed<'a>(this: NonNull) -> Rabitq4Borrowed<'a> { + unsafe { + let this = this.as_ptr(); + Rabitq4Borrowed::new( + (&raw const (*this).dim).read() as u32, + (&raw const (*this).sum_of_x2).read(), + (&raw const (*this).norm_of_lattice).read(), + (&raw const (*this).sum_of_code).read(), + (&raw const (*this).sum_of_abs_x).read(), + std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dim).read().div_ceil(2) as usize, + ), + ) + } + } +} + +pub struct Rabitq4Input<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl Rabitq4Input<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq4Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Rabitq4Input(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + unsafe { Rabitq4Header::as_borrowed(self.0) } + } +} + +impl Drop for Rabitq4Input<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct Rabitq4Output(NonNull); + +impl Rabitq4Output { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq4Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Self(q) + } + pub fn new(vector: Rabitq4Borrowed<'_>) -> Self { + unsafe { + let packed_code = vector.packed_code(); + let size = Rabitq4Header::size_of_by_len(packed_code.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut Rabitq4Header; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dim).write(vector.dim() as _); + (&raw mut (*ptr).unused).write(0); + (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); + (&raw mut (*ptr).norm_of_lattice).write(vector.norm_of_lattice()); + (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); + (&raw mut (*ptr).sum_of_abs_x).write(vector.sum_of_abs_x()); + std::ptr::copy_nonoverlapping( + packed_code.as_ptr(), + (&raw mut (*ptr).elements).cast(), + packed_code.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> Rabitq4Borrowed<'_> { + unsafe { Rabitq4Header::as_borrowed(self.0) } + } + fn into_raw(self) -> *mut Rabitq4Header { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for Rabitq4Output { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for Rabitq4Input<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for Rabitq4Output { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for Rabitq4Output { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl<'a> pgrx::datum::UnboxDatum for Rabitq4Input<'a> { + type As<'src> + = Rabitq4Input<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for Rabitq4Output { + type As<'src> = Rabitq4Output; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for Rabitq4Input<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq4"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq4")))) + } +} + +unsafe impl SqlTranslatable for Rabitq4Output { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq4"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq4")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Rabitq4Input<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxRet + +unsafe impl pgrx::callconv::BoxRet for Rabitq4Output { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/memory_rabitq8.rs b/src/datatype/memory_rabitq8.rs new file mode 100644 index 00000000..b5ef8a4d --- /dev/null +++ b/src/datatype/memory_rabitq8.rs @@ -0,0 +1,276 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::rabitq8::Rabitq8Borrowed; + +#[repr(C)] +struct Rabitq8Header { + varlena: u32, + dim: u16, + unused: u16, + sum_of_x2: f32, + norm_of_lattice: f32, + sum_of_code: f32, + sum_of_abs_x: f32, + elements: [u8; 0], +} + +impl Rabitq8Header { + fn size_of_by_dim(dim: usize) -> usize { + Self::size_of_by_len(dim.div_ceil(1)) + } + fn size_of_by_len(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + size_of::() + size_of::() * len + } + unsafe fn as_borrowed<'a>(this: NonNull) -> Rabitq8Borrowed<'a> { + unsafe { + let this = this.as_ptr(); + Rabitq8Borrowed::new( + (&raw const (*this).dim).read() as u32, + (&raw const (*this).sum_of_x2).read(), + (&raw const (*this).norm_of_lattice).read(), + (&raw const (*this).sum_of_code).read(), + (&raw const (*this).sum_of_abs_x).read(), + std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dim).read().div_ceil(1) as usize, + ), + ) + } + } +} + +pub struct Rabitq8Input<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl Rabitq8Input<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq8Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Rabitq8Input(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + unsafe { Rabitq8Header::as_borrowed(self.0) } + } +} + +impl Drop for Rabitq8Input<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct Rabitq8Output(NonNull); + +impl Rabitq8Output { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(Rabitq8Header::size_of_by_dim(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Self(q) + } + pub fn new(vector: Rabitq8Borrowed<'_>) -> Self { + unsafe { + let packed_code = vector.packed_code(); + let size = Rabitq8Header::size_of_by_len(packed_code.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut Rabitq8Header; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dim).write(vector.dim() as _); + (&raw mut (*ptr).unused).write(0); + (&raw mut (*ptr).sum_of_x2).write(vector.sum_of_x2()); + (&raw mut (*ptr).norm_of_lattice).write(vector.norm_of_lattice()); + (&raw mut (*ptr).sum_of_code).write(vector.sum_of_code()); + (&raw mut (*ptr).sum_of_abs_x).write(vector.sum_of_abs_x()); + std::ptr::copy_nonoverlapping( + packed_code.as_ptr(), + (&raw mut (*ptr).elements).cast(), + packed_code.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> Rabitq8Borrowed<'_> { + unsafe { Rabitq8Header::as_borrowed(self.0) } + } + fn into_raw(self) -> *mut Rabitq8Header { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for Rabitq8Output { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for Rabitq8Input<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for Rabitq8Output { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for Rabitq8Output { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl<'a> pgrx::datum::UnboxDatum for Rabitq8Input<'a> { + type As<'src> + = Rabitq8Input<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for Rabitq8Output { + type As<'src> = Rabitq8Output; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for Rabitq8Input<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq8"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq8")))) + } +} + +unsafe impl SqlTranslatable for Rabitq8Output { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("rabitq8"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("rabitq8")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Rabitq8Input<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxRet + +unsafe impl pgrx::callconv::BoxRet for Rabitq8Output { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/memory_vecf32.rs b/src/datatype/memory_vecf32.rs deleted file mode 100644 index 2802134c..00000000 --- a/src/datatype/memory_vecf32.rs +++ /dev/null @@ -1,222 +0,0 @@ -use base::vector::*; -use pgrx::datum::FromDatum; -use pgrx::datum::IntoDatum; -use pgrx::pg_sys::Datum; -use pgrx::pg_sys::Oid; -use pgrx::pgrx_sql_entity_graph::metadata::ArgumentError; -use pgrx::pgrx_sql_entity_graph::metadata::Returns; -use pgrx::pgrx_sql_entity_graph::metadata::ReturnsError; -use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping; -use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable; -use std::alloc::Layout; -use std::ops::Deref; -use std::ptr::NonNull; - -pub const HEADER_MAGIC: u16 = 0; - -#[repr(C, align(8))] -pub struct Vecf32Header { - varlena: u32, - dims: u16, - magic: u16, - phantom: [f32; 0], -} - -impl Vecf32Header { - fn varlena(size: usize) -> u32 { - (size << 2) as u32 - } - fn layout(len: usize) -> Layout { - u16::try_from(len).expect("Vector is too large."); - let layout_alpha = Layout::new::(); - let layout_beta = Layout::array::(len).unwrap(); - let layout = layout_alpha.extend(layout_beta).unwrap().0; - layout.pad_to_align() - } - #[allow(dead_code)] - pub fn dims(&self) -> u32 { - self.dims as u32 - } - pub fn slice(&self) -> &[f32] { - unsafe { std::slice::from_raw_parts(self.phantom.as_ptr(), self.dims as usize) } - } - pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { - unsafe { VectBorrowed::new_unchecked(self.slice()) } - } -} - -impl Deref for Vecf32Header { - type Target = [f32]; - - fn deref(&self) -> &Self::Target { - self.slice() - } -} - -pub enum Vecf32Input<'a> { - Owned(Vecf32Output), - Borrowed(&'a Vecf32Header), -} - -impl<'a> Vecf32Input<'a> { - unsafe fn new(p: NonNull) -> Self { - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - Vecf32Input::Owned(Vecf32Output(q)) - } else { - unsafe { Vecf32Input::Borrowed(p.as_ref()) } - } - } -} - -impl Deref for Vecf32Input<'_> { - type Target = Vecf32Header; - - fn deref(&self) -> &Self::Target { - match self { - Vecf32Input::Owned(x) => x, - Vecf32Input::Borrowed(x) => x, - } - } -} - -pub struct Vecf32Output(NonNull); - -impl Vecf32Output { - pub fn new(vector: VectBorrowed<'_, f32>) -> Vecf32Output { - unsafe { - let slice = vector.slice(); - let layout = Vecf32Header::layout(slice.len()); - let dims = vector.dims(); - let internal_dims = dims as u16; - let ptr = pgrx::pg_sys::palloc(layout.size()) as *mut Vecf32Header; - ptr.cast::().add(layout.size() - 8).write_bytes(0, 8); - std::ptr::addr_of_mut!((*ptr).varlena).write(Vecf32Header::varlena(layout.size())); - std::ptr::addr_of_mut!((*ptr).magic).write(HEADER_MAGIC); - std::ptr::addr_of_mut!((*ptr).dims).write(internal_dims); - std::ptr::copy_nonoverlapping(slice.as_ptr(), (*ptr).phantom.as_mut_ptr(), slice.len()); - Vecf32Output(NonNull::new(ptr).unwrap()) - } - } - pub fn into_raw(self) -> *mut Vecf32Header { - let result = self.0.as_ptr(); - std::mem::forget(self); - result - } -} - -impl Deref for Vecf32Output { - type Target = Vecf32Header; - - fn deref(&self) -> &Self::Target { - unsafe { self.0.as_ref() } - } -} - -impl Drop for Vecf32Output { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::pfree(self.0.as_ptr() as _); - } - } -} - -impl<'a> FromDatum for Vecf32Input<'a> { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let ptr = NonNull::new(datum.cast_mut_ptr::()).unwrap(); - unsafe { Some(Vecf32Input::new(ptr)) } - } - } -} - -impl IntoDatum for Vecf32Output { - fn into_datum(self) -> Option { - Some(Datum::from(self.into_raw() as *mut ())) - } - - fn type_oid() -> Oid { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"vectors").unwrap(); - let namespace = namespace.get().expect("pgvecto.rs is not installed."); - let t = pgrx::pg_catalog::PgType::search_typenamensp(c"vector", namespace.oid()).unwrap(); - let t = t.get().expect("pgvecto.rs is not installed."); - t.oid() - } -} - -impl FromDatum for Vecf32Output { - unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { - if is_null { - None - } else { - let p = NonNull::new(datum.cast_mut_ptr::())?; - let q = - unsafe { NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast())? }; - if p != q { - Some(Vecf32Output(q)) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Some(Vecf32Output::new(vector)) - } - } - } -} - -unsafe impl pgrx::datum::UnboxDatum for Vecf32Output { - type As<'src> = Vecf32Output; - #[inline] - unsafe fn unbox<'src>(d: pgrx::datum::Datum<'src>) -> Self::As<'src> - where - Self: 'src, - { - let p = NonNull::new(d.sans_lifetime().cast_mut_ptr::()).unwrap(); - let q = unsafe { - NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.cast().as_ptr()).cast()).unwrap() - }; - if p != q { - Vecf32Output(q) - } else { - let header = p.as_ptr(); - let vector = unsafe { (*header).as_borrowed() }; - Vecf32Output::new(vector) - } - } -} - -unsafe impl SqlTranslatable for Vecf32Input<'_> { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("vector"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("vector")))) - } -} - -unsafe impl SqlTranslatable for Vecf32Output { - fn argument_sql() -> Result { - Ok(SqlMapping::As(String::from("vector"))) - } - fn return_sql() -> Result { - Ok(Returns::One(SqlMapping::As(String::from("vector")))) - } -} - -unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for Vecf32Input<'fcx> { - unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { - unsafe { arg.unbox_arg_using_from_datum().unwrap() } - } -} - -unsafe impl pgrx::callconv::BoxRet for Vecf32Output { - unsafe fn box_into<'fcx>( - self, - fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, - ) -> pgrx::datum::Datum<'fcx> { - unsafe { fcinfo.return_raw_datum(Datum::from(self.into_raw() as *mut ())) } - } -} diff --git a/src/datatype/memory_vector.rs b/src/datatype/memory_vector.rs new file mode 100644 index 00000000..6cd6017c --- /dev/null +++ b/src/datatype/memory_vector.rs @@ -0,0 +1,258 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::datum::{FromDatum, IntoDatum}; +use pgrx::pg_sys::{Datum, Oid}; +use pgrx::pgrx_sql_entity_graph::metadata::*; +use std::marker::PhantomData; +use std::ptr::NonNull; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[repr(C)] +struct VectorHeader { + varlena: u32, + dim: u16, + unused: u16, + elements: [f32; 0], +} + +impl VectorHeader { + fn size_of(len: usize) -> usize { + if len > 65535 { + panic!("vector is too large"); + } + size_of::() + size_of::() * len + } + unsafe fn as_borrowed<'a>(this: NonNull) -> VectBorrowed<'a, f32> { + unsafe { + let this = this.as_ptr(); + VectBorrowed::new(std::slice::from_raw_parts( + (&raw const (*this).elements).cast(), + (&raw const (*this).dim).read() as usize, + )) + } + } +} + +pub struct VectorInput<'a>(NonNull, PhantomData<&'a ()>, bool); + +impl VectorInput<'_> { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + VectorInput(q, PhantomData, p != q) + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { + unsafe { VectorHeader::as_borrowed(self.0) } + } +} + +impl Drop for VectorInput<'_> { + fn drop(&mut self) { + if self.2 { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } + } +} + +pub struct VectorOutput(NonNull); + +impl VectorOutput { + unsafe fn from_ptr(p: NonNull) -> Self { + let q = unsafe { + NonNull::new(pgrx::pg_sys::pg_detoast_datum_copy(p.as_ptr().cast()).cast()).unwrap() + }; + unsafe { + let varlena = q.cast::().read(); + #[cfg(target_endian = "big")] + let size = varlena as usize; + #[cfg(target_endian = "little")] + let size = varlena as usize >> 2; + let dim = q.byte_add(4).cast::().read(); + assert_eq!(VectorHeader::size_of(dim as _), size); + let unused = q.byte_add(6).cast::().read(); + assert_eq!(unused, 0); + } + Self(q) + } + pub fn new(vector: VectBorrowed<'_, f32>) -> Self { + unsafe { + let slice = vector.slice(); + let size = VectorHeader::size_of(slice.len()); + + let ptr = pgrx::pg_sys::palloc0(size) as *mut VectorHeader; + // SET_VARSIZE_4B + #[cfg(target_endian = "big")] + (&raw mut (*ptr).varlena).write((size as u32) & 0x3FFFFFFF); + #[cfg(target_endian = "little")] + (&raw mut (*ptr).varlena).write((size << 2) as u32); + (&raw mut (*ptr).dim).write(vector.dim() as _); + (&raw mut (*ptr).unused).write(0); + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + (&raw mut (*ptr).elements).cast(), + slice.len(), + ); + Self(NonNull::new(ptr).unwrap()) + } + } + pub fn as_borrowed(&self) -> VectBorrowed<'_, f32> { + unsafe { VectorHeader::as_borrowed(self.0) } + } + fn into_raw(self) -> *mut VectorHeader { + let result = self.0.as_ptr(); + std::mem::forget(self); + result + } +} + +impl Drop for VectorOutput { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::pfree(self.0.as_ptr().cast()); + } + } +} + +// FromDatum + +impl FromDatum for VectorInput<'_> { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +impl FromDatum for VectorOutput { + unsafe fn from_polymorphic_datum(datum: Datum, is_null: bool, _typoid: Oid) -> Option { + if is_null { + None + } else { + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Some(Self::from_ptr(ptr)) } + } + } +} + +// IntoDatum + +impl IntoDatum for VectorOutput { + fn into_datum(self) -> Option { + Some(Datum::from(self.into_raw())) + } + + fn type_oid() -> Oid { + Oid::INVALID + } + + fn is_compatible_with(_: Oid) -> bool { + true + } +} + +// UnboxDatum + +unsafe impl<'a> pgrx::datum::UnboxDatum for VectorInput<'a> { + type As<'src> + = VectorInput<'src> + where + 'a: 'src; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +unsafe impl pgrx::datum::UnboxDatum for VectorOutput { + type As<'src> = VectorOutput; + #[inline] + unsafe fn unbox<'src>(datum: pgrx::datum::Datum<'src>) -> Self::As<'src> + where + Self: 'src, + { + let datum = datum.sans_lifetime(); + let ptr = NonNull::new(datum.cast_mut_ptr()).unwrap(); + unsafe { Self::from_ptr(ptr) } + } +} + +// SqlTranslatable + +unsafe impl SqlTranslatable for VectorInput<'_> { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("vector"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("vector")))) + } +} + +unsafe impl SqlTranslatable for VectorOutput { + fn argument_sql() -> Result { + Ok(SqlMapping::As(String::from("vector"))) + } + fn return_sql() -> Result { + Ok(Returns::One(SqlMapping::As(String::from("vector")))) + } +} + +// ArgAbi + +unsafe impl<'fcx> pgrx::callconv::ArgAbi<'fcx> for VectorInput<'fcx> { + unsafe fn unbox_arg_unchecked(arg: pgrx::callconv::Arg<'_, 'fcx>) -> Self { + let index = arg.index(); + unsafe { + arg.unbox_arg_using_from_datum() + .unwrap_or_else(|| panic!("argument {index} must not be null")) + } + } +} + +// BoxAbi + +unsafe impl pgrx::callconv::BoxRet for VectorOutput { + unsafe fn box_into<'fcx>( + self, + fcinfo: &mut pgrx::callconv::FcInfo<'fcx>, + ) -> pgrx::datum::Datum<'fcx> { + match self.into_datum() { + Some(datum) => unsafe { fcinfo.return_raw_datum(datum) }, + None => fcinfo.return_null(), + } + } +} diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 527ad4aa..ac550103 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -1,2 +1,45 @@ -pub mod memory_vecf32; +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod binary_rabitq4; +mod binary_rabitq8; +mod functions_rabitq4; +mod functions_rabitq8; +pub mod memory_halfvec; +pub mod memory_rabitq4; +pub mod memory_rabitq8; +pub mod memory_vector; +mod operators_halfvec; +mod operators_rabitq4; +mod operators_rabitq8; +mod operators_vector; +mod text_rabitq4; +mod text_rabitq8; pub mod typmod; +mod typmod_rabitq4; +mod typmod_rabitq8; + +pub(crate) const MAX_MAXSIM_VECTORS: usize = u16::MAX as usize + 1; + +pub(crate) fn validate_maxsim_array_len(len: usize) { + if len == 0 { + pgrx::error!("MaxSim arrays must contain at least one vector"); + } + if len > MAX_MAXSIM_VECTORS { + pgrx::error!( + "MaxSim arrays cannot contain more than {} vectors", + MAX_MAXSIM_VECTORS + ); + } +} diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs new file mode 100644 index 00000000..f20ebc8d --- /dev/null +++ b/src/datatype/operators_halfvec.rs @@ -0,0 +1,117 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use pgrx::datum::Array; +use std::num::NonZero; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_l2_in( + lhs: HalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_ip_in( + lhs: HalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_sphere_cosine_in( + lhs: HalfvecInput<'_>, + rhs: pgrx::composite_type!("sphere_halfvec"), +) -> bool { + let center: HalfvecOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_cos(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_halfvec_operator_maxsim( + lhs: Array<'_, HalfvecInput<'_>>, + rhs: Array<'_, HalfvecInput<'_>>, +) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs new file mode 100644 index 00000000..d1980fe9 --- /dev/null +++ b/src/datatype/operators_rabitq4.rs @@ -0,0 +1,147 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::datum::Array; +use std::num::NonZero; +use vector::VectorBorrowed; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_l2(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_ip(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_operator_cosine(lhs: Rabitq4Input<'_>, rhs: Rabitq4Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq4Borrowed::operator_cos(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_l2_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_ip_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_sphere_cosine_in( + lhs: Rabitq4Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq4"), +) -> bool { + let center: Rabitq4Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq4Borrowed::operator_cos(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq4_operator_maxsim( + lhs: Array<'_, Rabitq4Input<'_>>, + rhs: Array<'_, Rabitq4Input<'_>>, +) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + d = d.min(Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs new file mode 100644 index 00000000..7cac9aac --- /dev/null +++ b/src/datatype/operators_rabitq8.rs @@ -0,0 +1,147 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use pgrx::datum::Array; +use std::num::NonZero; +use vector::VectorBorrowed; +use vector::rabitq8::Rabitq8Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_operator_l2(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq8Borrowed::operator_l2s(lhs, rhs).to_f32().sqrt() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_operator_ip(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_operator_cosine(lhs: Rabitq8Input<'_>, rhs: Rabitq8Input<'_>) -> f32 { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + Rabitq8Borrowed::operator_cos(lhs, rhs).to_f32() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_sphere_l2_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), +) -> bool { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq8Borrowed::operator_l2s(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_sphere_ip_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), +) -> bool { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq8Borrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_sphere_cosine_in( + lhs: Rabitq8Input<'_>, + rhs: pgrx::composite_type!("sphere_rabitq8"), +) -> bool { + let center: Rabitq8Output = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = Rabitq8Borrowed::operator_cos(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(sql = "")] +fn _vchord_rabitq8_operator_maxsim( + lhs: Array<'_, Rabitq8Input<'_>>, + rhs: Array<'_, Rabitq8Input<'_>>, +) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + d = d.min(Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs new file mode 100644 index 00000000..c125a18b --- /dev/null +++ b/src/datatype/operators_vector.rs @@ -0,0 +1,117 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use pgrx::datum::Array; +use std::num::NonZero; +use vector::VectorBorrowed; +use vector::vect::VectBorrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_vector_sphere_l2_in( + lhs: VectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_l2s(lhs, center).to_f32().sqrt(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_vector_sphere_ip_in( + lhs: VectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_dot(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_vector_sphere_cosine_in( + lhs: VectorInput<'_>, + rhs: pgrx::composite_type!("sphere_vector"), +) -> bool { + let center: VectorOutput = match rhs.get_by_index(NonZero::new(1).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty center at sphere"), + Err(_) => unreachable!(), + }; + let radius: f32 = match rhs.get_by_index(NonZero::new(2).unwrap()) { + Ok(Some(s)) => s, + Ok(None) => pgrx::error!("Bad input: empty radius at sphere"), + Err(_) => unreachable!(), + }; + let lhs = lhs.as_borrowed(); + let center = center.as_borrowed(); + if lhs.dim() != center.dim() { + pgrx::error!("dimension is not matched"); + } + let d = VectBorrowed::operator_cos(lhs, center).to_f32(); + d < radius +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_vector_operator_maxsim( + lhs: Array<'_, VectorInput<'_>>, + rhs: Array<'_, VectorInput<'_>>, +) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } + let mut maxsim = 0.0f32; + for rhs in rhs.iter().flatten() { + let mut d = f32::INFINITY; + for lhs in lhs.iter().flatten() { + let lhs = lhs.as_borrowed(); + let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } + d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); + } + maxsim += d; + } + maxsim +} diff --git a/src/datatype/text_rabitq4.rs b/src/datatype/text_rabitq4.rs new file mode 100644 index 00000000..a239eb1c --- /dev/null +++ b/src/datatype/text_rabitq4.rs @@ -0,0 +1,160 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use pgrx::pg_sys::Oid; +use std::ffi::{CStr, CString}; +use vector::rabitq4::Rabitq4Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_in(input: &CStr, oid: Oid, typmod: i32) -> Rabitq4Output { + let _ = (oid, typmod); + let mut input = input.to_bytes().iter(); + let mut p0 = Vec::::new(); + let mut p1 = Vec::::new(); + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'(' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p0.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b')') => { + p0.push(s.parse().expect("failed to parse number")); + break; + } + (None, b')') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'[' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p1.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b']') => { + p1.push(s.parse().expect("failed to parse number")); + break; + } + (None, b']') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + if p0.len() != 4 { + pgrx::error!("incorrect vector"); + } + if p1.is_empty() { + pgrx::error!("vector must have at least 1 dimension"); + } + let sum_of_x2 = p0[0]; + let norm_of_lattice = p0[1]; + let sum_of_code = p0[2]; + let sum_of_abs_x = p0[3]; + let unpacked_code = p1; + let packed_code = rabitq::halfbyte::pack_code(&unpacked_code); + if let Some(x) = Rabitq4Borrowed::new_checked( + unpacked_code.len() as _, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq4Output::new(x) + } else { + pgrx::error!("incorrect vector"); + } +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_out(vector: Rabitq4Input<'_>) -> CString { + let vector = vector.as_borrowed(); + let mut buffer = String::new(); + buffer.push('('); + buffer.push_str(format!("{}", vector.sum_of_x2()).as_str()); + buffer.push_str(format!(", {}", vector.norm_of_lattice()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_code()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_abs_x()).as_str()); + buffer.push(')'); + buffer.push('['); + let mut unpacked_code = vector.unpacked_code(); + if let Some(x) = unpacked_code.next() { + buffer.push_str(format!("{x}").as_str()); + } + for x in unpacked_code { + buffer.push_str(format!(", {x}").as_str()); + } + buffer.push(']'); + CString::new(buffer).unwrap() +} diff --git a/src/datatype/text_rabitq8.rs b/src/datatype/text_rabitq8.rs new file mode 100644 index 00000000..f21e6586 --- /dev/null +++ b/src/datatype/text_rabitq8.rs @@ -0,0 +1,160 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use pgrx::pg_sys::Oid; +use std::ffi::{CStr, CString}; +use vector::rabitq8::Rabitq8Borrowed; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_in(input: &CStr, oid: Oid, typmod: i32) -> Rabitq8Output { + let _ = (oid, typmod); + let mut input = input.to_bytes().iter(); + let mut p0 = Vec::::new(); + let mut p1 = Vec::::new(); + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'(' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p0.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b')') => { + p0.push(s.parse().expect("failed to parse number")); + break; + } + (None, b')') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + { + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + match c { + b' ' => (), + b'[' => break, + _ => pgrx::error!("incorrect vector"), + } + } + } + { + let mut s = Option::::None; + loop { + let Some(c) = input.next().copied() else { + pgrx::error!("incorrect vector") + }; + s = match (s, c) { + (s, b' ') => s, + (None, c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + Some(String::from(c as char)) + } + (Some(s), c @ (b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-')) => { + let mut x = s; + x.push(c as char); + Some(x) + } + (Some(s), b',') => { + p1.push(s.parse().expect("failed to parse number")); + None + } + (None, b',') => { + pgrx::error!("incorrect vector") + } + (Some(s), b']') => { + p1.push(s.parse().expect("failed to parse number")); + break; + } + (None, b']') => break, + _ => pgrx::error!("incorrect vector"), + }; + } + } + if p0.len() != 4 { + pgrx::error!("incorrect vector"); + } + if p1.is_empty() { + pgrx::error!("vector must have at least 1 dimension"); + } + let sum_of_x2 = p0[0]; + let norm_of_lattice = p0[1]; + let sum_of_code = p0[2]; + let sum_of_abs_x = p0[3]; + let unpacked_code = p1; + let packed_code = rabitq::byte::pack_code(&unpacked_code); + if let Some(x) = Rabitq8Borrowed::new_checked( + unpacked_code.len() as _, + sum_of_x2, + norm_of_lattice, + sum_of_code, + sum_of_abs_x, + &packed_code, + ) { + Rabitq8Output::new(x) + } else { + pgrx::error!("incorrect vector"); + } +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_out(vector: Rabitq8Input<'_>) -> CString { + let vector = vector.as_borrowed(); + let mut buffer = String::new(); + buffer.push('('); + buffer.push_str(format!("{}", vector.sum_of_x2()).as_str()); + buffer.push_str(format!(", {}", vector.norm_of_lattice()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_code()).as_str()); + buffer.push_str(format!(", {}", vector.sum_of_abs_x()).as_str()); + buffer.push(')'); + buffer.push('['); + let mut unpacked_code = vector.unpacked_code(); + if let Some(x) = unpacked_code.next() { + buffer.push_str(format!("{x}").as_str()); + } + for x in unpacked_code { + buffer.push_str(format!(", {x}").as_str()); + } + buffer.push(']'); + CString::new(buffer).unwrap() +} diff --git a/src/datatype/typmod.rs b/src/datatype/typmod.rs index d67bb0be..6ab9b44c 100644 --- a/src/datatype/typmod.rs +++ b/src/datatype/typmod.rs @@ -1,44 +1,41 @@ -use serde::{Deserialize, Serialize}; -use std::num::NonZeroU32; +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +use std::num::NonZero; + +#[derive(Debug, Clone, Copy)] pub enum Typmod { Any, - Dims(NonZeroU32), + Dims(NonZero), } impl Typmod { - pub fn parse_from_i32(x: i32) -> Option { + pub fn new(x: i32) -> Option { use Typmod::*; if x == -1 { Some(Any) } else if x >= 1 { - Some(Dims(NonZeroU32::new(x as u32).unwrap())) + Some(Dims(NonZero::new(x as u32).unwrap())) } else { None } } - #[allow(dead_code)] - pub fn into_option_string(self) -> Option { - use Typmod::*; - match self { - Any => None, - Dims(x) => Some(x.get().to_string()), - } - } - #[allow(dead_code)] - pub fn into_i32(self) -> i32 { - use Typmod::*; - match self { - Any => -1, - Dims(x) => x.get() as i32, - } - } - pub fn dims(self) -> Option { + pub fn dim(self) -> Option> { use Typmod::*; match self { Any => None, - Dims(dims) => Some(dims), + Dims(dim) => Some(dim), } } } diff --git a/src/datatype/typmod_rabitq4.rs b/src/datatype/typmod_rabitq4.rs new file mode 100644 index 00000000..32ac0f93 --- /dev/null +++ b/src/datatype/typmod_rabitq4.rs @@ -0,0 +1,37 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::ffi::CStr; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq4_typmod_in(list: pgrx::datum::Array<&CStr>) -> i32 { + if list.is_empty() { + -1 + } else if list.len() == 1 { + let s = list.get(0).unwrap().unwrap().to_str().unwrap(); + if let Ok(d) = s.parse::() { + if d < 1 { + pgrx::error!("dimensions for type rabitq4 must be at least 1"); + } + if d > 65535 { + pgrx::error!("dimensions for type rabitq4 cannot exceed 65535"); + } + d + } else { + pgrx::error!("invalid type modifier") + } + } else { + pgrx::error!("invalid type modifier") + } +} diff --git a/src/datatype/typmod_rabitq8.rs b/src/datatype/typmod_rabitq8.rs new file mode 100644 index 00000000..881bb35d --- /dev/null +++ b/src/datatype/typmod_rabitq8.rs @@ -0,0 +1,37 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use std::ffi::CStr; + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchord_rabitq8_typmod_in(list: pgrx::datum::Array<&CStr>) -> i32 { + if list.is_empty() { + -1 + } else if list.len() == 1 { + let s = list.get(0).unwrap().unwrap().to_str().unwrap(); + if let Ok(d) = s.parse::() { + if d < 1 { + pgrx::error!("dimensions for type rabitq8 must be at least 1"); + } + if d > 65535 { + pgrx::error!("dimensions for type rabitq8 cannot exceed 65535"); + } + d + } else { + pgrx::error!("invalid type modifier") + } + } else { + pgrx::error!("invalid type modifier") + } +} diff --git a/src/gucs/executing.rs b/src/gucs/executing.rs deleted file mode 100644 index 421d5a40..00000000 --- a/src/gucs/executing.rs +++ /dev/null @@ -1,20 +0,0 @@ -use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting}; - -static NPROBE: GucSetting = GucSetting::::new(10); - -pub unsafe fn init() { - GucRegistry::define_int_guc( - "rabbithole.nprobe", - "`nprobe` argument of rabbithole.", - "`nprobe` argument of rabbithole.", - &NPROBE, - 1, - u16::MAX as _, - GucContext::Userset, - GucFlags::default(), - ); -} - -pub fn nprobe() -> u32 { - NPROBE.get() as u32 -} diff --git a/src/gucs/mod.rs b/src/gucs/mod.rs deleted file mode 100644 index cc0145ef..00000000 --- a/src/gucs/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod executing; - -pub unsafe fn init() { - unsafe { - executing::init(); - } -} diff --git a/src/index/am.rs b/src/index/am.rs deleted file mode 100644 index 329a257f..00000000 --- a/src/index/am.rs +++ /dev/null @@ -1,384 +0,0 @@ -use crate::algorithm; -use crate::algorithm::build::HeapRelation; -use crate::index::am_options::{Opfamily, Reloption}; -use crate::index::am_scan::Scanner; -use crate::index::utils::{ctid_to_pointer, pointer_to_ctid}; -use crate::index::{am_options, am_scan}; -use crate::postgres::Relation; -use crate::utils::cells::PgCell; -use base::search::Pointer; -use pgrx::datum::Internal; -use pgrx::pg_sys::Datum; - -static RELOPT_KIND_RABBITHOLE: PgCell = unsafe { PgCell::new(0) }; - -pub unsafe fn init() { - unsafe { - RELOPT_KIND_RABBITHOLE.set(pgrx::pg_sys::add_reloption_kind()); - pgrx::pg_sys::add_string_reloption( - RELOPT_KIND_RABBITHOLE.get(), - c"options".as_ptr(), - c"Vector index options, represented as a TOML string.".as_ptr(), - c"".as_ptr(), - None, - pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, - ); - } -} - -#[pgrx::pg_extern(sql = "\ -CREATE FUNCTION _rabbithole_amhandler(internal) RETURNS index_am_handler -IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';")] -fn _rabbithole_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { - type T = pgrx::pg_sys::IndexAmRoutine; - unsafe { - let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; - index_am_routine.write(AM_HANDLER); - Internal::from(Some(Datum::from(index_am_routine))) - } -} - -const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = { - let mut am_routine = - unsafe { std::mem::MaybeUninit::::zeroed().assume_init() }; - - am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; - - am_routine.amcanorderbyop = true; - - // Index access methods that set `amoptionalkey` to `false` - // must index all tuples, even if the first column is `NULL`. - // However, PostgreSQL does not generate a path if there is no - // index clauses, even if there is a `ORDER BY` clause. - // So we have to set it to `true` and set costs of every path - // for vector index scans without `ORDER BY` clauses a large number - // and throw errors if someone really wants such a path. - am_routine.amoptionalkey = true; - - am_routine.amvalidate = Some(amvalidate); - am_routine.amoptions = Some(amoptions); - am_routine.amcostestimate = Some(amcostestimate); - - am_routine.ambuild = Some(ambuild); - am_routine.ambuildempty = Some(ambuildempty); - am_routine.aminsert = Some(aminsert); - am_routine.ambulkdelete = Some(ambulkdelete); - am_routine.amvacuumcleanup = Some(amvacuumcleanup); - - am_routine.ambeginscan = Some(ambeginscan); - am_routine.amrescan = Some(amrescan); - am_routine.amgettuple = Some(amgettuple); - am_routine.amendscan = Some(amendscan); - - am_routine -}; - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvalidate(opclass_oid: pgrx::pg_sys::Oid) -> bool { - if am_options::convert_opclass_to_vd(opclass_oid).is_some() { - pgrx::info!("Vector indexes can only be built on built-in operator classes."); - true - } else { - false - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amoptions(reloptions: Datum, validate: bool) -> *mut pgrx::pg_sys::bytea { - let rdopts = unsafe { - pgrx::pg_sys::build_reloptions( - reloptions, - validate, - RELOPT_KIND_RABBITHOLE.get(), - size_of::(), - Reloption::TAB.as_ptr(), - Reloption::TAB.len() as _, - ) - }; - rdopts as *mut pgrx::pg_sys::bytea -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amcostestimate( - _root: *mut pgrx::pg_sys::PlannerInfo, - path: *mut pgrx::pg_sys::IndexPath, - _loop_count: f64, - index_startup_cost: *mut pgrx::pg_sys::Cost, - index_total_cost: *mut pgrx::pg_sys::Cost, - index_selectivity: *mut pgrx::pg_sys::Selectivity, - index_correlation: *mut f64, - index_pages: *mut f64, -) { - unsafe { - if (*path).indexorderbys.is_null() && (*path).indexclauses.is_null() { - *index_startup_cost = f64::MAX; - *index_total_cost = f64::MAX; - *index_selectivity = 0.0; - *index_correlation = 0.0; - *index_pages = 0.0; - return; - } - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 1.0; - *index_pages = 0.0; - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuild( - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, -) -> *mut pgrx::pg_sys::IndexBuildResult { - pub struct Heap { - heap: pgrx::pg_sys::Relation, - index: pgrx::pg_sys::Relation, - index_info: *mut pgrx::pg_sys::IndexInfo, - opfamily: Opfamily, - } - impl HeapRelation for Heap { - fn traverse(&self, callback: F) - where - F: FnMut((Pointer, Vec)), - { - pub struct State<'a, F> { - pub this: &'a Heap, - pub callback: F, - } - #[pgrx::pg_guard] - unsafe extern "C" fn call( - _index: pgrx::pg_sys::Relation, - ctid: pgrx::pg_sys::ItemPointer, - values: *mut Datum, - is_null: *mut bool, - _tuple_is_alive: bool, - state: *mut core::ffi::c_void, - ) where - F: FnMut((Pointer, Vec)), - { - pgrx::check_for_interrupts!(); - use base::vector::OwnedVector; - let state = unsafe { &mut *state.cast::>() }; - let vector = unsafe { - state - .this - .opfamily - .datum_to_vector(*values.add(0), *is_null.add(0)) - }; - let pointer = unsafe { ctid_to_pointer(ctid.read()) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; - (state.callback)((pointer, vector.into_vec())); - } - } - let table_am = unsafe { &*(*self.heap).rd_tableam }; - let mut state = State { - this: self, - callback, - }; - unsafe { - table_am.index_build_range_scan.unwrap()( - self.heap, - self.index, - self.index_info, - true, - false, - true, - 0, - pgrx::pg_sys::InvalidBlockNumber, - Some(call::), - (&mut state) as *mut State as *mut _, - std::ptr::null_mut(), - ); - } - } - } - let (vector_options, rabbithole_options) = unsafe { am_options::options(index) }; - let heap_relation = Heap { - heap, - index, - index_info, - opfamily: unsafe { am_options::opfamily(index) }, - }; - let index_relation = unsafe { Relation::new(index) }; - algorithm::build::build( - vector_options, - rabbithole_options, - heap_relation, - index_relation, - ); - unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambuildempty(_index: pgrx::pg_sys::Relation) { - pgrx::error!("Unlogged indexes are not supported."); -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn aminsert( - index: pgrx::pg_sys::Relation, - values: *mut Datum, - is_null: *mut bool, - heap_tid: pgrx::pg_sys::ItemPointer, - _heap: pgrx::pg_sys::Relation, - _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, - _index_unchanged: bool, - _index_info: *mut pgrx::pg_sys::IndexInfo, -) -> bool { - use base::vector::OwnedVector; - let opfamily = unsafe { am_options::opfamily(index) }; - let vector = unsafe { opfamily.datum_to_vector(*values.add(0), *is_null.add(0)) }; - if let Some(vector) = vector { - let vector = match vector { - OwnedVector::Vecf32(x) => x, - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }; - let pointer = ctid_to_pointer(unsafe { heap_tid.read() }); - algorithm::insert::insert(unsafe { Relation::new(index) }, pointer, vector.into_vec()); - } - false -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambeginscan( - index: pgrx::pg_sys::Relation, - n_keys: std::os::raw::c_int, - n_orderbys: std::os::raw::c_int, -) -> pgrx::pg_sys::IndexScanDesc { - use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; - - let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index, n_keys, n_orderbys) }; - unsafe { - let scanner = am_scan::scan_make(None, None, false); - (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); - } - scan -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amrescan( - scan: pgrx::pg_sys::IndexScanDesc, - keys: pgrx::pg_sys::ScanKey, - _n_keys: std::os::raw::c_int, - orderbys: pgrx::pg_sys::ScanKey, - _n_orderbys: std::os::raw::c_int, -) { - unsafe { - if !keys.is_null() && (*scan).numberOfKeys > 0 { - std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); - } - if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { - std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); - } - let opfamily = am_options::opfamily((*scan).indexRelation); - let (orderbys, spheres) = { - let mut orderbys = Vec::new(); - let mut spheres = Vec::new(); - if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { - pgrx::error!( - "vector search with no WHERE clause and no ORDER BY clause is not supported" - ); - } - for i in 0..(*scan).numberOfOrderBys { - let data = (*scan).orderByData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 1 => orderbys.push(opfamily.datum_to_vector(value, is_null)), - _ => unreachable!(), - } - } - for i in 0..(*scan).numberOfKeys { - let data = (*scan).keyData.add(i as usize); - let value = (*data).sk_argument; - let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; - match (*data).sk_strategy { - 2 => spheres.push(opfamily.datum_to_sphere(value, is_null)), - _ => unreachable!(), - } - } - (orderbys, spheres) - }; - let (vector, threshold, recheck) = am_scan::scan_build(orderbys, spheres, opfamily); - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(vector, threshold, recheck)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amgettuple( - scan: pgrx::pg_sys::IndexScanDesc, - direction: pgrx::pg_sys::ScanDirection::Type, -) -> bool { - if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { - pgrx::error!("vector search without a forward scan direction is not supported"); - } - // https://www.postgresql.org/docs/current/index-locking.html - // If heap entries referenced physical pointers are deleted before - // they are consumed by PostgreSQL, PostgreSQL will received wrong - // physical pointers: no rows or irreverent rows are referenced. - if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC - { - pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); - } - let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; - let relation = unsafe { Relation::new((*scan).indexRelation) }; - if let Some((pointer, recheck)) = am_scan::scan_next(scanner, relation) { - let ctid = pointer_to_ctid(pointer); - unsafe { - (*scan).xs_heaptid = ctid; - (*scan).xs_recheckorderby = false; - (*scan).xs_recheck = recheck; - } - true - } else { - false - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { - unsafe { - let scanner = (*scan).opaque.cast::().as_mut().unwrap_unchecked(); - let scanner = std::mem::replace(scanner, am_scan::scan_make(None, None, false)); - am_scan::scan_release(scanner); - } -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn ambulkdelete( - info: *mut pgrx::pg_sys::IndexVacuumInfo, - stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, - callback: pgrx::pg_sys::IndexBulkDeleteCallback, - callback_state: *mut std::os::raw::c_void, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - let mut stats = stats; - if stats.is_null() { - stats = unsafe { - pgrx::pg_sys::palloc0(size_of::()).cast() - }; - } - let callback = callback.unwrap(); - let callback = |p: Pointer| unsafe { callback(&mut pointer_to_ctid(p), callback_state) }; - algorithm::vacuum::vacuum(unsafe { Relation::new((*info).index) }, callback); - stats -} - -#[pgrx::pg_guard] -pub unsafe extern "C" fn amvacuumcleanup( - _info: *mut pgrx::pg_sys::IndexVacuumInfo, - _stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, -) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { - std::ptr::null_mut() -} diff --git a/src/index/am_options.rs b/src/index/am_options.rs deleted file mode 100644 index a3486a73..00000000 --- a/src/index/am_options.rs +++ /dev/null @@ -1,208 +0,0 @@ -use crate::datatype::memory_vecf32::Vecf32Input; -use crate::datatype::memory_vecf32::Vecf32Output; -use crate::datatype::typmod::Typmod; -use crate::types::RabbitholeIndexingOptions; -use base::distance::*; -use base::index::*; -use base::vector::*; -use pgrx::datum::FromDatum; -use pgrx::heap_tuple::PgHeapTuple; -use serde::Deserialize; -use std::ffi::CStr; -use std::num::NonZero; - -#[derive(Copy, Clone, Debug, Default)] -#[repr(C)] -pub struct Reloption { - vl_len_: i32, - pub options: i32, -} - -impl Reloption { - pub const TAB: &'static [pgrx::pg_sys::relopt_parse_elt] = &[pgrx::pg_sys::relopt_parse_elt { - optname: c"options".as_ptr(), - opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, - offset: std::mem::offset_of!(Reloption, options) as i32, - }]; - unsafe fn options(&self) -> &CStr { - unsafe { - let ptr = std::ptr::addr_of!(*self) - .cast::() - .offset(self.options as _); - CStr::from_ptr(ptr) - } - } -} - -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum PgDistanceKind { - L2, -} - -impl PgDistanceKind { - pub fn to_distance(self) -> DistanceKind { - match self { - PgDistanceKind::L2 => DistanceKind::L2, - } - } -} - -pub fn convert_opclass_to_vd( - opclass_oid: pgrx::pg_sys::Oid, -) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); - let namespace = namespace.get().expect("rabbithole is not installed."); - let opclass = pgrx::pg_catalog::PgOpclass::search_claoid(opclass_oid).unwrap(); - let opclass = opclass.get().expect("rabbithole is not installed."); - if opclass.opcnamespace() == namespace.oid() { - if let Ok(name) = opclass.opcname().to_str() { - if let Some(p) = convert_name_to_vd(name) { - return Some(p); - } - } - } - None -} - -pub fn convert_opfamily_to_vd( - opfamily_oid: pgrx::pg_sys::Oid, -) -> Option<(VectorKind, PgDistanceKind)> { - let namespace = pgrx::pg_catalog::PgNamespace::search_namespacename(c"rabbithole").unwrap(); - let namespace = namespace.get().expect("rabbithole is not installed."); - let opfamily = pgrx::pg_catalog::PgOpfamily::search_opfamilyoid(opfamily_oid).unwrap(); - let opfamily = opfamily.get().expect("rabbithole is not installed."); - if opfamily.opfnamespace() == namespace.oid() { - if let Ok(name) = opfamily.opfname().to_str() { - if let Some(p) = convert_name_to_vd(name) { - return Some(p); - } - } - } - None -} - -fn convert_name_to_vd(name: &str) -> Option<(VectorKind, PgDistanceKind)> { - match name.strip_suffix("_ops") { - Some("vector_l2") => Some((VectorKind::Vecf32, PgDistanceKind::L2)), - _ => None, - } -} - -unsafe fn convert_reloptions_to_options( - reloptions: *const pgrx::pg_sys::varlena, -) -> RabbitholeIndexingOptions { - #[derive(Debug, Clone, Deserialize, Default)] - #[serde(deny_unknown_fields)] - struct Parsed { - #[serde(flatten)] - rabitq: RabbitholeIndexingOptions, - } - let reloption = reloptions as *const Reloption; - if reloption.is_null() || unsafe { (*reloption).options == 0 } { - return Default::default(); - } - let s = unsafe { (*reloption).options() }.to_string_lossy(); - match toml::from_str::(&s) { - Ok(p) => p.rabitq, - Err(e) => pgrx::error!("failed to parse options: {}", e), - } -} - -pub unsafe fn options(index: pgrx::pg_sys::Relation) -> (VectorOptions, RabbitholeIndexingOptions) { - let opfamily = unsafe { (*index).rd_opfamily.read() }; - let att = unsafe { &mut *(*index).rd_att }; - let atts = unsafe { att.attrs.as_slice(att.natts as _) }; - if atts.is_empty() { - pgrx::error!("indexing on no columns is not supported"); - } - if atts.len() != 1 { - pgrx::error!("multicolumn index is not supported"); - } - // get dims - let typmod = Typmod::parse_from_i32(atts[0].type_mod()).unwrap(); - let dims = if let Some(dims) = typmod.dims() { - dims.get() - } else { - pgrx::error!( - "Dimensions type modifier of a vector column is needed for building the index." - ); - }; - // get v, d - let (v, pg_d) = convert_opfamily_to_vd(opfamily).unwrap(); - let vector = VectorOptions { - dims, - v, - d: pg_d.to_distance(), - }; - // get indexing, segment, optimizing - let rabitq = unsafe { convert_reloptions_to_options((*index).rd_options) }; - (vector, rabitq) -} - -#[derive(Debug, Clone, Copy)] -pub struct Opfamily { - vector: VectorKind, - pg_distance: PgDistanceKind, -} - -impl Opfamily { - pub unsafe fn datum_to_vector( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> Option { - if is_null || datum.is_null() { - return None; - } - let vector = match self.vector { - VectorKind::Vecf32 => { - let vector = unsafe { Vecf32Input::from_datum(datum, false).unwrap() }; - self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed())) - } - _ => unreachable!(), - }; - Some(vector) - } - pub unsafe fn datum_to_sphere( - self, - datum: pgrx::pg_sys::Datum, - is_null: bool, - ) -> (Option, Option) { - if is_null || datum.is_null() { - return (None, None); - } - let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; - let center = match self.vector { - VectorKind::Vecf32 => tuple - .get_by_index::(NonZero::new(1).unwrap()) - .unwrap() - .map(|vector| self.preprocess(BorrowedVector::Vecf32(vector.as_borrowed()))), - _ => unreachable!(), - }; - let radius = tuple.get_by_index::(NonZero::new(2).unwrap()).unwrap(); - (center, radius) - } - pub fn preprocess(self, vector: BorrowedVector<'_>) -> OwnedVector { - use BorrowedVector as B; - use OwnedVector as O; - match (vector, self.pg_distance) { - (B::Vecf32(x), _) => O::Vecf32(x.own()), - (B::Vecf16(x), _) => O::Vecf16(x.own()), - (B::SVecf32(x), _) => O::SVecf32(x.own()), - (B::BVector(x), _) => O::BVector(x.own()), - } - } - pub fn process(self, x: Distance) -> f32 { - f32::from(x) - } -} - -pub unsafe fn opfamily(index: pgrx::pg_sys::Relation) -> Opfamily { - let opfamily = unsafe { (*index).rd_opfamily.read() }; - let (vector, pg_distance) = convert_opfamily_to_vd(opfamily).unwrap(); - Opfamily { - vector, - pg_distance, - } -} diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs deleted file mode 100644 index 9792729a..00000000 --- a/src/index/am_scan.rs +++ /dev/null @@ -1,124 +0,0 @@ -use super::am_options::Opfamily; -use crate::algorithm::scan::scan; -use crate::gucs::executing::nprobe; -use crate::postgres::Relation; -use base::distance::Distance; -use base::search::*; -use base::vector::*; - -pub enum Scanner { - Initial { - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, - }, - Vbase { - vbase: Box>, - threshold: Option, - recheck: bool, - opfamily: Opfamily, - }, - Empty {}, -} - -pub fn scan_build( - orderbys: Vec>, - spheres: Vec<(Option, Option)>, - opfamily: Opfamily, -) -> (Option<(OwnedVector, Opfamily)>, Option, bool) { - let mut pair = None; - let mut threshold = None; - let mut recheck = false; - for orderby_vector in orderbys { - if pair.is_none() { - pair = orderby_vector; - } else if orderby_vector.is_some() && pair != orderby_vector { - pgrx::error!("vector search with multiple vectors is not supported"); - } - } - for (sphere_vector, sphere_threshold) in spheres { - if pair.is_none() { - pair = sphere_vector; - threshold = sphere_threshold; - } else if pair == sphere_vector { - if threshold.is_none() || sphere_threshold < threshold { - threshold = sphere_threshold; - } - } else { - recheck = true; - break; - } - } - (pair.map(|x| (x, opfamily)), threshold, recheck) -} - -pub fn scan_make( - vector: Option<(OwnedVector, Opfamily)>, - threshold: Option, - recheck: bool, -) -> Scanner { - Scanner::Initial { - vector, - threshold, - recheck, - } -} - -pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer, bool)> { - if let Scanner::Initial { - vector, - threshold, - recheck, - } = scanner - { - if let Some((vector, opfamily)) = vector.as_ref() { - let vbase = scan( - relation, - match vector { - OwnedVector::Vecf32(x) => x.slice().to_vec(), - OwnedVector::Vecf16(_) => unreachable!(), - OwnedVector::SVecf32(_) => unreachable!(), - OwnedVector::BVector(_) => unreachable!(), - }, - nprobe(), - ); - *scanner = Scanner::Vbase { - vbase: Box::new(vbase), - threshold: *threshold, - recheck: *recheck, - opfamily: *opfamily, - }; - } else { - *scanner = Scanner::Empty {}; - } - } - match scanner { - Scanner::Initial { .. } => unreachable!(), - Scanner::Vbase { - vbase, - threshold, - recheck, - opfamily, - } => match ( - vbase.next().map(|(d, p)| (opfamily.process(d), p)), - threshold, - ) { - (Some((_, ptr)), None) => Some((ptr, *recheck)), - (Some((distance, ptr)), Some(t)) if distance < *t => Some((ptr, *recheck)), - _ => { - let scanner = std::mem::replace(scanner, Scanner::Empty {}); - scan_release(scanner); - None - } - }, - Scanner::Empty {} => None, - } -} - -pub fn scan_release(scanner: Scanner) { - match scanner { - Scanner::Initial { .. } => {} - Scanner::Vbase { .. } => {} - Scanner::Empty {} => {} - } -} diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs new file mode 100644 index 00000000..fa836316 --- /dev/null +++ b/src/index/fetcher.rs @@ -0,0 +1,352 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::pg_sys::{BlockIdData, Datum, ItemPointerData}; +use std::cell::LazyCell; +use std::num::NonZero; +use std::ops::DerefMut; +use std::ptr::NonNull; + +pub trait FilterableTuple: Tuple { + fn filter(&mut self) -> bool; +} + +pub trait Tuple { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + + /// Read one user attribute from the already fetched heap slot. + /// + /// Callers that may expose the value outside PostgreSQL must call + /// [`FilterableTuple::filter`] first. Attribute numbers are PostgreSQL + /// one-based `attnum` values, not zero-based Rust indexes. + #[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] + fn attribute(&mut self, attnum: i16) -> Option; +} + +#[derive(Clone, Copy)] +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub struct TupleAttribute { + pub datum: Datum, + pub is_null: bool, +} + +pub trait Fetcher { + type Tuple<'a>: FilterableTuple + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option>; +} + +impl T> Fetcher for LazyCell { + type Tuple<'a> + = T::Tuple<'a> + where + Self: 'a; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { + self.deref_mut().fetch(key) + } +} + +pub struct HeapFetcher { + index_info: *mut pgrx::pg_sys::IndexInfo, + estate: *mut pgrx::pg_sys::EState, + econtext: *mut pgrx::pg_sys::ExprContext, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, + owns_heapfetch: bool, + slot: *mut pgrx::pg_sys::TupleTableSlot, + values: [Datum; 32], + is_nulls: [bool; 32], + hack: *mut pgrx::pg_sys::IndexScanState, +} + +impl HeapFetcher { + pub unsafe fn new( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, + hack: *mut pgrx::pg_sys::IndexScanState, + ) -> Self { + unsafe { + let index_info = pgrx::pg_sys::BuildIndexInfo(index_relation); + let estate = pgrx::pg_sys::CreateExecutorState(); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); + Self { + index_info, + estate, + econtext, + heap_relation, + snapshot, + heapfetch, + owns_heapfetch: false, + slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), + values: [Datum::null(); 32], + is_nulls: [true; 32], + hack, + } + } + } + + /// Create a heap fetch state that is not owned by an `IndexScanDesc`. + /// + /// This is used by the restricted external MaxSim executor to resolve the + /// root TIDs stored in the index through HOT chains before SQL-visible + /// descriptor projection. The table AM owns the rules for doing that; + /// looking up `ctid` directly does not follow a HOT chain. + pub unsafe fn new_standalone( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + unsafe { + let table_am = (*heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_fetch_begin = (*table_am) + .index_fetch_begin + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let heapfetch = pg_guard_ffi_boundary(|| index_fetch_begin(heap_relation)); + if heapfetch.is_null() { + panic!("heap access method returned a null index fetch state"); + } + let mut fetcher = Self::new( + index_relation, + heap_relation, + snapshot, + heapfetch, + std::ptr::null_mut(), + ); + fetcher.owns_heapfetch = true; + fetcher + } + } +} + +impl Drop for HeapFetcher { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // free common resources + pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); + pgrx::pg_sys::FreeExecutorState(self.estate); + if self.owns_heapfetch { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let table_am = (*self.heap_relation).rd_tableam; + let index_fetch_end = (*table_am) + .index_fetch_end + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| index_fetch_end(self.heapfetch)); + } + } + } +} + +impl Fetcher for HeapFetcher { + type Tuple<'a> = HeapTuple<'a>; + + fn fetch(&mut self, key: [u16; 3]) -> Option> { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let mut ctid = key_to_ctid(key); + let table_am = (*self.heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_fetch_tuple = (*table_am) + .index_fetch_tuple + .expect("unsupported heap access method"); + let found = 'a: { + let mut call_again = false; + let mut all_dead = false; + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let found = pg_guard_ffi_boundary(|| { + index_fetch_tuple( + self.heapfetch, + &mut ctid, + self.snapshot, + self.slot, + &mut call_again, + &mut all_dead, + ) + }); + if found { + break 'a true; + } + while call_again { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let found = pg_guard_ffi_boundary(|| { + index_fetch_tuple( + self.heapfetch, + &mut ctid, + self.snapshot, + self.slot, + &mut call_again, + &mut all_dead, + ) + }); + if found { + break 'a true; + } + } + false + }; + if found { + // The heap table AM rewrites the requested root TID to the + // snapshot-visible HOT-chain member. The slot itself keeps + // the rewritten index TID as well, but carrying it explicitly + // avoids depending on slot representation details. + Some(HeapTuple { + this: self, + current_ctid: ctid, + }) + } else { + None + } + } + } +} + +pub struct HeapTuple<'a> { + this: &'a mut HeapFetcher, + current_ctid: ItemPointerData, +} + +impl HeapTuple<'_> { + /// Return the physical TID of the tuple version materialized in the slot. + /// This may differ from the root TID supplied to `Fetcher::fetch` after a + /// HOT update. + pub fn ctid(&self) -> ItemPointerData { + self.current_ctid + } +} + +impl Tuple for HeapTuple<'_> { + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { + unsafe { + let this = &mut self.this; + (*this.econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); + pgrx::pg_sys::FormIndexDatum( + this.index_info, + this.slot, + this.estate, + this.values.as_mut_ptr(), + this.is_nulls.as_mut_ptr(), + ); + (&this.values, &this.is_nulls) + } + } + + fn attribute(&mut self, attnum: i16) -> Option { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let slot = self.this.slot; + let tuple_descriptor = (*slot).tts_tupleDescriptor; + if attnum <= 0 + || tuple_descriptor.is_null() + || i32::from(attnum) > (*tuple_descriptor).natts + { + return None; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + pgrx::pg_sys::slot_getsomeattrs_int(slot, i32::from(attnum)); + }); + let offset = usize::try_from(attnum - 1).ok()?; + Some(TupleAttribute { + datum: *(*slot).tts_values.add(offset), + is_null: *(*slot).tts_isnull.add(offset), + }) + } + } +} + +impl FilterableTuple for HeapTuple<'_> { + fn filter(&mut self) -> bool { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let this = &mut self.this; + if !this.hack.is_null() { + if let Some(qual) = NonNull::new((*this.hack).ss.ps.qual) { + use pgrx::datum::FromDatum; + use pgrx::memcxt::PgMemoryContexts; + assert!(qual.as_ref().flags & pgrx::pg_sys::EEO_FLAG_IS_QUAL as u8 != 0); + let evalfunc = qual.as_ref().evalfunc.expect("no evalfunc for qual"); + if !(*this.hack).ss.ps.ps_ExprContext.is_null() { + let econtext = (*this.hack).ss.ps.ps_ExprContext; + (*econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*econtext).ecxt_per_tuple_memory); + let result = PgMemoryContexts::For((*econtext).ecxt_per_tuple_memory) + .switch_to(|_| { + let mut is_null = true; + #[allow( + ffi_unwind_calls, + reason = "protected by pg_guard_ffi_boundary" + )] + let datum = pg_guard_ffi_boundary(|| { + evalfunc(qual.as_ptr(), econtext, &mut is_null) + }); + bool::from_datum(datum, is_null) + }); + if result != Some(true) { + return false; + } + } + } + } + true + } + } +} + +pub const fn ctid_to_key( + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + }: ItemPointerData, +) -> [u16; 3] { + [bi_hi, bi_lo, ip_posid] +} + +pub const fn key_to_ctid([bi_hi, bi_lo, ip_posid]: [u16; 3]) -> ItemPointerData { + ItemPointerData { + ip_blkid: BlockIdData { bi_hi, bi_lo }, + ip_posid, + } +} + +pub const fn pointer_to_kv(pointer: NonZero) -> ([u16; 3], u16) { + let value = pointer.get(); + let bi_hi = ((value >> 48) & 0xffff) as u16; + let bi_lo = ((value >> 32) & 0xffff) as u16; + let ip_posid = ((value >> 16) & 0xffff) as u16; + let extra = value as u16; + ([bi_hi, bi_lo, ip_posid], extra) +} + +pub const fn kv_to_pointer((key, value): ([u16; 3], u16)) -> NonZero { + let x = (key[0] as u64) << 48 | (key[1] as u64) << 32 | (key[2] as u64) << 16 | value as u64; + NonZero::new(x).expect("invalid key") +} diff --git a/src/index/functions.rs b/src/index/functions.rs new file mode 100644 index 00000000..d21d08b3 --- /dev/null +++ b/src/index/functions.rs @@ -0,0 +1,110 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::storage::PostgresRelation; +use crate::recorder::dump; +use pgrx::iter::SetOfIterator; +use pgrx::pg_sys::Oid; +use pgrx_catalog::{PgAm, PgClass, PgClassRelkind}; + +#[pgrx::pg_extern(sql = "")] +fn _vchordg_prewarm(indexrelid: Oid) -> String { + let pg_am = PgAm::search_amname(c"vchordg").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("the relation does not exist"); + }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } + if pg_class.relam() != pg_am.oid() { + pgrx::error!("the index {:?} is not a vchordg index", pg_class.relname()); + } + let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let opfamily = unsafe { crate::index::vchordg::opclass::opfamily(relation.raw()) }; + let index = unsafe { PostgresRelation::new(relation.raw()) }; + crate::index::vchordg::dispatch::prewarm(opfamily, &index) +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_prewarm(indexrelid: Oid, height: i32) -> String { + let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("the relation does not exist"); + }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } + if pg_class.relam() != pg_am.oid() { + pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); + } + let relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let opfamily = unsafe { crate::index::vchordrq::opclass::opfamily(relation.raw()) }; + let index = unsafe { PostgresRelation::new(relation.raw()) }; + crate::index::vchordrq::dispatch::prewarm(opfamily, &index, height) +} + +struct Index { + raw: *mut pgrx::pg_sys::RelationData, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl Index { + fn open(indexrelid: Oid, lockmode: pgrx::pg_sys::LOCKMASK) -> Self { + Self { + raw: unsafe { pgrx::pg_sys::index_open(indexrelid, lockmode) }, + lockmode, + } + } + fn raw(&self) -> *mut pgrx::pg_sys::RelationData { + self.raw + } +} + +impl Drop for Index { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::index_close(self.raw, self.lockmode); + } + } +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_sampled_values(indexrelid: Oid) -> SetOfIterator<'static, String> { + let pg_am = PgAm::search_amname(c"vchordrq").unwrap(); + let Some(pg_am) = pg_am.get() else { + pgrx::error!("vchord is not installed"); + }; + let pg_class = PgClass::search_reloid(indexrelid).unwrap(); + let Some(pg_class) = pg_class.get() else { + pgrx::error!("the relation does not exist"); + }; + if pg_class.relkind() != PgClassRelkind::Index { + pgrx::error!("the relation {:?} is not an index", pg_class.relname()); + } + if pg_class.relam() != pg_am.oid() { + pgrx::error!("the index {:?} is not a vchordrq index", pg_class.relname()); + } + // The user must have access to the index, if not, raise an error from Postgres. + let _relation = Index::open(indexrelid, pgrx::pg_sys::AccessShareLock as _); + let queries = dump(indexrelid.to_u32()); + SetOfIterator::new(queries) +} diff --git a/src/index/gucs.rs b/src/index/gucs.rs new file mode 100644 index 00000000..a0c831a4 --- /dev/null +++ b/src/index/gucs.rs @@ -0,0 +1,718 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::scanners::Io; +use pgrx::guc::{GucContext, GucFlags, GucRegistry, GucSetting, PostgresGucEnum}; +use std::ffi::{CStr, CString}; + +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresIo { + #[name = c"read_buffer"] + ReadBuffer, + #[name = c"prefetch_buffer"] + PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + #[name = c"read_stream"] + ReadStream, +} + +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresMaxsimBackend { + #[name = c"coarse_only"] + CoarseOnly, + #[name = c"cpu_exact"] + CpuExact, + #[name = c"gpu"] + Gpu, + #[name = c"auto"] + Auto, +} + +static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting = GucSetting::::new(false); + +static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); + +static VCHORDRQ_QUERY_SAMPLING_RATE: GucSetting = GucSetting::::new(0.0); + +static VCHORDG_ENABLE_SCAN: GucSetting = GucSetting::::new(true); + +static VCHORDG_EF_SEARCH: GucSetting = GucSetting::::new(64); + +static mut VCHORDG_EF_SEARCH_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); + +static VCHORDG_BEAM_SEARCH: GucSetting = GucSetting::::new(1); + +static VCHORDG_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); + +static VCHORDG_IO_SEARCH: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); + +static VCHORDG_IO_RERANK: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); + +static VCHORDRQ_ENABLE_SCAN: GucSetting = GucSetting::::new(true); + +static VCHORDRQ_PROBES: GucSetting> = GucSetting::>::new(Some(c"")); + +static mut VCHORDRQ_PROBES_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); + +static VCHORDRQ_EPSILON: GucSetting = GucSetting::::new(1.9); + +static mut VCHORDRQ_EPSILON_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); + +static VCHORDRQ_MAX_SCAN_TUPLES: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_REFINE: GucSetting = GucSetting::::new(0); + +static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); + +static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); + +static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = + core::ptr::null_mut(); + +static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_PROFILE: GucSetting = GucSetting::::new(false); + +static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); + +static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); + +static VCHORDRQ_MAXSIM_BACKEND: GucSetting = + GucSetting::::new(PostgresMaxsimBackend::CoarseOnly); + +static VCHORDRQ_MAXSIM_GPU_ENDPOINT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS: GucSetting = GucSetting::::new(2000); + +static VCHORDRQ_MAXSIM_TENANT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_PRIORITY: GucSetting = GucSetting::::new(0); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS: GucSetting = GucSetting::::new(1_000_000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES: GucSetting = GucSetting::::new(1_073_741_824); + +static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); + +static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); + +static VCHORDRQ_IO_RERANK: GucSetting = GucSetting::::new( + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + PostgresIo::PrefetchBuffer, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream, +); + +pub fn init() { + GucRegistry::define_bool_guc( + c"vchordrq.enable_scan", + c"`enable_scan` argument of vchordrq.", + c"`enable_scan` argument of vchordrq.", + &VCHORDRQ_ENABLE_SCAN, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.probes", + c"`probes` argument of vchordrq.", + c"`probes` argument of vchordrq.", + &VCHORDRQ_PROBES, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_float_guc( + c"vchordrq.epsilon", + c"`epsilon` argument of vchordrq.", + c"`epsilon` argument of vchordrq.", + &VCHORDRQ_EPSILON, + 0.0, + 4.0, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.max_scan_tuples", + c"`max_scan_tuples` argument of vchordrq.", + c"`max_scan_tuples` argument of vchordrq.", + &VCHORDRQ_MAX_SCAN_TUPLES, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_refine", + c"`maxsim_refine` argument of vchordrq.", + c"`maxsim_refine` argument of vchordrq.", + &VCHORDRQ_MAXSIM_REFINE, + 0, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_threshold", + c"`maxsim_threshold` argument of vchordrq.", + c"`maxsim_threshold` argument of vchordrq.", + &VCHORDRQ_MAXSIM_THRESHOLD, + 0, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_candidate_limit", + c"Maximum number of page candidates produced by MaxSim aggregation.", + c"Use -1 for no page-candidate limit.", + &VCHORDRQ_MAXSIM_CANDIDATE_LIMIT, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_bool_guc( + c"vchordrq.maxsim_profile", + c"Emit one client NOTICE with MaxSim phase timings and counters.", + c"Disabled by default; emitted profiles contain no tensor references or application IDs.", + &VCHORDRQ_MAXSIM_PROFILE, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_query_tokens", + c"Expected MaxSim query-token count used by the planner.", + c"Set this to the measured deployment average until expression statistics are available.", + &VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_document_tokens", + c"Fallback MaxSim document-token count used by the planner.", + c"Used for indexes that predate the native indexed-vector statistic; set it to the measured deployment average.", + &VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.maxsim_backend", + c"Page-level MaxSim rerank backend.", + c"GPU and auto modes require the native sidecar integration.", + &VCHORDRQ_MAXSIM_BACKEND, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_gpu_endpoint", + c"Unix-socket endpoint for the TileMaxSim sidecar.", + c"An empty endpoint disables GPU transport.", + &VCHORDRQ_MAXSIM_GPU_ENDPOINT, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_timeout_ms", + c"Overall TileMaxSim sidecar deadline in milliseconds.", + c"The deadline covers connection, request write, and response read.", + &VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS, + 1, + 600_000, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_tenant", + c"Logical tenant used only by the TileMaxSim GPU scheduler.", + c"Authorization and source filtering must be completed before setting this request-local value.", + &VCHORDRQ_MAXSIM_TENANT, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_priority", + c"TileMaxSim request priority; larger values run first subject to tenant fairness and aging.", + c"The allowed request-local range is -100 through 100.", + &VCHORDRQ_MAXSIM_PRIORITY, + -100, + 100, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_tokens", + c"Maximum query plus document tokens in one TileMaxSim request.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS, + 1, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_bytes", + c"Maximum encoded TileMaxSim request size in bytes.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES, + 1024, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_bool_guc( + c"vchordrq.prefilter", + c"`prefilter` argument of vchordrq.", + c"`prefilter` argument of vchordrq.", + &VCHORDRQ_PREFILTER, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.io_search", + c"`io_search` argument of vchordrq.", + c"`io_search` argument of vchordrq.", + &VCHORDRQ_IO_SEARCH, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.io_rerank", + c"`io_rerank` argument of vchordrq.", + c"`io_rerank` argument of vchordrq.", + &VCHORDRQ_IO_RERANK, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_bool_guc( + c"vchordrq.query_sampling_enable", + c"`query_sampling_enable` argument of vchordrq.", + c"`query_sampling_enable` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_ENABLE, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.query_sampling_max_records", + c"`query_sampling_max_records` argument of vchordrq.", + c"`query_sampling_max_records` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS, + 0, + 10000, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_float_guc( + c"vchordrq.query_sampling_rate", + c"`query_sampling_rate` argument of vchordrq.", + c"`query_sampling_rate` argument of vchordrq.", + &VCHORDRQ_QUERY_SAMPLING_RATE, + 0.0, + 1.0, + GucContext::Suset, + GucFlags::default(), + ); + unsafe { + #[cfg(feature = "pg14")] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordrq".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordrq".as_ptr()); + } + GucRegistry::define_bool_guc( + c"vchordg.enable_scan", + c"`enable_scan` argument of vchordg.", + c"`enable_scan` argument of vchordg.", + &VCHORDG_ENABLE_SCAN, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordg.ef_search", + c"`ef_search` argument of vchordg.", + c"`ef_search` argument of vchordg.", + &VCHORDG_EF_SEARCH, + 1, + 65535, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordg.beam_search", + c"`beam_search` argument of vchordg.", + c"`beam_search` argument of vchordg.", + &VCHORDG_BEAM_SEARCH, + 1, + 65535, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordg.max_scan_tuples", + c"`max_scan_tuples` argument of vchordg.", + c"`max_scan_tuples` argument of vchordg.", + &VCHORDG_MAX_SCAN_TUPLES, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordg.io_search", + c"`io_search` argument of vchordg.", + c"`io_search` argument of vchordg.", + &VCHORDG_IO_SEARCH, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordg.io_rerank", + c"`io_rerank` argument of vchordg.", + c"`io_rerank` argument of vchordg.", + &VCHORDG_IO_RERANK, + GucContext::Userset, + GucFlags::default(), + ); + unsafe { + #[cfg(feature = "pg14")] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchordg".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchordg".as_ptr()); + } + assert!(crate::is_main()); + let targets = vec![ + (c"vchordg.ef_search", &raw mut VCHORDG_EF_SEARCH_CONFIG), + (c"vchordrq.epsilon", &raw mut VCHORDRQ_EPSILON_CONFIG), + ( + c"vchordrq.maxsim_refine", + &raw mut VCHORDRQ_MAXSIM_REFINE_CONFIG, + ), + ( + c"vchordrq.maxsim_threshold", + &raw mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG, + ), + (c"vchordrq.probes", &raw mut VCHORDRQ_PROBES_CONFIG), + ]; + #[cfg(any(feature = "pg14", feature = "pg15"))] + unsafe { + let len = pgrx::pg_sys::GetNumConfigOptions() as usize; + let arr = pgrx::pg_sys::get_guc_variables(); + let mut sources = (0..len).map(|i| arr.add(i).read()); + debug_assert!(targets.is_sorted_by(|(a, _), (b, _)| guc_name_compare(a, b).is_le())); + for (name, ptr) in targets { + *ptr = loop { + if let Some(source) = sources.next() { + if !(*source).name.is_null() && CStr::from_ptr((*source).name) == name { + break source; + } else { + continue; + } + } else { + pgrx::error!("failed to find GUC {name:?}"); + } + }; + assert!(check(*ptr, name), "failed to find GUC {name:?}"); + } + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + unsafe { + use pgrx::pg_sys::PGERROR; + for (name, ptr) in targets { + *ptr = pgrx::pg_sys::find_option(name.as_ptr(), false, false, PGERROR as _); + assert!(check(*ptr, name), "failed to find GUC {name:?}"); + } + } +} + +unsafe fn check(p: *mut pgrx::pg_sys::config_generic, name: &CStr) -> bool { + if p.is_null() { + return false; + } + if unsafe { (*p).flags } & pgrx::pg_sys::GUC_CUSTOM_PLACEHOLDER as core::ffi::c_int != 0 { + return false; + } + if unsafe { (*p).name }.is_null() { + return false; + } + if unsafe { CStr::from_ptr((*p).name) != name } { + return false; + } + true +} + +pub fn vchordg_enable_scan() -> bool { + VCHORDG_ENABLE_SCAN.get() +} + +pub fn vchordg_ef_search(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 64; + if unsafe { (*VCHORDG_EF_SEARCH_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDG_EF_SEARCH.get(); + parse(value) + } else { + use crate::index::vchordg::am::Reloption; + let value = unsafe { Reloption::ef_search((*index).rd_options as _, DEFAULT) }; + parse(value) + } +} + +pub fn vchordg_beam_search() -> u32 { + VCHORDG_BEAM_SEARCH.get() as u32 +} + +pub fn vchordg_max_scan_tuples() -> Option { + let x = VCHORDG_MAX_SCAN_TUPLES.get(); + if x < 0 { None } else { Some(x as u32) } +} + +pub fn vchordg_io_search() -> Io { + match VCHORDG_IO_SEARCH.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } +} + +pub fn vchordg_io_rerank() -> Io { + match VCHORDG_IO_RERANK.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } +} + +pub fn vchordrq_enable_scan() -> bool { + VCHORDRQ_ENABLE_SCAN.get() +} + +pub unsafe fn vchordrq_probes(index: pgrx::pg_sys::Relation) -> Vec { + fn parse(value: &CStr) -> Vec { + let mut result = Vec::new(); + let mut current = None; + for &c in value.to_bytes() { + match c { + b' ' => continue, + b',' => match current.take() { + Some(value) => result.push(value), + None => pgrx::error!("empty entry in probes"), + }, + b'0'..=b'9' => { + if let Some(x) = current.as_mut() { + *x = *x * 10 + (c - b'0') as u32; + } else { + current = Some((c - b'0') as u32); + } + } + c => pgrx::error!("unknown character in probes: ASCII = {c}"), + } + } + if let Some(current) = current { + result.push(current); + } + result + } + assert!(crate::is_main()); + const DEFAULT: &CStr = c""; + if unsafe { (*VCHORDRQ_PROBES_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDRQ_PROBES.get(); + parse(value.as_deref().unwrap_or(DEFAULT)) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::probes((*index).rd_options as _, DEFAULT) }; + parse(value) + } +} + +pub unsafe fn vchordrq_epsilon(index: pgrx::pg_sys::Relation) -> f32 { + fn parse(x: f64) -> f32 { + x as f32 + } + assert!(crate::is_main()); + const DEFAULT: f64 = 1.9; + if unsafe { (*VCHORDRQ_EPSILON_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT { + let value = VCHORDRQ_EPSILON.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::epsilon((*index).rd_options as _, DEFAULT) }; + parse(value) + } +} + +pub fn vchordrq_max_scan_tuples() -> Option { + let x = VCHORDRQ_MAX_SCAN_TUPLES.get(); + if x < 0 { None } else { Some(x as u32) } +} + +pub fn vchordrq_maxsim_refine(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 0; + if unsafe { (*VCHORDRQ_MAXSIM_REFINE_CONFIG).source } != pgrx::pg_sys::GucSource::PGC_S_DEFAULT + { + let value = VCHORDRQ_MAXSIM_REFINE.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::maxsim_refine((*index).rd_options as _, DEFAULT) }; + parse(value) + } +} + +pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 { + fn parse(x: i32) -> u32 { + x as u32 + } + assert!(crate::is_main()); + const DEFAULT: i32 = 0; + if unsafe { (*VCHORDRQ_MAXSIM_THRESHOLD_CONFIG).source } + != pgrx::pg_sys::GucSource::PGC_S_DEFAULT + { + let value = VCHORDRQ_MAXSIM_THRESHOLD.get(); + parse(value) + } else { + use crate::index::vchordrq::am::Reloption; + let value = unsafe { Reloption::maxsim_threshold((*index).rd_options as _, DEFAULT) }; + parse(value) + } +} + +pub fn vchordrq_maxsim_candidate_limit() -> Option { + let value = VCHORDRQ_MAXSIM_CANDIDATE_LIMIT.get(); + if value < 0 { None } else { Some(value as u32) } +} + +pub fn vchordrq_maxsim_profile() -> bool { + VCHORDRQ_MAXSIM_PROFILE.get() +} + +pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_planner_document_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_backend() -> PostgresMaxsimBackend { + VCHORDRQ_MAXSIM_BACKEND.get() +} + +pub fn vchordrq_maxsim_gpu_endpoint() -> Option { + VCHORDRQ_MAXSIM_GPU_ENDPOINT.get() +} + +pub fn vchordrq_maxsim_gpu_timeout_ms() -> u32 { + VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS.get() as u32 +} + +pub fn vchordrq_maxsim_tenant() -> Option { + VCHORDRQ_MAXSIM_TENANT.get().and_then(|tenant| { + let tenant = tenant.to_string_lossy().into_owned(); + (!tenant.is_empty()).then_some(tenant) + }) +} + +pub fn vchordrq_maxsim_priority() -> i32 { + VCHORDRQ_MAXSIM_PRIORITY.get() +} + +pub fn vchordrq_maxsim_gpu_max_batch_tokens() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_bytes() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES.get() as u32 +} + +pub fn vchordrq_prefilter() -> bool { + VCHORDRQ_PREFILTER.get() +} + +pub fn vchordrq_io_search() -> Io { + match VCHORDRQ_IO_SEARCH.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } +} + +pub fn vchordrq_io_rerank() -> Io { + match VCHORDRQ_IO_RERANK.get() { + PostgresIo::ReadBuffer => Io::Plain, + PostgresIo::PrefetchBuffer => Io::Simple, + #[cfg(any(feature = "pg17", feature = "pg18"))] + PostgresIo::ReadStream => Io::Stream, + } +} + +pub fn vchordrq_query_sampling_enable() -> bool { + VCHORDRQ_QUERY_SAMPLING_ENABLE.get() +} + +pub fn vchordrq_query_sampling_max_records() -> u32 { + VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS.get() as u32 +} + +pub fn vchordrq_query_sampling_rate() -> f64 { + VCHORDRQ_QUERY_SAMPLING_RATE.get() +} + +#[allow(dead_code)] +fn guc_name_compare(a: &CStr, b: &CStr) -> std::cmp::Ordering { + let (a, b) = (a.to_bytes_with_nul(), b.to_bytes_with_nul()); + let mut i = 0; + while a[i] != 0 && b[i] != 0 { + let a = a[i].to_ascii_lowercase(); + let b = b[i].to_ascii_lowercase(); + if a != b { + return Ord::cmp(&a, &b); + } + i += 1; + } + if b[i] != 0 { + std::cmp::Ordering::Less + } else if a[i] != 0 { + std::cmp::Ordering::Greater + } else { + std::cmp::Ordering::Equal + } +} diff --git a/src/index/hook.rs b/src/index/hook.rs new file mode 100644 index 00000000..7a35f0b0 --- /dev/null +++ b/src/index/hook.rs @@ -0,0 +1,145 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn rewrite_plan_state( + node: *mut pgrx::pg_sys::PlanState, + context: *mut core::ffi::c_void, +) -> bool { + unsafe fn dirty_check_vchordg(index_relation: *mut pgrx::pg_sys::RelationData) -> Option { + type FnPtr = unsafe extern "C-unwind" fn( + *mut pgrx::pg_sys::RelationData, + i32, + i32, + ) -> *mut pgrx::pg_sys::IndexScanDescData; + unsafe { + let index_relation = index_relation.as_ref()?; + let indam = index_relation.rd_indam.as_ref()?; + let ambeginscan = indam.ambeginscan.as_ref()?; + Some(core::ptr::fn_addr_eq::( + *ambeginscan, + crate::index::vchordg::am::ambeginscan, + )) + } + } + + unsafe fn dirty_check_vchordrq( + index_relation: *mut pgrx::pg_sys::RelationData, + ) -> Option { + type FnPtr = unsafe extern "C-unwind" fn( + *mut pgrx::pg_sys::RelationData, + i32, + i32, + ) -> *mut pgrx::pg_sys::IndexScanDescData; + unsafe { + let index_relation = index_relation.as_ref()?; + let indam = index_relation.rd_indam.as_ref()?; + let ambeginscan = indam.ambeginscan.as_ref()?; + Some(core::ptr::fn_addr_eq::( + *ambeginscan, + crate::index::vchordrq::am::ambeginscan, + )) + } + } + + unsafe { + if (*node).type_ == pgrx::pg_sys::NodeTag::T_IndexScanState { + let node = node as *mut pgrx::pg_sys::IndexScanState; + let index_relation = (*node).iss_RelationDesc; + if (*node).iss_ScanDesc.is_null() { + if Some(true) == dirty_check_vchordg(index_relation) { + use crate::index::vchordg::am::Scanner; + + (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( + (*node).ss.ss_currentRelation, + (*node).iss_RelationDesc, + (*(*node).ss.ps.state).es_snapshot, + #[cfg(feature = "pg18")] + std::ptr::null_mut(), + (*node).iss_NumScanKeys, + (*node).iss_NumOrderByKeys, + ); + + let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); + scanner.hack = std::ptr::NonNull::new(node); + + if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { + pgrx::pg_sys::index_rescan( + (*node).iss_ScanDesc, + (*node).iss_ScanKeys, + (*node).iss_NumScanKeys, + (*node).iss_OrderByKeys, + (*node).iss_NumOrderByKeys, + ); + } + } + if Some(true) == dirty_check_vchordrq(index_relation) { + use crate::index::vchordrq::am::Scanner; + + (*node).iss_ScanDesc = pgrx::pg_sys::index_beginscan( + (*node).ss.ss_currentRelation, + (*node).iss_RelationDesc, + (*(*node).ss.ps.state).es_snapshot, + #[cfg(feature = "pg18")] + std::ptr::null_mut(), + (*node).iss_NumScanKeys, + (*node).iss_NumOrderByKeys, + ); + + let scanner = &mut *((*(*node).iss_ScanDesc).opaque as *mut Scanner); + scanner.hack = std::ptr::NonNull::new(node); + + if (*node).iss_NumRuntimeKeys == 0 || (*node).iss_RuntimeKeysReady { + pgrx::pg_sys::index_rescan( + (*node).iss_ScanDesc, + (*node).iss_ScanKeys, + (*node).iss_NumScanKeys, + (*node).iss_OrderByKeys, + (*node).iss_NumOrderByKeys, + ); + } + } + } + } + pgrx::pg_sys::planstate_tree_walker(node, Some(rewrite_plan_state), context) + } +} + +static mut PREV_EXECUTOR_START: pgrx::pg_sys::ExecutorStart_hook_type = None; + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn executor_start( + query_desc: *mut pgrx::pg_sys::QueryDesc, + eflags: core::ffi::c_int, +) { + unsafe { + use core::ptr::null_mut; + use pgrx::pg_sys::submodules::ffi::pg_guard_ffi_boundary; + if let Some(prev_executor_start) = PREV_EXECUTOR_START { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| prev_executor_start(query_desc, eflags)) + } else { + pgrx::pg_sys::standard_ExecutorStart(query_desc, eflags) + } + pg_guard_ffi_boundary(|| rewrite_plan_state((*query_desc).planstate, null_mut())); + } +} + +pub fn init() { + assert!(crate::is_main()); + unsafe { + PREV_EXECUTOR_START = pgrx::pg_sys::ExecutorStart_hook; + pgrx::pg_sys::ExecutorStart_hook = Some(executor_start); + } +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 489f18f7..154e7489 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -1,10 +1,32 @@ -pub mod am; -pub mod am_options; -pub mod am_scan; -pub mod utils; +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. -pub unsafe fn init() { - unsafe { - am::init(); - } +mod fetcher; +mod functions; +mod gucs; +mod hook; +mod opclass; +mod sample; +mod scanners; +mod storage; +mod traverse; +mod vchordg; +mod vchordrq; + +pub fn init() { + gucs::init(); + hook::init(); + vchordrq::am::init(); + vchordg::am::init(); } diff --git a/src/index/opclass.rs b/src/index/opclass.rs new file mode 100644 index 00000000..32ee8b20 --- /dev/null +++ b/src/index/opclass.rs @@ -0,0 +1,158 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_l2_ops() -> String { + "vchordg_vector_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_cosine_ops() -> String { + "vchordg_vector_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_vector_ip_ops() -> String { + "vchordg_vector_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_l2_ops() -> String { + "vchordg_halfvec_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_cosine_ops() -> String { + "vchordg_halfvec_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_halfvec_ip_ops() -> String { + "vchordg_halfvec_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_l2_ops() -> String { + "vchordg_rabitq8_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_cosine_ops() -> String { + "vchordg_rabitq8_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq8_ip_ops() -> String { + "vchordg_rabitq8_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_l2_ops() -> String { + "vchordg_rabitq4_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_cosine_ops() -> String { + "vchordg_rabitq4_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordg_support_rabitq4_ip_ops() -> String { + "vchordg_rabitq4_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_vector_l2_ops() -> String { + "vchordrq_vector_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_vector_ip_ops() -> String { + "vchordrq_vector_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_vector_cosine_ops() -> String { + "vchordrq_vector_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_l2_ops() -> String { + "vchordrq_halfvec_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_ip_ops() -> String { + "vchordrq_halfvec_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_cosine_ops() -> String { + "vchordrq_halfvec_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_l2_ops() -> String { + "vchordrq_rabitq8_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_ip_ops() -> String { + "vchordrq_rabitq8_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_cosine_ops() -> String { + "vchordrq_rabitq8_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_l2_ops() -> String { + "vchordrq_rabitq4_l2_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_ip_ops() -> String { + "vchordrq_rabitq4_ip_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_cosine_ops() -> String { + "vchordrq_rabitq4_cosine_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_vector_maxsim_ops() -> String { + "vchordrq_vector_maxsim_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_halfvec_maxsim_ops() -> String { + "vchordrq_halfvec_maxsim_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq8_maxsim_ops() -> String { + "vchordrq_rabitq8_maxsim_ops".to_string() +} + +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vchordrq_support_rabitq4_maxsim_ops() -> String { + "vchordrq_rabitq4_maxsim_ops".to_string() +} + +pub struct Sphere { + pub center: T, + pub radius: f32, +} diff --git a/src/index/sample.rs b/src/index/sample.rs new file mode 100644 index 00000000..3bd25493 --- /dev/null +++ b/src/index/sample.rs @@ -0,0 +1,262 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::pg_sys::{Datum, ItemPointerData}; +use std::ptr::NonNull; + +pub trait Tuple { + #[expect(dead_code)] + fn id(&mut self) -> ItemPointerData; + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); +} + +pub trait Sample { + type Tuple<'a>: Tuple + where + Self: 'a; + + fn next(&mut self) -> Option>; +} + +pub trait Sampler { + type Sample: Sample; + + fn sample(&self) -> Self::Sample; +} + +pub struct HeapSampler { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, +} + +impl HeapSampler { + pub unsafe fn new( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + Self { + heap_relation, + index_relation, + snapshot, + } + } +} + +impl Drop for HeapSampler { + fn drop(&mut self) {} +} + +impl Sampler for HeapSampler { + type Sample = HeapSample; + + fn sample(&self) -> Self::Sample { + unsafe { + let state = NonNull::new_unchecked(Box::into_raw(Box::new(State { + blocks: None, + tuples: None, + }))); + let sample_scan_state = + NonNull::new_unchecked(Box::into_raw(Box::new(pgrx::pg_sys::SampleScanState { + tsmroutine: (&raw const TSM.0).cast_mut().cast(), + tsm_state: state.as_ptr().cast(), + ..core::mem::zeroed() + }))); + let table_scan_desc = pgrx::pg_sys::table_beginscan_sampling( + self.heap_relation, + self.snapshot, + 0, + std::ptr::null_mut(), + true, + false, + true, + ); + let index_info = pgrx::pg_sys::BuildIndexInfo(self.index_relation); + let estate = pgrx::pg_sys::CreateExecutorState(); + let econtext = pgrx::pg_sys::MakePerTupleExprContext(estate); + HeapSample { + index_info, + estate, + econtext, + slot: pgrx::pg_sys::table_slot_create(self.heap_relation, std::ptr::null_mut()), + values: [Datum::null(); 32], + is_nulls: [true; 32], + state, + sample_scan_state, + table_scan_desc, + done: false, + have_block: false, + } + } + } +} + +pub struct HeapSample { + index_info: *mut pgrx::pg_sys::IndexInfo, + estate: *mut pgrx::pg_sys::EState, + econtext: *mut pgrx::pg_sys::ExprContext, + slot: *mut pgrx::pg_sys::TupleTableSlot, + values: [Datum; 32], + is_nulls: [bool; 32], + state: NonNull, + sample_scan_state: NonNull, + table_scan_desc: pgrx::pg_sys::TableScanDesc, + done: bool, + have_block: bool, +} + +impl Drop for HeapSample { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::MemoryContextReset((*self.econtext).ecxt_per_tuple_memory); + // free common resources + pgrx::pg_sys::table_endscan(self.table_scan_desc); + let _ = Box::from_raw(self.sample_scan_state.as_ptr()); + let _ = Box::from_raw(self.state.as_ptr()); + pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); + pgrx::pg_sys::FreeExecutorState(self.estate); + } + } +} + +impl Sample for HeapSample { + type Tuple<'a> = HeapTuple<'a>; + + fn next(&mut self) -> Option> { + unsafe { + use pgrx::pg_sys::{table_scan_sample_next_block, table_scan_sample_next_tuple}; + + if self.done { + return None; + } + + loop { + if !self.have_block { + if !table_scan_sample_next_block( + self.table_scan_desc, + self.sample_scan_state.as_ptr(), + ) { + self.have_block = false; + self.done = true; + return None; + } + + self.have_block = true; + } + + if !table_scan_sample_next_tuple( + self.table_scan_desc, + self.sample_scan_state.as_ptr(), + self.slot, + ) { + self.have_block = false; + continue; + } + + break; + } + + Some(HeapTuple { this: self }) + } + } +} + +pub struct HeapTuple<'a> { + this: &'a mut HeapSample, +} + +impl Tuple for HeapTuple<'_> { + fn id(&mut self) -> ItemPointerData { + unsafe { + let this = &mut self.this; + (*this.slot).tts_tid + } + } + fn build(&mut self) -> (&[Datum; 32], &[bool; 32]) { + unsafe { + let this = &mut self.this; + (*this.econtext).ecxt_scantuple = this.slot; + pgrx::pg_sys::MemoryContextReset((*this.econtext).ecxt_per_tuple_memory); + pgrx::pg_sys::FormIndexDatum( + this.index_info, + this.slot, + this.estate, + this.values.as_mut_ptr(), + this.is_nulls.as_mut_ptr(), + ); + (&this.values, &this.is_nulls) + } + } +} + +fn sample(n: u32) -> Box> { + let width = (n.ilog2() + 1).next_multiple_of(2); + let key_0 = rand::RngExt::random(&mut rand::rng()); + let key_1 = rand::RngExt::random(&mut rand::rng()); + let secret = move |round: u32, x: u32| { + let buffer = [round.to_le_bytes(), x.to_le_bytes(), key_0, key_1]; + wyhash::wyhash(buffer.as_flattened(), 0) as u32 + }; + let permutation = (0..1 << width) + .map(move |i| feistel::feistel(width, i, 8, secret)) + .filter(move |&x| x < n); + Box::new(permutation) +} + +pub struct State { + blocks: Option>>, + tuples: Option>, +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn feistel_rows_nextsampleblock( + node: *mut pgrx::pg_sys::SampleScanState, + nblocks: pgrx::pg_sys::BlockNumber, +) -> pgrx::pg_sys::BlockNumber { + let state: &mut State = unsafe { &mut *(*node).tsm_state.cast() }; + let iter = state.blocks.get_or_insert_with(|| sample(nblocks)); + if let Some(number) = iter.next() { + number + } else { + pgrx::pg_sys::InvalidBlockNumber + } +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn feistel_rows_nextsampletuple( + node: *mut pgrx::pg_sys::SampleScanState, + _blockno: pgrx::pg_sys::BlockNumber, + maxoffset: pgrx::pg_sys::OffsetNumber, +) -> pgrx::pg_sys::OffsetNumber { + let state: &mut State = unsafe { &mut *(*node).tsm_state.cast() }; + let iter = state.tuples.get_or_insert(1..=maxoffset); + if let Some(number) = iter.next() { + number + } else { + state.tuples = None; + pgrx::pg_sys::InvalidOffsetNumber + } +} + +struct AssertSync(T); + +unsafe impl Sync for AssertSync {} + +static TSM: AssertSync = AssertSync(pgrx::pg_sys::TsmRoutine { + type_: pgrx::pg_sys::NodeTag::T_TsmRoutine, + NextSampleBlock: Some(feistel_rows_nextsampleblock), + NextSampleTuple: Some(feistel_rows_nextsampletuple), + ..unsafe { core::mem::zeroed() } +}); diff --git a/src/index/scanners.rs b/src/index/scanners.rs new file mode 100644 index 00000000..dd48848b --- /dev/null +++ b/src/index/scanners.rs @@ -0,0 +1,54 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::Fetcher; +use crate::recorder::Recorder; +use index::bump::Bump; +use index::relation::{Page, RelationPrefetch, RelationRead, RelationReadStream}; +use pgrx::pg_sys::Datum; + +#[derive(Debug, Clone, Copy)] +pub enum Io { + Plain, + Simple, + #[cfg_attr( + any(feature = "pg14", feature = "pg15", feature = "pg16"), + expect(dead_code) + )] + Stream, +} + +pub trait SearchBuilder: 'static { + type Options; + + type Opfamily; + + type Opaque: Copy; + + fn new(opfamily: Self::Opfamily) -> Self; + + unsafe fn add(&mut self, strategy: u16, datum: Option); + + fn build<'b, R>( + self, + relation: &'b R, + options: Self::Options, + fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + recorder: impl Recorder, + ) -> Box + 'b> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page; +} diff --git a/src/index/storage.rs b/src/index/storage.rs new file mode 100644 index 00000000..946df2c2 --- /dev/null +++ b/src/index/storage.rs @@ -0,0 +1,723 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod buffered; + +use index::fetch::Fetch; +use index::relation::{ + Hints, Opaque, Page, PageGuard, ReadStream, Relation, RelationPrefetch, RelationRead, + RelationReadStream, RelationReadStreamTypes, RelationReadTypes, RelationWrite, + RelationWriteTypes, +}; +use std::collections::VecDeque; +use std::iter::{Chain, Flatten}; +use std::marker::PhantomData; +use std::mem::{MaybeUninit, offset_of}; +use std::ops::{Deref, DerefMut}; +use std::ptr::NonNull; + +#[repr(C, align(8))] +#[derive(Debug)] +pub struct PostgresPage { + header: pgrx::pg_sys::PageHeaderData, + content: [u8; pgrx::pg_sys::BLCKSZ as usize - size_of::()], + _opaque: PhantomData O>, +} + +// It is a non-guaranteed detection. +// If `PageHeaderData` contains padding bytes, const-eval probably fails. +const _: () = { + use pgrx::pg_sys::PageHeaderData as T; + use std::mem::{transmute, zeroed}; + const _ZERO: &[u8; size_of::()] = unsafe { transmute(&zeroed::()) }; +}; + +// Layout checks of header. +const _: () = { + use pgrx::pg_sys::{MAXIMUM_ALIGNOF, PageHeaderData as T}; + assert!(size_of::() == offset_of!(T, pd_linp)); + assert!(size_of::() % MAXIMUM_ALIGNOF as usize == 0); +}; + +const _: () = assert!(align_of::>() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); +const _: () = assert!(size_of::>() == pgrx::pg_sys::BLCKSZ as usize); + +impl PostgresPage { + pub fn clone_into_boxed(&self) -> Box { + let mut result = Box::new_uninit(); + unsafe { + std::ptr::copy(self as *const Self, result.as_mut_ptr(), 1); + result.assume_init() + } + } +} + +impl Page for PostgresPage { + type Opaque = O; + fn get_opaque(&self) -> &O { + assert!(self.header.pd_special as usize + size_of::() == size_of::()); + unsafe { &*((self as *const _ as *const O).byte_add(self.header.pd_special as _)) } + } + fn get_opaque_mut(&mut self) -> &mut O { + assert!(self.header.pd_special as usize + size_of::() == size_of::()); + unsafe { &mut *((self as *mut _ as *mut O).byte_add(self.header.pd_special as _)) } + } + fn len(&self) -> u16 { + use pgrx::pg_sys::{ItemIdData, PageHeaderData}; + assert!(self.header.pd_lower as usize <= size_of::()); + assert!(self.header.pd_upper as usize <= size_of::()); + let lower = self.header.pd_lower as usize; + let upper = self.header.pd_upper as usize; + assert!(offset_of!(PageHeaderData, pd_linp) <= lower && lower <= upper); + ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16 + } + fn get(&self, i: u16) -> Option<&[u8]> { + use pgrx::pg_sys::{ItemIdData, PageHeaderData}; + if i == 0 { + return None; + } + assert!(self.header.pd_lower as usize <= size_of::()); + let lower = self.header.pd_lower as usize; + assert!(offset_of!(PageHeaderData, pd_linp) <= lower); + let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; + if i > n { + return None; + } + let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; + let lp_off = iid.lp_off() as usize; + let lp_len = iid.lp_len() as usize; + match lp_flags(iid) { + pgrx::pg_sys::LP_UNUSED => return None, + pgrx::pg_sys::LP_NORMAL => (), + pgrx::pg_sys::LP_REDIRECT => unimplemented!(), + pgrx::pg_sys::LP_DEAD => unimplemented!(), + _ => unreachable!(), + } + assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off); + assert!(lp_off <= size_of::()); + assert!(lp_len <= size_of::()); + assert!(lp_off + lp_len <= size_of::()); + unsafe { + let ptr = (self as *const Self).cast::().add(lp_off as _); + Some(std::slice::from_raw_parts(ptr, lp_len as _)) + } + } + fn get_mut(&mut self, i: u16) -> Option<&mut [u8]> { + use pgrx::pg_sys::{ItemIdData, PageHeaderData}; + if i == 0 { + return None; + } + assert!(self.header.pd_lower as usize <= size_of::()); + let lower = self.header.pd_lower as usize; + assert!(offset_of!(PageHeaderData, pd_linp) <= lower); + let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; + if i > n { + return None; + } + let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; + let lp_off = iid.lp_off() as usize; + let lp_len = iid.lp_len() as usize; + match lp_flags(iid) { + pgrx::pg_sys::LP_UNUSED => return None, + pgrx::pg_sys::LP_NORMAL => (), + pgrx::pg_sys::LP_REDIRECT => unimplemented!(), + pgrx::pg_sys::LP_DEAD => unimplemented!(), + _ => unreachable!(), + } + assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off); + assert!(lp_off <= size_of::()); + assert!(lp_len <= size_of::()); + assert!(lp_off + lp_len <= size_of::()); + unsafe { + let ptr = (self as *mut Self).cast::().add(lp_off as _); + Some(std::slice::from_raw_parts_mut(ptr, lp_len as _)) + } + } + fn alloc(&mut self, data: &[u8]) -> Option { + unsafe { + let i = pgrx::pg_sys::PageAddItemExtended( + (self as *const Self).cast_mut().cast(), + data.as_ptr().cast_mut().cast(), + data.len(), + 0, + 0, + ); + if i == 0 { None } else { Some(i) } + } + } + fn free(&mut self, i: u16) { + unsafe { + pgrx::pg_sys::PageIndexTupleDeleteNoCompact((self as *mut Self).cast(), i); + } + } + fn freespace(&self) -> u16 { + unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } + } + fn clear(&mut self, opaque: O) { + unsafe { + page_init(self, opaque); + } + } +} + +unsafe fn page_init(this: *mut PostgresPage, opaque: O) { + unsafe { + use pgrx::pg_sys::{BLCKSZ, PageHeaderData, PageInit}; + PageInit(this.cast(), BLCKSZ as usize, size_of::()); + assert_eq!( + (*this.cast::()).pd_special as usize + size_of::(), + size_of::>() + ); + this.cast::() + .byte_add(size_of::>() - size_of::()) + .write(opaque); + } +} + +pub struct PostgresBufferReadGuard { + buf: i32, + page: NonNull>, + id: u32, +} + +impl PageGuard for PostgresBufferReadGuard { + fn id(&self) -> u32 { + self.id + } +} + +impl Deref for PostgresBufferReadGuard { + type Target = PostgresPage; + + fn deref(&self) -> &PostgresPage { + unsafe { self.page.as_ref() } + } +} + +impl Drop for PostgresBufferReadGuard { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::UnlockReleaseBuffer(self.buf); + } + } +} + +pub struct PostgresBufferWriteGuard { + raw: pgrx::pg_sys::Relation, + buf: i32, + page: NonNull>, + state: *mut pgrx::pg_sys::GenericXLogState, + id: u32, + tracking_freespace: bool, +} + +impl PageGuard for PostgresBufferWriteGuard { + fn id(&self) -> u32 { + self.id + } +} + +impl Deref for PostgresBufferWriteGuard { + type Target = PostgresPage; + + fn deref(&self) -> &PostgresPage { + unsafe { self.page.as_ref() } + } +} + +impl DerefMut for PostgresBufferWriteGuard { + fn deref_mut(&mut self) -> &mut PostgresPage { + unsafe { self.page.as_mut() } + } +} + +impl Drop for PostgresBufferWriteGuard { + fn drop(&mut self) { + unsafe { + if std::thread::panicking() { + pgrx::pg_sys::GenericXLogAbort(self.state); + } else { + if self.tracking_freespace { + pgrx::pg_sys::RecordPageWithFreeSpace(self.raw, self.id, self.freespace() as _); + pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, self.id, self.id + 1); + } + pgrx::pg_sys::GenericXLogFinish(self.state); + } + pgrx::pg_sys::UnlockReleaseBuffer(self.buf); + } + } +} + +#[derive(Debug, Clone)] +pub struct PostgresRelation { + raw: pgrx::pg_sys::Relation, + _phantom: PhantomData Opaque>, +} + +impl PostgresRelation { + pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { + Self { + raw, + _phantom: PhantomData, + } + } +} + +impl Relation for PostgresRelation { + type Page = PostgresPage; +} + +impl RelationReadTypes for PostgresRelation { + type ReadGuard<'a> = PostgresBufferReadGuard; +} + +impl RelationRead for PostgresRelation { + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + assert!(id != u32::MAX, "no such page"); + unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_SHARE, BufferGetPage, ForkNumber, LockBuffer, ReadBufferExtended, + ReadBufferMode, + }; + let buf = ReadBufferExtended( + self.raw, + ForkNumber::MAIN_FORKNUM, + id, + ReadBufferMode::RBM_NORMAL, + std::ptr::null_mut(), + ); + LockBuffer(buf, BUFFER_LOCK_SHARE as _); + let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); + PostgresBufferReadGuard { buf, page, id } + } + } +} + +impl RelationWriteTypes for PostgresRelation { + type WriteGuard<'a> = PostgresBufferWriteGuard; +} + +impl RelationWrite for PostgresRelation { + fn write(&self, id: u32, tracking_freespace: bool) -> PostgresBufferWriteGuard { + assert!(id != u32::MAX, "no such page"); + unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_EXCLUSIVE, ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, + LockBuffer, ReadBufferExtended, ReadBufferMode, + }; + let buf = ReadBufferExtended( + self.raw, + ForkNumber::MAIN_FORKNUM, + id, + ReadBufferMode::RBM_NORMAL, + std::ptr::null_mut(), + ); + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); + let state = GenericXLogStart(self.raw); + let page = NonNull::new( + GenericXLogRegisterBuffer(state, buf, 0).cast::>>(), + ) + .expect("failed to get page"); + PostgresBufferWriteGuard { + raw: self.raw, + buf, + page: page.cast(), + state, + id, + tracking_freespace, + } + } + } + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> PostgresBufferWriteGuard { + unsafe { + use pgrx::pg_sys::{ + GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, GenericXLogStart, + }; + let buf; + #[cfg(any(feature = "pg14", feature = "pg15"))] + { + use pgrx::pg_sys::{ + BUFFER_LOCK_EXCLUSIVE, ExclusiveLock, ForkNumber, LockBuffer, + LockRelationForExtension, ReadBufferExtended, ReadBufferMode, + UnlockRelationForExtension, + }; + LockRelationForExtension(self.raw, ExclusiveLock as _); + buf = ReadBufferExtended( + self.raw, + ForkNumber::MAIN_FORKNUM, + u32::MAX, + ReadBufferMode::RBM_NORMAL, + std::ptr::null_mut(), + ); + UnlockRelationForExtension(self.raw, ExclusiveLock as _); + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + { + use pgrx::pg_sys::{ + BufferManagerRelation, ExtendBufferedFlags, ExtendBufferedRel, ForkNumber, + }; + let bmr = BufferManagerRelation { + rel: self.raw, + smgr: std::ptr::null_mut(), + relpersistence: 0, + }; + buf = ExtendBufferedRel( + bmr, + ForkNumber::MAIN_FORKNUM, + std::ptr::null_mut(), + ExtendBufferedFlags::EB_LOCK_FIRST as _, + ); + } + let state = GenericXLogStart(self.raw); + let mut page = NonNull::new( + GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) + .cast::>>(), + ) + .expect("failed to get page"); + page_init(page.as_mut().as_mut_ptr(), opaque); + PostgresBufferWriteGuard { + raw: self.raw, + buf, + page: page.cast(), + state, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + tracking_freespace, + } + } + } + fn search(&self, freespace: usize) -> Option> { + unsafe { + loop { + let id = pgrx::pg_sys::GetPageWithFreeSpace(self.raw, freespace); + if id == u32::MAX { + return None; + } + let write = self.write(id, true); + if write.freespace() < freespace as _ { + // the free space is recorded incorrectly + pgrx::pg_sys::RecordPageWithFreeSpace(self.raw, id, write.freespace() as _); + pgrx::pg_sys::FreeSpaceMapVacuumRange(self.raw, id, id + 1); + continue; + } + return Some(write); + } + } + } +} + +impl RelationPrefetch for PostgresRelation { + fn prefetch(&self, id: u32) { + assert!(id != u32::MAX, "no such page"); + unsafe { + use pgrx::pg_sys::PrefetchBuffer; + PrefetchBuffer(self.raw, 0, id); + } + } +} + +pub struct Cache<'b, I: Iterator> { + window: VecDeque, + tail: VecDeque, + iter: Option, + _phantom: PhantomData<&'b mut ()>, +} + +impl<'b, I: Iterator> Default for Cache<'b, I> { + fn default() -> Self { + Self { + window: Default::default(), + tail: Default::default(), + iter: Default::default(), + _phantom: PhantomData, + } + } +} + +impl<'b, I: Iterator> Cache<'b, I> +where + I::Item: Fetch<'b>, +{ + #[allow(dead_code)] + pub fn pop_id(&mut self) -> Option { + while self.tail.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + self.tail.pop_front() + } + #[allow(dead_code)] + pub fn pop_item_if(&mut self, predicate: impl FnOnce(&I::Item) -> bool) -> Option { + while self.window.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + self.window.pop_front_if(move |x| predicate(x)) + } + #[allow(dead_code)] + pub fn pop_item(&mut self) -> Option { + while self.window.is_empty() + && let Some(iter) = self.iter.as_mut() + && let Some(e) = iter.next() + { + for id in e.fetch() { + self.tail.push_back(id); + } + self.window.push_back(e); + } + self.window.pop_front() + } +} + +pub struct PostgresReadStream<'b, O: Opaque, I: Iterator> { + #[cfg(any(feature = "pg17", feature = "pg18"))] + raw: *mut pgrx::pg_sys::ReadStream, + // Because of `Box`'s special alias rules, `Box` cannot be used here. + cache: NonNull>, + _phantom: PhantomData O>, +} + +pub struct PostgresReadStreamGuards { + #[cfg(any(feature = "pg17", feature = "pg18"))] + raw: *mut pgrx::pg_sys::ReadStream, + list: L, + _phantom: PhantomData (O, I)>, +} + +impl Iterator for PostgresReadStreamGuards +where + L: Iterator, +{ + type Item = PostgresBufferReadGuard; + + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next(&mut self) -> Option { + unimplemented!() + } + + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next(&mut self) -> Option { + let _id = self.list.next()?; + unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_SHARE, BufferGetPage, LockBuffer, read_stream_next_buffer, + }; + let buf = read_stream_next_buffer(self.raw, core::ptr::null_mut()); + LockBuffer(buf, BUFFER_LOCK_SHARE as _); + let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); + Some(PostgresBufferReadGuard { + buf, + page, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + }) + } + } + + fn size_hint(&self) -> (usize, Option) { + self.list.size_hint() + } +} + +impl ExactSizeIterator for PostgresReadStreamGuards where + L: ExactSizeIterator +{ +} + +#[cfg(any(feature = "pg17", feature = "pg18"))] +impl<'b, O: Opaque, I: Iterator> PostgresReadStream<'b, O, I> +where + I::Item: Fetch<'b>, +{ + fn read>(&mut self, fetch: L) -> PostgresReadStreamGuards { + PostgresReadStreamGuards { + raw: self.raw, + list: fetch, + _phantom: PhantomData, + } + } +} + +impl<'b, O: Opaque, I: Iterator> ReadStream<'b> for PostgresReadStream<'b, O, I> +where + I::Item: Fetch<'b>, +{ + type Relation = PostgresRelation; + + type Guards = PostgresReadStreamGuards>::Iter>; + + type Item = I::Item; + + type Inner = + Chain, Flatten>>; + + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next(&mut self) -> Option<(I::Item, Self::Guards)> { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + fn next_if bool>( + &mut self, + _predicate: P, + ) -> Option<(I::Item, Self::Guards)> { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next(&mut self) -> Option<(I::Item, Self::Guards)> { + if let Some(e) = unsafe { self.cache.as_mut().pop_item() } { + let list = self.read(e.fetch()); + Some((e, list)) + } else { + None + } + } + + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn next_if bool>( + &mut self, + predicate: P, + ) -> Option<(I::Item, Self::Guards)> { + if let Some(e) = unsafe { self.cache.as_mut().pop_item_if(predicate) } { + let list = self.read(e.fetch()); + Some((e, list)) + } else { + None + } + } + + fn into_inner(mut self) -> Self::Inner { + let cache = unsafe { std::mem::take(self.cache.as_mut()) }; + cache + .window + .into_iter() + .chain(cache.iter.into_iter().flatten()) + } +} + +impl<'b, O: Opaque, I: Iterator> Drop for PostgresReadStream<'b, O, I> { + fn drop(&mut self) { + unsafe { + let _ = std::mem::take(self.cache.as_mut()); + #[cfg(any(feature = "pg17", feature = "pg18"))] + if !std::thread::panicking() && pgrx::pg_sys::IsTransactionState() { + pgrx::pg_sys::read_stream_end(self.raw); + } + let _ = box_from_non_null(self.cache); + } + } +} + +impl RelationReadStreamTypes for PostgresRelation { + type ReadStream<'b, I: Iterator> + = PostgresReadStream<'b, O, I> + where + I::Item: Fetch<'b>; +} + +impl RelationReadStream for PostgresRelation { + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))] + fn read_stream<'b, I: Iterator>(&'b self, _iter: I, _hints: Hints) -> Self::ReadStream<'b, I> + where + I::Item: Fetch<'b>, + { + panic!("read_stream is not supported on PostgreSQL versions earlier than 17."); + } + + #[cfg(any(feature = "pg17", feature = "pg18"))] + fn read_stream<'b, I: Iterator>(&'b self, iter: I, hints: Hints) -> Self::ReadStream<'b, I> + where + I::Item: Fetch<'b>, + { + #[pgrx::pg_guard] + unsafe extern "C-unwind" fn callback<'b, I: Iterator>( + _stream: *mut pgrx::pg_sys::ReadStream, + callback_private_data: *mut core::ffi::c_void, + _per_buffer_data: *mut core::ffi::c_void, + ) -> pgrx::pg_sys::BlockNumber + where + I::Item: Fetch<'b>, + { + unsafe { + use pgrx::pg_sys::InvalidBlockNumber; + let inner = callback_private_data.cast::>(); + (*inner).pop_id().unwrap_or(InvalidBlockNumber) + } + } + let cache = box_into_non_null(Box::new(Cache { + window: VecDeque::new(), + tail: VecDeque::new(), + iter: Some(iter), + _phantom: PhantomData, + })); + let raw = unsafe { + let mut flags = pgrx::pg_sys::READ_STREAM_DEFAULT; + if hints.full { + flags |= pgrx::pg_sys::READ_STREAM_FULL; + } + #[cfg(feature = "pg18")] + if hints.batch { + flags |= pgrx::pg_sys::READ_STREAM_USE_BATCHING; + } + pgrx::pg_sys::read_stream_begin_relation( + flags as i32, + core::ptr::null_mut(), + self.raw, + pgrx::pg_sys::ForkNumber::MAIN_FORKNUM, + Some(callback::), + cache.as_ptr().cast(), + 0, + ) + }; + PostgresReadStream { + raw, + cache, + _phantom: PhantomData, + } + } +} + +#[inline(always)] +fn lp_flags(x: pgrx::pg_sys::ItemIdData) -> u32 { + let x: u32 = unsafe { std::mem::transmute(x) }; + (x >> 15) & 0b11 +} + +// Emulate unstable library feature `box_vec_non_null`. +// See https://github.com/rust-lang/rust/issues/130364. + +#[allow(dead_code)] +#[must_use] +fn box_into_non_null(b: Box) -> NonNull { + unsafe { NonNull::new_unchecked(Box::into_raw(b)) } +} + +#[must_use] +unsafe fn box_from_non_null(ptr: NonNull) -> Box { + unsafe { Box::from_raw(ptr.as_ptr()) } +} diff --git a/src/index/storage/buffered.rs b/src/index/storage/buffered.rs new file mode 100644 index 00000000..3748dd52 --- /dev/null +++ b/src/index/storage/buffered.rs @@ -0,0 +1,135 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::storage::{ + PostgresBufferReadGuard, PostgresBufferWriteGuard, PostgresPage, PostgresRelation, +}; +use index::relation::{ + Opaque, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; + +#[derive(Debug)] +pub struct BufferedPostgresRelation { + postgres: PostgresRelation, + list: std::cell::RefCell>, +} + +impl BufferedPostgresRelation { + pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { + Self { + postgres: unsafe { PostgresRelation::new(raw) }, + list: std::cell::RefCell::new(arrayvec::ArrayVec::new()), + } + } + pub fn release(&self) -> bool { + let list = self.list.borrow(); + !list.is_empty() + } +} + +impl Drop for BufferedPostgresRelation { + fn drop(&mut self) { + let free = self.list.get_mut(); + #[cfg(debug_assertions)] + assert!(free.is_empty()); + for &mut buf in free { + unsafe { + pgrx::pg_sys::ReleaseBuffer(buf); + } + } + } +} + +impl Relation for BufferedPostgresRelation { + type Page = PostgresPage; +} + +impl RelationReadTypes for BufferedPostgresRelation { + type ReadGuard<'a> = PostgresBufferReadGuard; +} + +impl RelationWriteTypes for BufferedPostgresRelation { + type WriteGuard<'a> = PostgresBufferWriteGuard; +} + +impl RelationRead for BufferedPostgresRelation { + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + self.postgres.read(id) + } +} + +impl RelationWrite for BufferedPostgresRelation { + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.postgres.write(id, tracking_freespace) + } + fn search(&self, freespace: usize) -> Option> { + self.postgres.search(freespace) + } + #[cfg(any(feature = "pg14", feature = "pg15"))] + fn extend(&self, opaque: O, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.postgres.extend(opaque, tracking_freespace) + } + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + fn extend(&self, opaque: O, tracking_freespace: bool) -> Self::WriteGuard<'_> { + unsafe { + use pgrx::pg_sys::{ + BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, GenericXLogRegisterBuffer, + GenericXLogStart, LockBuffer, + }; + use std::mem::MaybeUninit; + use std::ptr::NonNull; + let buf = { + let mut list = self.list.borrow_mut(); + if list.is_empty() { + use pgrx::pg_sys::{BufferManagerRelation, ExtendBufferedRelBy, ForkNumber}; + let bmr = BufferManagerRelation { + rel: self.postgres.raw, + smgr: std::ptr::null_mut(), + relpersistence: 0, + }; + let mut len = 0_u32; + _ = ExtendBufferedRelBy( + bmr, + ForkNumber::MAIN_FORKNUM, + std::ptr::null_mut(), + 0, + list.capacity() as _, + list.as_mut_ptr(), + &raw mut len, + ); + assert!(1 <= len && len <= list.capacity() as _); + list.set_len(len as _); + list.reverse(); + } + list.pop().expect("number of allocated pages is zero") + }; + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); + let state = GenericXLogStart(self.postgres.raw); + let mut page = NonNull::new( + GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) + .cast::>>(), + ) + .expect("failed to get page"); + crate::index::storage::page_init(page.as_mut().as_mut_ptr(), opaque); + PostgresBufferWriteGuard { + raw: self.postgres.raw, + buf, + page: page.cast(), + state, + id: pgrx::pg_sys::BufferGetBlockNumber(buf), + tracking_freespace, + } + } + } +} diff --git a/src/index/traverse.rs b/src/index/traverse.rs new file mode 100644 index 00000000..c8d35699 --- /dev/null +++ b/src/index/traverse.rs @@ -0,0 +1,127 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use pgrx::pg_sys::{Datum, ItemPointerData}; + +pub trait Tuple { + fn id(&mut self) -> ItemPointerData; + fn build(&mut self) -> (*const Datum, *const bool); +} + +pub trait Traverse { + fn next(&mut self, tuple: T); +} + +impl Traverse for F { + fn next(&mut self, mut tuple: T) { + self(&mut tuple); + } +} + +pub trait Traverser { + fn traverse(&self, progress: bool, traverse: T); +} + +#[derive(Debug, Clone)] +pub struct HeapTraverser { + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + scan: *mut pgrx::pg_sys::TableScanDescData, +} + +impl HeapTraverser { + pub unsafe fn new( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + scan: *mut pgrx::pg_sys::TableScanDescData, + ) -> Self { + Self { + heap_relation, + index_relation, + index_info, + scan, + } + } +} + +impl Drop for HeapTraverser { + fn drop(&mut self) {} +} + +impl Traverser for HeapTraverser { + fn traverse(&self, progress: bool, mut traverse: T) { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let table_am = (*self.heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_build_range_scan = (*table_am) + .index_build_range_scan + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + index_build_range_scan( + self.heap_relation, + self.index_relation, + self.index_info, + true, + false, + progress, + 0, + pgrx::pg_sys::InvalidBlockNumber, + Some(callback::), + (&raw mut traverse).cast(), + self.scan, + ) + }); + } + } +} + +struct HeapTuple { + id: ItemPointerData, + values: *const Datum, + is_nulls: *const bool, +} + +impl Tuple for HeapTuple { + fn id(&mut self) -> ItemPointerData { + self.id + } + + fn build(&mut self) -> (*const Datum, *const bool) { + (self.values, self.is_nulls) + } +} + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn callback( + _index_relation: pgrx::pg_sys::Relation, + ctid: pgrx::pg_sys::ItemPointer, + values: *mut Datum, + is_null: *mut bool, + _tuple_is_alive: bool, + state: *mut core::ffi::c_void, +) { + let state = unsafe { &mut *state.cast::() }; + + state.next(HeapTuple { + id: unsafe { *ctid }, + values: values.cast_const(), + is_nulls: is_null.cast_const(), + }); +} diff --git a/src/index/utils.rs b/src/index/utils.rs deleted file mode 100644 index a5d85a3f..00000000 --- a/src/index/utils.rs +++ /dev/null @@ -1,20 +0,0 @@ -use base::search::*; - -pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData { - let value = pointer.as_u64(); - pgrx::pg_sys::ItemPointerData { - ip_blkid: pgrx::pg_sys::BlockIdData { - bi_hi: ((value >> 32) & 0xffff) as u16, - bi_lo: ((value >> 16) & 0xffff) as u16, - }, - ip_posid: (value & 0xffff) as u16, - } -} - -pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer { - let mut value = 0; - value |= (ctid.ip_blkid.bi_hi as u64) << 32; - value |= (ctid.ip_blkid.bi_lo as u64) << 16; - value |= ctid.ip_posid as u64; - Pointer::new(value) -} diff --git a/src/index/vchordg/am/am_build.rs b/src/index/vchordg/am/am_build.rs new file mode 100644 index 00000000..a5ab6361 --- /dev/null +++ b/src/index/vchordg/am/am_build.rs @@ -0,0 +1,664 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::typmod::Typmod; +use crate::index::storage::PostgresRelation; +use crate::index::traverse::{HeapTraverser, Traverser}; +use crate::index::vchordg::am::{Reloption, ctid_to_key, kv_to_pointer}; +use crate::index::vchordg::opclass::opfamily; +use crate::index::vchordg::types::VchordgIndexingOptions; +use std::ffi::CStr; +use std::marker::PhantomData; +use vchordg::types::*; + +#[derive(Debug, Clone, Copy)] +#[repr(u16)] +pub enum BuildPhaseCode { + Initializing = 0, + Inserting = 1, +} + +pub struct BuildPhase(BuildPhaseCode, u16); + +impl BuildPhase { + pub const fn new(code: BuildPhaseCode, k: u16) -> Option { + match (code, k) { + (BuildPhaseCode::Initializing, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Inserting, 0) => Some(BuildPhase(code, k)), + _ => None, + } + } + pub const fn name(self) -> &'static CStr { + match self { + BuildPhase(BuildPhaseCode::Initializing, k) => { + static RAW: [&CStr; 1] = [c"initializing"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Inserting, k) => { + static RAW: [&CStr; 1] = [c"inserting tuples from table to index"]; + RAW[k as usize] + } + } + } + pub const fn from_code(code: BuildPhaseCode) -> Self { + Self(code, 0) + } + pub const fn from_value(value: u32) -> Option { + const INITIALIZING: u16 = BuildPhaseCode::Initializing as _; + const INSERTING: u16 = BuildPhaseCode::Inserting as _; + let k = value as u16; + match (value >> 16) as u16 { + INITIALIZING => Self::new(BuildPhaseCode::Initializing, k), + INSERTING => Self::new(BuildPhaseCode::Inserting, k), + _ => None, + } + } + pub const fn into_value(self) -> u32 { + (self.0 as u32) << 16 | (self.1 as u32) + } +} + +#[pgrx::pg_guard] +pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { + if let Ok(x) = u32::try_from(x.wrapping_sub(1)) { + if let Some(x) = BuildPhase::from_value(x) { + x.name().as_ptr().cast_mut() + } else { + std::ptr::null_mut() + } + } else { + std::ptr::null_mut() + } +} + +#[derive(Debug, Clone)] +struct PostgresReporter { + _phantom: PhantomData<*mut ()>, +} + +impl PostgresReporter { + fn phase(&self, phase: BuildPhase) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, + (phase.into_value() as i64) + 1, + ); + } + } + fn tuples_total(&self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } + fn tuples_done(&self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuild( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, +) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; + let (vector_options, vchordg_options) = unsafe { options(index_relation) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if let Err(errors) = Validate::validate(&vchordg_options) { + pgrx::error!("error while validating options: {}", errors); + } + if vector_options.d != DistanceKind::L2S + && (vchordg_options.index.alpha != [1.0] && vchordg_options.index.alpha != [1.0, 1.2]) + { + let warnings = "alpha not equal to `1.0` are only applicable to l2 and cosine distance."; + pgrx::warning!("warning while validating options: {warnings}"); + } + let index = unsafe { PostgresRelation::new(index_relation) }; + let reporter = PostgresReporter { + _phantom: PhantomData, + }; + crate::index::vchordg::dispatch::build(vector_options, vchordg_options.index, &index); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + let cached = vchordg_cached::VchordgCached::_0 {}.serialize(); + if let Some(leader) = unsafe { + VchordgLeader::enter( + c"vchordg_parallel_build_main", + heap_relation, + index_relation, + (*index_info).ii_Concurrent, + &cached, + ) + } { + drop(cached); + unsafe { + parallel_build( + index_relation, + heap_relation, + index_info, + leader.tablescandesc, + leader.vchordgshared, + leader.vchordgcached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + ); + leader.wait(); + let nparticipants = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*leader.vchordgshared).mutex); + if (*leader.vchordgshared).workers_done == nparticipants { + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*leader.vchordgshared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*leader.vchordgshared).condvar_workers_done, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + reporter.tuples_done((*leader.vchordgshared).indtuples); + reporter.tuples_total((*leader.vchordgshared).indtuples); + pgrx::pg_sys::ConditionVariableCancelSleep(); + } + } else { + unsafe { + let indtuples = sequential_build( + index_relation, + heap_relation, + index_info, + &cached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + ); + reporter.tuples_total(indtuples); + } + } + unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } +} + +struct VchordgShared { + /* immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* locking */ + mutex: pgrx::pg_sys::slock_t, + condvar_workers_done: pgrx::pg_sys::ConditionVariable, + + /* mutable state */ + workers_done: i32, + indtuples: u64, +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct VchordgLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipants: i32, + snapshot: pgrx::pg_sys::Snapshot, + vchordgshared: *mut VchordgShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordgcached: *const u8, +} + +impl VchordgLeader { + pub unsafe fn enter( + main: &'static CStr, + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + isconcurrent: bool, + cached: &[u8], + ) -> Option { + unsafe fn compute_parallel_workers( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers( + (*heap_relation).rd_id, + (*index_relation).rd_id, + ) == 0 + { + return 0; + } + if !(*heap_relation).rd_options.is_null() { + let std_options = (*heap_relation) + .rd_options + .cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap_relation, index_relation) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext(c"vchord".as_ptr(), main.as_ptr(), request) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap_relation, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, 8 + cached.len()); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let vchordgshared = unsafe { + let vchordgshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordgshared.write(VchordgShared { + heaprelid: (*heap_relation).rd_id, + indexrelid: (*index_relation).rd_id, + isconcurrent, + condvar_workers_done: std::mem::zeroed(), + mutex: std::mem::zeroed(), + workers_done: 0, + indtuples: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordgshared).condvar_workers_done); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordgshared).mutex); + vchordgshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap_relation, tablescandesc, snapshot); + tablescandesc + }; + + let vchordgcached = unsafe { + let x = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + cached.len()).cast::(); + (x as *mut u64).write_unaligned(cached.len() as _); + std::ptr::copy(cached.as_ptr(), x.add(8), cached.len()); + x + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordgshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000003, vchordgcached.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nworkers_launched == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + Some(Self { + pcxt, + nparticipants: nworkers_launched + 1, + snapshot, + vchordgshared, + tablescandesc, + vchordgcached, + }) + } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } +} + +impl Drop for VchordgLeader { + fn drop(&mut self) { + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } + } +} + +#[pgrx::pg_guard] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vchordg_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let _ = rand::rng().reseed(); + let vchordgshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let vchordgcached = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000003, false) + .cast::() + .cast_const() + }; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordgshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + } else { + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; + let heap = unsafe { pgrx::pg_sys::table_open((*vchordgshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*vchordgshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*vchordgshared).isconcurrent; + } + + unsafe { + parallel_build( + index, + heap, + index_info, + tablescandesc, + vchordgshared, + vchordgcached, + |_| (), + ); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordgshared: *mut VchordgShared, + vchordgcached: *const u8, + mut callback: impl FnMut(u64), +) { + use vchordg_cached::VchordgCachedReader; + let cached = VchordgCachedReader::deserialize_ref(unsafe { + let bytes = (vchordgcached as *const u64).read_unaligned(); + std::slice::from_raw_parts(vchordgcached.add(8), bytes as _) + }); + + let index = unsafe { PostgresRelation::new(index_relation) }; + + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; + let opfamily = unsafe { opfamily(index_relation) }; + let traverser = unsafe { HeapTraverser::new(heap_relation, index_relation, index_info, scan) }; + match cached { + VchordgCachedReader::_0(_) => { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); + } + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordgshared).mutex); + (*vchordgshared).indtuples += 1; + indtuples = (*vchordgshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordgshared).mutex); + } + callback(indtuples); + } + }); + } + } + unsafe { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordgshared).mutex); + (*vchordgshared).workers_done += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordgshared).mutex); + pgrx::pg_sys::ConditionVariableSignal(&raw mut (*vchordgshared).condvar_workers_done); + } +} + +unsafe fn sequential_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + vchordgcached: &[u8], + mut callback: impl FnMut(u64), +) -> u64 { + use vchordg_cached::VchordgCachedReader; + let cached = VchordgCachedReader::deserialize_ref(vchordgcached); + let index = unsafe { PostgresRelation::new(index_relation) }; + + let opfamily = unsafe { opfamily(index_relation) }; + let traverser = unsafe { + HeapTraverser::new( + heap_relation, + index_relation, + index_info, + std::ptr::null_mut(), + ) + }; + + let mut indtuples = 0; + match cached { + VchordgCachedReader::_0(_) => { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); + } + indtuples += 1; + callback(indtuples); + }); + } + } + indtuples +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { + pgrx::error!("Unlogged indexes are not supported."); +} + +unsafe fn options( + index_relation: pgrx::pg_sys::Relation, +) -> (VectorOptions, VchordgIndexingOptions) { + let att = unsafe { &mut *(*index_relation).rd_att }; + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + #[cfg(feature = "pg18")] + let atts = unsafe { + let ptr = att + .compact_attrs + .as_ptr() + .add(att.natts as _) + .cast::(); + std::slice::from_raw_parts(ptr, att.natts as _) + }; + if atts.is_empty() { + pgrx::error!("indexing on no columns is not supported"); + } + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); + } + // get dim + let typmod = Typmod::new(atts[0].atttypmod).unwrap(); + let dim = if let Some(dim) = typmod.dim() { + dim.get() + } else { + pgrx::error!( + "Dimensions type modifier of a vector column is needed for building the index." + ); + }; + // get v, d + let opfamily = unsafe { opfamily(index_relation) }; + let vector = VectorOptions { + dim, + v: opfamily.vector_kind(), + d: opfamily.distance_kind(), + }; + // get indexing options + let indexing_options = { + let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; + let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); + match toml::from_str::(&s) { + Ok(p) => p, + Err(e) => pgrx::error!("failed to parse options: {}", e), + } + }; + (vector, indexing_options) +} + +mod vchordg_cached { + pub type Tag = u64; + + use index::tuples::RefChecker; + use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + + #[repr(C, align(8))] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VchordgCachedHeader0 {} + + pub enum VchordgCached { + _0 {}, + } + + impl VchordgCached { + pub fn serialize(&self) -> Vec { + let mut buffer = Vec::new(); + match self { + VchordgCached::_0 {} => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + buffer[size_of::()..][..size_of::()] + .copy_from_slice(VchordgCachedHeader0 {}.as_bytes()); + } + } + buffer + } + } + + #[derive(Debug, Clone, Copy)] + pub enum VchordgCachedReader<'a> { + #[allow(dead_code)] + _0(VchordgCachedReader0<'a>), + } + + #[derive(Debug, Clone, Copy)] + pub struct VchordgCachedReader0<'a> { + #[allow(dead_code)] + header: &'a VchordgCachedHeader0, + } + + impl<'a> VchordgCachedReader<'a> { + pub fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = u64::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VchordgCachedHeader0 = checker.prefix(size_of::()); + Self::_0(VchordgCachedReader0 { header }) + } + _ => panic!("bad bytes"), + } + } + } +} diff --git a/src/index/vchordg/am/mod.rs b/src/index/vchordg/am/mod.rs new file mode 100644 index 00000000..39707100 --- /dev/null +++ b/src/index/vchordg/am/mod.rs @@ -0,0 +1,456 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod am_build; + +use crate::index::fetcher::*; +use crate::index::gucs; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordg::opclass::opfamily; +use crate::index::vchordg::scanners::*; +use crate::recorder::DefaultRecorder; +use pgrx::datum::Internal; +use pgrx::pg_sys::Datum; +use std::cell::LazyCell; +use std::ffi::CStr; +use std::num::NonZero; +use std::ops::DerefMut; +use std::ptr::NonNull; +use std::sync::OnceLock; + +#[repr(C)] +pub struct Reloption { + vl_len_: i32, + options: i32, + ef_search: i32, +} + +impl Reloption { + unsafe fn options<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { + unsafe { + if this.is_null() { + return default; + } + let count = (&raw const (*this).options).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); + CStr::from_ptr(ptr.cast()) + } + } + pub unsafe fn ef_search(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).ef_search + } + } +} + +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[ + pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"ef_search".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, ef_search) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, +]; + +static RELOPT_KIND: OnceLock = OnceLock::new(); + +pub fn init() { + RELOPT_KIND.get_or_init(|| { + let kind; + unsafe { + kind = pgrx::pg_sys::add_reloption_kind(); + pgrx::pg_sys::add_string_reloption( + kind as _, + c"options".as_ptr(), + c"Vector index options, represented as a TOML string.".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"ef_search".as_ptr(), + c"Search parameter `vchordg.ef_search`".as_ptr(), + 64, + 1, + 65535, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + } + kind + }); +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordg_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { + type T = pgrx::pg_sys::IndexAmRoutine; + unsafe { + let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; + index_am_routine.write(AM_HANDLER); + Internal::from(Some(Datum::from(index_am_routine))) + } +} + +const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { + let mut am_routine = unsafe { std::mem::zeroed::() }; + + am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + + am_routine.amsupport = 1; + am_routine.amcanorderbyop = true; + + #[cfg(any(feature = "pg17", feature = "pg18"))] + { + am_routine.amcanbuildparallel = true; + } + + // Index access methods that set `amoptionalkey` to `false` + // must index all tuples, even if the first column is `NULL`. + // However, PostgreSQL does not generate a path if there is no + // index clauses, even if there is a `ORDER BY` clause. + // So we have to set it to `true` and set costs of every path + // for vector index scans without `ORDER BY` clauses a large number + // and throw errors if someone really wants such a path. + am_routine.amoptionalkey = true; + + am_routine.amvalidate = Some(amvalidate); + am_routine.amoptions = Some(amoptions); + am_routine.amcostestimate = Some(amcostestimate); + + am_routine.ambuildphasename = Some(am_build::ambuildphasename); + am_routine.ambuild = Some(am_build::ambuild); + am_routine.ambuildempty = Some(am_build::ambuildempty); + am_routine.aminsert = Some(aminsert); + am_routine.ambulkdelete = Some(ambulkdelete); + am_routine.amvacuumcleanup = Some(amvacuumcleanup); + + am_routine.ambeginscan = Some(ambeginscan); + am_routine.amrescan = Some(amrescan); + am_routine.amgettuple = Some(amgettuple); + am_routine.amendscan = Some(amendscan); + + am_routine.amparallelvacuumoptions = pgrx::pg_sys::VACUUM_OPTION_PARALLEL_BULKDEL as u8 + | pgrx::pg_sys::VACUUM_OPTION_PARALLEL_CLEANUP as u8; + + am_routine +}; + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amoptions( + reloptions: Datum, + validate: bool, +) -> *mut pgrx::pg_sys::bytea { + let relopt_kind = RELOPT_KIND.get().copied().expect("init is not called"); + let rdopts = unsafe { + pgrx::pg_sys::build_reloptions( + reloptions, + validate, + relopt_kind, + size_of::(), + TABLE.as_ptr(), + TABLE.len() as _, + ) + }; + rdopts as *mut pgrx::pg_sys::bytea +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amcostestimate( + _root: *mut pgrx::pg_sys::PlannerInfo, + _path: *mut pgrx::pg_sys::IndexPath, + _loop_count: f64, + index_startup_cost: *mut pgrx::pg_sys::Cost, + index_total_cost: *mut pgrx::pg_sys::Cost, + index_selectivity: *mut pgrx::pg_sys::Selectivity, + index_correlation: *mut f64, + index_pages: *mut f64, +) { + unsafe { + use pgrx::pg_sys::disable_cost; + if !gucs::vchordg_enable_scan() { + *index_startup_cost = disable_cost; + *index_total_cost = disable_cost; + *index_selectivity = 0.0; + *index_correlation = 0.0; + *index_pages = 1.0; + return; + } + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = 1.0; + *index_correlation = 0.0; + *index_pages = 1.0; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_unchanged: bool, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { heap_tid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + crate::index::vchordg::dispatch::insert(opfamily, &index, payload, vector); + } + } + false +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambulkdelete( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, + callback: pgrx::pg_sys::IndexBulkDeleteCallback, + callback_state: *mut std::os::raw::c_void, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + let callback = callback.expect("null function pointer"); + let callback = |pointer: NonZero| { + let (key, _) = pointer_to_kv(pointer); + let mut ctid = key_to_ctid(key); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) + } + }; + crate::index::vchordg::dispatch::bulkdelete(opfamily, &index, check, callback); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvacuumcleanup( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + crate::index::vchordg::dispatch::maintain(opfamily, &index, check); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambeginscan( + index_relation: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; + let scanner: Scanner = Scanner { + hack: None, + scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), + bump: Box::new(bumpalo::Bump::new()), + }; + unsafe { + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + use crate::index::vchordg::opclass::Opfamily; + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); + } + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); + } + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); + } + let scanner = &mut *(*scan).opaque.cast::(); + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); + let opfamily = opfamily((*scan).indexRelation); + let index = PostgresRelation::new((*scan).indexRelation); + let options = SearchOptions { + ef_search: gucs::vchordg_ef_search((*scan).indexRelation), + beam_search: gucs::vchordg_beam_search(), + max_scan_tuples: gucs::vchordg_max_scan_tuples(), + io_search: gucs::vchordg_io_search(), + io_rerank: gucs::vchordg_io_rerank(), + }; + let fetcher = { + let hack = scanner.hack; + LazyCell::new(move || { + HeapFetcher::new( + (*scan).indexRelation, + (*scan).heapRelation, + (*scan).xs_snapshot, + (*scan).xs_heapfetch, + if let Some(hack) = hack { + hack.as_ptr() + } else { + std::ptr::null_mut() + }, + ) + }) + }; + // Query recorde is disable for vchordg indexes for now. + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: (*(*scan).indexRelation).rd_id.to_u32(), + }; + // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. + let bump = scanner.bump.as_ref(); + scanner.scanning = match opfamily { + Opfamily::VectorL2 + | Opfamily::VectorCosine + | Opfamily::VectorIp + | Opfamily::HalfvecL2 + | Opfamily::HalfvecCosine + | Opfamily::HalfvecIp + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4Ip => { + let mut builder = DefaultBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump, recorder) + })) + } + }; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + if let Some((_, key, recheck)) = scanner.scanning.deref_mut().next() { + unsafe { + (*scan).xs_heaptid = key_to_ctid(key); + (*scan).xs_recheck = recheck; + (*scan).xs_recheckorderby = false; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + let scanner = unsafe { &mut *(*scan).opaque.cast::() }; + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); +} + +type Iter = Box>; + +pub struct Scanner { + pub hack: Option>, + scanning: LazyCell Iter>>, + bump: Box, +} diff --git a/src/index/vchordg/dispatch.rs b/src/index/vchordg/dispatch.rs new file mode 100644 index 00000000..02115da3 --- /dev/null +++ b/src/index/vchordg/dispatch.rs @@ -0,0 +1,388 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::vchordg::opclass::Opfamily; +use index::fetch::Fetch; +use index::prefetcher::*; +use index::relation::{ + Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, +}; +use index_accessor::{Dot, L2S}; +use simd::f16; +use std::num::NonZero; +use vchordg::operator::Op; +use vchordg::types::*; +use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::{VectBorrowed, VectOwned}; + +pub fn prewarm(opfamily: Opfamily, index: &R) -> String +where + R: RelationRead, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op, L2S>>(index) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op, Dot>>(index) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op, L2S>>(index) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op, Dot>>(index) + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op>(index) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op>(index) + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::prewarm::<_, Op>(index) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::prewarm::<_, Op>(index) + } + } +} + +pub fn bulkdelete( + opfamily: Opfamily, + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::bulkdelete::<_, Op>(index, &check, &callback); + } + } +} + +pub fn maintain(opfamily: Opfamily, index: &R, check: impl Fn()) +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::maintain::<_, Op, L2S>>(index, &check); + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::maintain::<_, Op, Dot>>(index, &check); + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::maintain::<_, Op, L2S>>(index, &check); + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::maintain::<_, Op, Dot>>(index, &check); + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::maintain::<_, Op>(index, &check); + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::maintain::<_, Op>(index, &check); + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::maintain::<_, Op>(index, &check); + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::maintain::<_, Op>(index, &check); + } + } +} + +pub fn build(vector_options: VectorOptions, vchordg_options: VchordgIndexOptions, index: &R) +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (vector_options.v, vector_options.d) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordg::build::<_, Op, L2S>>(vector_options, vchordg_options, index) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordg::build::<_, Op, Dot>>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordg::build::<_, Op>(vector_options, vchordg_options, index) + } + } +} + +pub fn insert(opfamily: Opfamily, index: &R, payload: NonZero, vector: OwnedVector) +where + R: RelationRead + RelationWrite + RelationReadStream, + R::Page: Page, +{ + let bump = bumpalo::Bump::new(); + let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; + match (vector, opfamily.distance_kind()) { + (OwnedVector::Vecf32(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, L2S>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf32(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, Dot>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf16(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, L2S>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Vecf16(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(unprojected.as_borrowed()); + vchordg::insert::<_, Op, Dot>>( + index, + projected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq8(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq8(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq4(unprojected), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + (OwnedVector::Rabitq4(unprojected), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + vchordg::insert::<_, Op>( + index, + unprojected.as_borrowed(), + payload, + &bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ) + } + } +} + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + let input = self.slice(); + VectOwned::new(rotate(input)) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + use simd::Floating; + let input = f16::vector_to_f32(self.slice()); + VectOwned::new(f16::vector_from_f32(&rotate(&input))) + } +} + +#[derive(Debug)] +pub struct MakePlainPrefetcher<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakePlainPrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead> PrefetcherSequenceFamily<'b, R> for MakePlainPrefetcher<'b, R> { + type P + = PlainPrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + PlainPrefetcher::new(self.index, seq) + } + + fn is_not_plain(&self) -> bool { + false + } +} + +#[derive(Debug)] +pub struct MakeSimplePrefetcher<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakeSimplePrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'b, R> + for MakeSimplePrefetcher<'b, R> +{ + type P + = SimplePrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + SimplePrefetcher::new(self.index, seq) + } + + fn is_not_plain(&self) -> bool { + true + } +} + +#[derive(Debug)] +pub struct MakeStreamPrefetcher<'b, R> { + pub index: &'b R, + pub hints: Hints, +} + +impl<'b, R> Clone for MakeStreamPrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { + index: self.index, + hints: self.hints, + } + } +} + +impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> + for MakeStreamPrefetcher<'b, R> +{ + type P + = StreamPrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + StreamPrefetcher::new(self.index, seq, self.hints) + } + + fn is_not_plain(&self) -> bool { + true + } +} diff --git a/src/index/vchordg/mod.rs b/src/index/vchordg/mod.rs new file mode 100644 index 00000000..b01977c7 --- /dev/null +++ b/src/index/vchordg/mod.rs @@ -0,0 +1,19 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod am; +pub mod dispatch; +pub mod opclass; +mod scanners; +pub mod types; diff --git a/src/index/vchordg/opclass.rs b/src/index/vchordg/opclass.rs new file mode 100644 index 00000000..2132ae22 --- /dev/null +++ b/src/index/vchordg/opclass.rs @@ -0,0 +1,242 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use crate::index::opclass::Sphere; +use distance::Distance; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use pgrx::pg_sys::Datum; +use std::num::NonZero; +use vchordg::types::*; +use vector::VectorBorrowed; + +#[derive(Debug, Clone, Copy)] +pub enum Opfamily { + VectorL2, + VectorCosine, + VectorIp, + HalfvecL2, + HalfvecCosine, + HalfvecIp, + Rabitq8L2, + Rabitq8Cosine, + Rabitq8Ip, + Rabitq4L2, + Rabitq4Cosine, + Rabitq4Ip, +} + +impl Opfamily { + fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { + use BorrowedVector as B; + use OwnedVector as O; + match (self, vector) { + (Self::VectorL2, B::Vecf32(x)) => O::Vecf32(x.own()), + (Self::VectorL2, _) => unreachable!(), + (Self::VectorCosine, B::Vecf32(x)) => O::Vecf32(x.function_normalize()), + (Self::VectorCosine, _) => unreachable!(), + (Self::VectorIp, B::Vecf32(x)) => O::Vecf32(x.own()), + (Self::VectorIp, _) => unreachable!(), + (Self::HalfvecL2, B::Vecf16(x)) => O::Vecf16(x.own()), + (Self::HalfvecL2, _) => unreachable!(), + (Self::HalfvecCosine, B::Vecf16(x)) => O::Vecf16(x.function_normalize()), + (Self::HalfvecCosine, _) => unreachable!(), + (Self::HalfvecIp, B::Vecf16(x)) => O::Vecf16(x.own()), + (Self::HalfvecIp, _) => unreachable!(), + (Self::Rabitq8L2, B::Rabitq8(x)) => O::Rabitq8(x.own()), + (Self::Rabitq8L2, _) => unreachable!(), + (Self::Rabitq8Cosine, B::Rabitq8(x)) => O::Rabitq8(x.function_normalize()), + (Self::Rabitq8Cosine, _) => unreachable!(), + (Self::Rabitq8Ip, B::Rabitq8(x)) => O::Rabitq8(x.own()), + (Self::Rabitq8Ip, _) => unreachable!(), + (Self::Rabitq4L2, B::Rabitq4(x)) => O::Rabitq4(x.own()), + (Self::Rabitq4L2, _) => unreachable!(), + (Self::Rabitq4Cosine, B::Rabitq4(x)) => O::Rabitq4(x.function_normalize()), + (Self::Rabitq4Cosine, _) => unreachable!(), + (Self::Rabitq4Ip, B::Rabitq4(x)) => O::Rabitq4(x.own()), + (Self::Rabitq4Ip, _) => unreachable!(), + } + } + pub unsafe fn store(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let store = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] + } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] + } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), 0)] + } + }; + Some(store) + } + pub unsafe fn input_sphere(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let attno_1 = NonZero::new(1_usize).unwrap(); + let attno_2 = NonZero::new(2_usize).unwrap(); + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } + }; + let radius = tuple.get_by_index::(attno_2).unwrap()?; + Some(Sphere { center, radius }) + } + pub unsafe fn input_vector(self, datum: Datum) -> Option { + if datum.is_null() { + return None; + } + let vector = match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } + }; + Some(vector) + } + pub fn output(self, x: Distance) -> f32 { + match self { + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => x.to_f32() * 0.5, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + x.to_f32().sqrt() + } + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip | Self::Rabitq4Ip => x.to_f32(), + } + } + pub const fn distance_kind(self) -> DistanceKind { + match self { + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + DistanceKind::L2S + } + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => DistanceKind::L2S, + Self::VectorIp | Self::HalfvecIp | Self::Rabitq8Ip | Self::Rabitq4Ip => { + DistanceKind::Dot + } + } + } + pub const fn vector_kind(self) -> VectorKind { + match self { + Self::VectorL2 | Self::VectorCosine | Self::VectorIp => VectorKind::Vecf32, + Self::HalfvecL2 | Self::HalfvecCosine | Self::HalfvecIp => VectorKind::Vecf16, + Self::Rabitq8L2 | Self::Rabitq8Cosine | Self::Rabitq8Ip => VectorKind::Rabitq8, + Self::Rabitq4L2 | Self::Rabitq4Cosine | Self::Rabitq4Ip => VectorKind::Rabitq4, + } + } +} + +pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) + }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null return value"); + + let result = match result_string.as_str() { + "vchordg_vector_l2_ops" => Opfamily::VectorL2, + "vchordg_vector_ip_ops" => Opfamily::VectorIp, + "vchordg_vector_cosine_ops" => Opfamily::VectorCosine, + "vchordg_halfvec_l2_ops" => Opfamily::HalfvecL2, + "vchordg_halfvec_ip_ops" => Opfamily::HalfvecIp, + "vchordg_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordg_rabitq8_l2_ops" => Opfamily::Rabitq8L2, + "vchordg_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, + "vchordg_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, + "vchordg_rabitq4_l2_ops" => Opfamily::Rabitq4L2, + "vchordg_rabitq4_ip_ops" => Opfamily::Rabitq4Ip, + "vchordg_rabitq4_cosine_ops" => Opfamily::Rabitq4Cosine, + _ => pgrx::error!("unknown operator class"), + }; + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + result +} diff --git a/src/index/vchordg/scanners/default.rs b/src/index/vchordg/scanners/default.rs new file mode 100644 index 00000000..22fd5b33 --- /dev/null +++ b/src/index/vchordg/scanners/default.rs @@ -0,0 +1,943 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::{Fetcher, pointer_to_kv}; +use crate::index::opclass::Sphere; +use crate::index::scanners::{Io, SearchBuilder}; +use crate::index::vchordg::dispatch::*; +use crate::index::vchordg::opclass::Opfamily; +use crate::index::vchordg::scanners::SearchOptions; +use crate::recorder::{Recorder, text}; +use distance::Distance; +use index::bump::Bump; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::{Dot, L2S}; +use simd::f16; +use std::num::NonZero; +use vchordg::operator::{self}; +use vchordg::search; +use vchordg::types::{DistanceKind, OwnedVector, VectorKind}; +use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; +use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; +use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; + +pub struct DefaultBuilder { + opfamily: Opfamily, + orderbys: Vec>, + spheres: Vec>>, +} + +impl SearchBuilder for DefaultBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + + type Opaque = vchordg::Opaque; + + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::VectorCosine + | Opfamily::VectorL2 + | Opfamily::VectorIp + | Opfamily::HalfvecCosine + | Opfamily::HalfvecL2 + | Opfamily::HalfvecIp + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Ip + )); + Self { + opfamily, + orderbys: Vec::new(), + spheres: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 1 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vector(x)) }; + self.orderbys.push(x); + } + 2 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_sphere(x)) }; + self.spheres.push(x); + } + _ => unreachable!(), + } + } + + fn build<'b, R>( + self, + index: &'b R, + options: SearchOptions, + _fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + recorder: impl Recorder, + ) -> Box + 'b> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, + { + let mut vector = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in self.orderbys.into_iter().flatten() { + if vector.is_none() { + vector = Some(orderby_vector); + } else { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for Sphere { center, radius } in self.spheres.into_iter().flatten() { + if vector.is_none() { + (vector, threshold) = (Some(center), Some(radius)); + } else { + recheck = true; + } + } + let opfamily = self.opfamily; + let Some(vector) = vector else { + return Box::new(std::iter::empty()) as Box>; + }; + let vertex_hints = Hints::default().full(true); + let vector_hints = Hints::default().full(true); + let make_vertex_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vertex_simple_prefetcher = MakeSimplePrefetcher { index }; + let make_vertex_stream_prefetcher = MakeStreamPrefetcher { + index, + hints: vertex_hints, + }; + let make_vector_plain_prefetcher = MakePlainPrefetcher { index }; + let make_vector_simple_prefetcher = MakeSimplePrefetcher { index }; + let make_vector_stream_prefetcher = MakeStreamPrefetcher { + index, + hints: vector_hints, + }; + let iter: Box)>> = + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { + VectBorrowed::new(bump.alloc_slice(vector.slice())) + } else { + unreachable!() + }; + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type Op = operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { + VectBorrowed::new(bump.alloc_slice(vector.slice())) + } else { + unreachable!() + }; + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + type Op = operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { + VectBorrowed::new(bump.alloc_slice(vector.slice())) + } else { + unreachable!() + }; + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type Op = operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { + VectBorrowed::new(bump.alloc_slice(vector.slice())) + } else { + unreachable!() + }; + let projected = { + let projected = RandomProject::project(unprojected); + VectBorrowed::new(bump.alloc_slice(projected.slice())) + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + projected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq8Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq8Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq4Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + type Op = operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + let vector = vector.as_borrowed(); + Rabitq4Borrowed::new( + vector.dim(), + vector.sum_of_x2(), + vector.norm_of_lattice(), + vector.sum_of_code(), + vector.sum_of_abs_x(), + bump.alloc_slice(vector.packed_code()), + ) + } else { + unreachable!() + }; + match (options.io_search, options.io_rerank) { + (Io::Plain, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Plain, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_plain_prefetcher, + ), + (Io::Simple, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Simple, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_simple_prefetcher, + ), + (Io::Stream, Io::Plain) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_plain_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Simple) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_simple_prefetcher, + make_vector_stream_prefetcher, + ), + (Io::Stream, Io::Stream) => search::<_, Op>( + index, + unprojected, + options.ef_search, + options.beam_search, + bump, + make_vertex_stream_prefetcher, + make_vector_stream_prefetcher, + ), + } + } + }; + let iter = if let Some(threshold) = threshold { + Box::new(iter.take_while(move |(distance, _)| distance.to_f32() < threshold)) + } else { + iter + }; + let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { + Box::new(iter.take(max_scan_tuples as _)) + } else { + iter + }; + if recorder.is_enabled() { + match &vector { + OwnedVector::Vecf32(v) => { + recorder.send(&text::vector_out(v.as_borrowed())); + } + OwnedVector::Vecf16(v) => { + recorder.send(&text::halfvec_out(v.as_borrowed())); + } + OwnedVector::Rabitq8(v) => { + recorder.send(&text::rabitq8_out(v.as_borrowed())); + } + OwnedVector::Rabitq4(v) => { + recorder.send(&text::rabitq4_out(v.as_borrowed())); + } + } + } + Box::new(iter.map(move |(distance, pointer)| { + let (key, _) = pointer_to_kv(pointer); + (opfamily.output(distance), key, recheck) + })) + } +} diff --git a/src/index/vchordg/scanners/mod.rs b/src/index/vchordg/scanners/mod.rs new file mode 100644 index 00000000..9809809d --- /dev/null +++ b/src/index/vchordg/scanners/mod.rs @@ -0,0 +1,26 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod default; + +pub use default::DefaultBuilder; + +#[derive(Debug)] +pub struct SearchOptions { + pub ef_search: u32, + pub beam_search: u32, + pub max_scan_tuples: Option, + pub io_search: crate::index::scanners::Io, + pub io_rerank: crate::index::scanners::Io, +} diff --git a/src/index/vchordg/types.rs b/src/index/vchordg/types.rs new file mode 100644 index 00000000..6be44ea3 --- /dev/null +++ b/src/index/vchordg/types.rs @@ -0,0 +1,25 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use serde::{Deserialize, Serialize}; +use validator::Validate; +use vchordg::types::VchordgIndexOptions; + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordgIndexingOptions { + #[serde(flatten)] + #[validate(nested)] + pub index: VchordgIndexOptions, +} diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs new file mode 100644 index 00000000..15a5d982 --- /dev/null +++ b/src/index/vchordrq/am/am_build.rs @@ -0,0 +1,1849 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::typmod::Typmod; +use crate::index::fetcher::*; +use crate::index::sample::{HeapSampler, Sample, Sampler, Tuple}; +use crate::index::storage::buffered::BufferedPostgresRelation; +use crate::index::storage::{PostgresPage, PostgresRelation}; +use crate::index::traverse::{HeapTraverser, Traverser}; +use crate::index::vchordrq::am::Reloption; +use crate::index::vchordrq::build::{Normalize, Normalized}; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::types::*; +use index::relation::{ + Page, PageGuard, Relation, RelationRead, RelationReadTypes, RelationWrite, RelationWriteTypes, +}; +use k_means::k_means_lookup; +use k_means::square::Square; +use simd::Floating; +use std::ffi::CStr; +use std::marker::PhantomData; +use std::num::NonZero; +use std::ops::Deref; +use vchordrq::types::*; +use vchordrq::{InsertChooser, MaintainChooser}; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; + +#[derive(Debug, Clone, Copy)] +#[repr(u16)] +pub enum BuildPhaseCode { + Initializing = 0, + DefaultBuild = 1, + InternalBuild = 2, + ExternalBuild = 3, + Build = 4, + Inserting = 5, + Compacting = 6, +} + +pub struct BuildPhase(BuildPhaseCode, u16); + +impl BuildPhase { + pub const fn new(code: BuildPhaseCode, k: u16) -> Option { + match (code, k) { + (BuildPhaseCode::Initializing, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::DefaultBuild, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::InternalBuild, 0..102) => Some(BuildPhase(code, k)), + (BuildPhaseCode::ExternalBuild, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Build, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Inserting, 0) => Some(BuildPhase(code, k)), + (BuildPhaseCode::Compacting, 0) => Some(BuildPhase(code, k)), + _ => None, + } + } + pub const fn name(self) -> &'static CStr { + match self { + BuildPhase(BuildPhaseCode::Initializing, k) => { + static RAW: [&CStr; 1] = [c"initializing"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::DefaultBuild, k) => { + static RAW: [&CStr; 1] = [c"initializing index, by default build"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::InternalBuild, k) => { + static RAWS: [&[&CStr]; 2] = [ + &[c"initializing index, by internal build"], + seq_macro::seq!( + N in 0..=100 { + &[#( + const { + let bytes = concat!("initializing index, by internal build (", N, " %)\0").as_bytes(); + if let Ok(cstr) = CStr::from_bytes_with_nul(bytes) { + cstr + } else { + unreachable!() + } + }, + )*] + } + ), + ]; + static RAW: [&CStr; 102] = { + let mut result = [c""; 102]; + let mut offset = 0_usize; + let mut i = 0_usize; + while i < RAWS.len() { + let mut j = 0_usize; + while j < RAWS[i].len() { + { + result[offset] = RAWS[i][j]; + offset += 1; + } + j += 1; + } + i += 1; + } + assert!(offset == result.len()); + result + }; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::ExternalBuild, k) => { + static RAW: [&CStr; 1] = [c"initializing index, by external build"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Build, k) => { + static RAW: [&CStr; 1] = [c"initializing index"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Inserting, k) => { + static RAW: [&CStr; 1] = [c"inserting tuples from table to index"]; + RAW[k as usize] + } + BuildPhase(BuildPhaseCode::Compacting, k) => { + static RAW: [&CStr; 1] = [c"compacting tuples in index"]; + RAW[k as usize] + } + } + } + pub const fn from_code(code: BuildPhaseCode) -> Self { + Self(code, 0) + } + pub const fn from_value(value: u32) -> Option { + const INITIALIZING: u16 = BuildPhaseCode::Initializing as _; + const DEFAULT_BUILD: u16 = BuildPhaseCode::DefaultBuild as _; + const INTERNAL_BUILD: u16 = BuildPhaseCode::InternalBuild as _; + const EXTERNAL_BUILD: u16 = BuildPhaseCode::ExternalBuild as _; + const BUILD: u16 = BuildPhaseCode::Build as _; + const INSERTING: u16 = BuildPhaseCode::Inserting as _; + const COMPACTING: u16 = BuildPhaseCode::Compacting as _; + let k = value as u16; + match (value >> 16) as u16 { + INITIALIZING => Self::new(BuildPhaseCode::Initializing, k), + DEFAULT_BUILD => Self::new(BuildPhaseCode::DefaultBuild, k), + INTERNAL_BUILD => Self::new(BuildPhaseCode::InternalBuild, k), + EXTERNAL_BUILD => Self::new(BuildPhaseCode::ExternalBuild, k), + BUILD => Self::new(BuildPhaseCode::Build, k), + INSERTING => Self::new(BuildPhaseCode::Inserting, k), + COMPACTING => Self::new(BuildPhaseCode::Compacting, k), + _ => None, + } + } + pub const fn into_value(self) -> u32 { + (self.0 as u32) << 16 | (self.1 as u32) + } +} + +#[pgrx::pg_guard] +pub extern "C-unwind" fn ambuildphasename(x: i64) -> *mut core::ffi::c_char { + if let Ok(x) = u32::try_from(x.wrapping_sub(1)) { + if let Some(x) = BuildPhase::from_value(x) { + x.name().as_ptr().cast_mut() + } else { + std::ptr::null_mut() + } + } else { + std::ptr::null_mut() + } +} + +#[derive(Debug, Clone)] +struct PostgresReporter { + _phantom: PhantomData<*mut ()>, +} + +impl PostgresReporter { + fn phase(&self, phase: BuildPhase) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_SUBPHASE as _, + (phase.into_value() as i64) + 1, + ); + } + } + fn tuples_total(&self, tuples_total: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as _, + tuples_total as _, + ); + } + } + fn tuples_done(&self, tuples_done: u64) { + unsafe { + pgrx::pg_sys::pgstat_progress_update_param( + pgrx::pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as _, + tuples_done as _, + ); + } + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuild( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, +) -> *mut pgrx::pg_sys::IndexBuildResult { + use validator::Validate; + let (vector_options, vchordrq_options) = unsafe { options(index_relation) }; + if let Err(errors) = Validate::validate(&vector_options) { + pgrx::error!("error while validating options: {}", errors); + } + if let Err(errors) = Validate::validate(&vchordrq_options) { + pgrx::error!("error while validating options: {}", errors); + } + if vector_options.v == VectorKind::Rabitq8 && vchordrq_options.index.residual_quantization { + let errors = "residual_quantization is not supported for rabitq8 type"; + pgrx::error!("error while validating options: {errors}"); + } + if vector_options.v == VectorKind::Rabitq4 && vchordrq_options.index.residual_quantization { + let errors = "residual_quantization is not supported for rabitq4 type"; + pgrx::error!("error while validating options: {errors}"); + } + let opfamily = unsafe { opfamily(index_relation) }; + let reporter = PostgresReporter { + _phantom: PhantomData, + }; + reporter.tuples_total(unsafe { (*(*index_relation).rd_rel).reltuples as u64 }); + let mut structures = match vchordrq_options.build.source.clone() { + VchordrqBuildSourceOptions::Default(default_build) => { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::DefaultBuild)); + make_default_build(vector_options, default_build) + } + VchordrqBuildSourceOptions::Internal(internal_build) => { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::InternalBuild)); + let snapshot = if unsafe { (*index_info).ii_Concurrent } { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + let sampler = unsafe { HeapSampler::new(index_relation, heap_relation, snapshot) }; + let result = + make_internal_build(vector_options, opfamily, internal_build, sampler, &reporter); + if is_mvcc_snapshot(snapshot) { + unsafe { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + } + result + } + VchordrqBuildSourceOptions::External(external_build) => { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::ExternalBuild)); + make_external_build(vector_options, opfamily, external_build) + } + }; + for structure in structures.iter_mut() { + for centroid in structure.centroids.iter_mut() { + rabitq::rotate::rotate_inplace(centroid); + } + } + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Build)); + let index = unsafe { PostgresRelation::new(index_relation) }; + crate::index::vchordrq::dispatch::build( + vector_options, + vchordrq_options.index, + &index, + structures, + ); + let cached = if vchordrq_options.build.pin >= 0 { + let mut trace = vchordrq::cache(&index, vchordrq_options.build.pin); + trace.sort(); + trace.dedup(); + if let Some(max) = trace.last().copied() { + let mut mapping = vec![u32::MAX; 1 + max as usize]; + let mut pages = Vec::>>::with_capacity(trace.len()); + for id in trace { + mapping[id as usize] = pages.len() as u32; + pages.push(index.read(id).clone_into_boxed()); + } + vchordrq_cached::VchordrqCached::_1 { mapping, pages } + } else { + vchordrq_cached::VchordrqCached::_0 {} + } + } else { + vchordrq_cached::VchordrqCached::_0 {} + } + .serialize(); + if let Some(leader) = unsafe { + VchordrqLeader::enter( + c"vchordrq_parallel_build_main", + heap_relation, + index_relation, + (*index_info).ii_Concurrent, + &cached, + ) + } { + drop(cached); + unsafe { + leader.wait(); + parallel_build( + index_relation, + heap_relation, + index_info, + leader.tablescandesc, + leader.vchordrqshared, + leader.vchordrqcached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + || { + #[allow(clippy::needless_late_init)] + let order; + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).nparticipants = leader.nparticipants as u32; + order = (*shared).barrier_enter_0 as u32; + (*shared).barrier_enter_0 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_0, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_0 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_0, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_0 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_0, + ); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + order + }, + |indtuples, indexed_vectors| { + reporter.tuples_done(indtuples); + reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_1 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_1, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_1 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_1, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_1 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_1, + ); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); + }, + || { + // enter the barrier + let shared = leader.vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_2 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_2, + ); + // leave the barrier + let total = leader.nparticipants; + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_enter_2 == total { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_enter_2, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_leave_2 = true; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_leave_2, + ); + }, + ); + } + } else { + unsafe { + sequential_build( + index_relation, + heap_relation, + index_info, + &cached, + |indtuples| { + reporter.tuples_done(indtuples); + }, + || { + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); + }, + |indtuples, indexed_vectors| { + reporter.tuples_done(indtuples); + reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); + reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); + }, + || {}, + ); + } + } + unsafe { pgrx::pgbox::PgBox::::alloc0().into_pg() } +} + +struct VchordrqShared { + /* immutable state */ + heaprelid: pgrx::pg_sys::Oid, + indexrelid: pgrx::pg_sys::Oid, + isconcurrent: bool, + + /* locking */ + mutex: pgrx::pg_sys::slock_t, + condvar_barrier_enter_0: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_0: pgrx::pg_sys::ConditionVariable, + condvar_barrier_enter_1: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_1: pgrx::pg_sys::ConditionVariable, + condvar_barrier_enter_2: pgrx::pg_sys::ConditionVariable, + condvar_barrier_leave_2: pgrx::pg_sys::ConditionVariable, + + /* mutable state */ + barrier_enter_0: i32, + nparticipants: u32, + indtuples: u64, + indexed_vectors: u64, + barrier_leave_0: bool, + barrier_enter_1: i32, + barrier_leave_1: bool, + barrier_enter_2: i32, + barrier_leave_2: bool, +} + +mod vchordrq_cached { + pub const ALIGN: usize = 8; + pub type Tag = u64; + + use crate::index::storage::PostgresPage; + use index::tuples::RefChecker; + use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + + #[repr(C, align(8))] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VchordrqCachedHeader0 {} + + #[repr(C, align(8))] + #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] + struct VchordrqCachedHeader1 { + mapping_s: usize, + mapping_e: usize, + pages_s: usize, + pages_e: usize, + } + + pub enum VchordrqCached { + _0 {}, + _1 { + mapping: Vec, + pages: Vec>>, + }, + } + + impl VchordrqCached { + pub fn serialize(&self) -> Vec { + let mut buffer = Vec::new(); + match self { + VchordrqCached::_0 {} => { + buffer.extend((0 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + buffer[size_of::()..][..size_of::()] + .copy_from_slice(VchordrqCachedHeader0 {}.as_bytes()); + } + VchordrqCached::_1 { mapping, pages } => { + buffer.extend((1 as Tag).to_ne_bytes()); + buffer.extend(std::iter::repeat_n(0, size_of::())); + let mapping_s = buffer.len(); + buffer.extend(mapping.as_bytes()); + let mapping_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0u8); + } + let pages_s = buffer.len(); + buffer.extend(pages.iter().flat_map(|x| unsafe { + std::mem::transmute::< + &PostgresPage, + &[u8; pgrx::pg_sys::BLCKSZ as usize], + >(x.as_ref()) + })); + let pages_e = buffer.len(); + while buffer.len() % ALIGN != 0 { + buffer.push(0u8); + } + buffer[size_of::()..][..size_of::()] + .copy_from_slice( + VchordrqCachedHeader1 { + mapping_s, + mapping_e, + pages_s, + pages_e, + } + .as_bytes(), + ); + } + } + buffer + } + } + + #[derive(Debug, Clone, Copy)] + pub enum VchordrqCachedReader<'a> { + #[allow(dead_code)] + _0(VchordrqCachedReader0<'a>), + _1(VchordrqCachedReader1<'a>), + } + + #[derive(Debug, Clone, Copy)] + pub struct VchordrqCachedReader0<'a> { + #[allow(dead_code)] + header: &'a VchordrqCachedHeader0, + } + + #[derive(Debug, Clone, Copy)] + pub struct VchordrqCachedReader1<'a> { + #[allow(dead_code)] + header: &'a VchordrqCachedHeader1, + mapping: &'a [u32], + pages: &'a [PostgresPage], + } + + impl<'a> VchordrqCachedReader1<'a> { + pub fn get(&self, id: u32) -> Option<&'a PostgresPage> { + let index = *self.mapping.get(id as usize)?; + if index == u32::MAX { + return None; + } + Some(&self.pages[index as usize]) + } + } + + impl<'a> VchordrqCachedReader<'a> { + pub fn deserialize_ref(source: &'a [u8]) -> Self { + let tag = u64::from_ne_bytes(std::array::from_fn(|i| source[i])); + match tag { + 0 => { + let checker = RefChecker::new(source); + let header: &VchordrqCachedHeader0 = checker.prefix(size_of::()); + Self::_0(VchordrqCachedReader0 { header }) + } + 1 => { + let checker = RefChecker::new(source); + let header: &VchordrqCachedHeader1 = checker.prefix(size_of::()); + let mapping = checker.bytes(header.mapping_s, header.mapping_e); + let pages = + unsafe { checker.bytes_slice_unchecked(header.pages_s, header.pages_e) }; + Self::_1(VchordrqCachedReader1 { + header, + mapping, + pages, + }) + } + _ => panic!("bad bytes"), + } + } + } +} + +fn is_mvcc_snapshot(snapshot: *mut pgrx::pg_sys::SnapshotData) -> bool { + matches!( + unsafe { (*snapshot).snapshot_type }, + pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + | pgrx::pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC + ) +} + +struct VchordrqLeader { + pcxt: *mut pgrx::pg_sys::ParallelContext, + nparticipants: i32, + snapshot: pgrx::pg_sys::Snapshot, + vchordrqshared: *mut VchordrqShared, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordrqcached: *const u8, +} + +impl VchordrqLeader { + pub unsafe fn enter( + main: &'static CStr, + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + isconcurrent: bool, + vchordrq_cached: &[u8], + ) -> Option { + unsafe fn compute_parallel_workers( + heap_relation: pgrx::pg_sys::Relation, + index_relation: pgrx::pg_sys::Relation, + ) -> i32 { + unsafe { + if pgrx::pg_sys::plan_create_index_workers( + (*heap_relation).rd_id, + (*index_relation).rd_id, + ) == 0 + { + return 0; + } + if !(*heap_relation).rd_options.is_null() { + let std_options = (*heap_relation) + .rd_options + .cast::(); + std::cmp::min( + (*std_options).parallel_workers, + pgrx::pg_sys::max_parallel_maintenance_workers, + ) + } else { + pgrx::pg_sys::max_parallel_maintenance_workers + } + } + } + + let request = unsafe { compute_parallel_workers(heap_relation, index_relation) }; + if request <= 0 { + return None; + } + + unsafe { + pgrx::pg_sys::EnterParallelMode(); + } + let pcxt = unsafe { + pgrx::pg_sys::CreateParallelContext(c"vchord".as_ptr(), main.as_ptr(), request) + }; + + let snapshot = if isconcurrent { + unsafe { pgrx::pg_sys::RegisterSnapshot(pgrx::pg_sys::GetTransactionSnapshot()) } + } else { + &raw mut pgrx::pg_sys::SnapshotAnyData + }; + + fn estimate_chunk(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.space_for_chunks += x.next_multiple_of(pgrx::pg_sys::ALIGNOF_BUFFER as _); + } + fn estimate_keys(e: &mut pgrx::pg_sys::shm_toc_estimator, x: usize) { + e.number_of_keys += x; + } + let est_tablescandesc = + unsafe { pgrx::pg_sys::table_parallelscan_estimate(heap_relation, snapshot) }; + unsafe { + estimate_chunk(&mut (*pcxt).estimator, size_of::()); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, est_tablescandesc); + estimate_keys(&mut (*pcxt).estimator, 1); + estimate_chunk(&mut (*pcxt).estimator, 8 + vchordrq_cached.len()); + estimate_keys(&mut (*pcxt).estimator, 1); + } + + unsafe { + pgrx::pg_sys::InitializeParallelDSM(pcxt); + if (*pcxt).seg.is_null() { + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + let vchordrqshared = unsafe { + let vchordrqshared = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) + .cast::(); + vchordrqshared.write(VchordrqShared { + heaprelid: (*heap_relation).rd_id, + indexrelid: (*index_relation).rd_id, + isconcurrent, + nparticipants: 0, + condvar_barrier_enter_0: std::mem::zeroed(), + condvar_barrier_leave_0: std::mem::zeroed(), + condvar_barrier_enter_1: std::mem::zeroed(), + condvar_barrier_leave_1: std::mem::zeroed(), + condvar_barrier_enter_2: std::mem::zeroed(), + condvar_barrier_leave_2: std::mem::zeroed(), + barrier_enter_0: 0, + barrier_leave_0: false, + barrier_enter_1: 0, + barrier_leave_1: false, + barrier_enter_2: 0, + barrier_leave_2: false, + mutex: std::mem::zeroed(), + indtuples: 0, + indexed_vectors: 0, + }); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_0); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_0); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_1); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_1); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_2); + pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_2); + pgrx::pg_sys::SpinLockInit(&raw mut (*vchordrqshared).mutex); + vchordrqshared + }; + + let tablescandesc = unsafe { + let tablescandesc = pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, est_tablescandesc) + .cast::(); + pgrx::pg_sys::table_parallelscan_initialize(heap_relation, tablescandesc, snapshot); + tablescandesc + }; + + let vchordrqcached = unsafe { + let x = + pgrx::pg_sys::shm_toc_allocate((*pcxt).toc, 8 + vchordrq_cached.len()).cast::(); + (x as *mut u64).write_unaligned(vchordrq_cached.len() as _); + std::ptr::copy(vchordrq_cached.as_ptr(), x.add(8), vchordrq_cached.len()); + x + }; + + unsafe { + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000001, vchordrqshared.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000002, tablescandesc.cast()); + pgrx::pg_sys::shm_toc_insert((*pcxt).toc, 0xA000000000000003, vchordrqcached.cast()); + } + + unsafe { + pgrx::pg_sys::LaunchParallelWorkers(pcxt); + } + + let nworkers_launched = unsafe { (*pcxt).nworkers_launched }; + + unsafe { + if nworkers_launched == 0 { + pgrx::pg_sys::WaitForParallelWorkersToFinish(pcxt); + if is_mvcc_snapshot(snapshot) { + pgrx::pg_sys::UnregisterSnapshot(snapshot); + } + pgrx::pg_sys::DestroyParallelContext(pcxt); + pgrx::pg_sys::ExitParallelMode(); + return None; + } + } + + Some(Self { + pcxt, + nparticipants: nworkers_launched + 1, + snapshot, + vchordrqshared, + tablescandesc, + vchordrqcached, + }) + } + + pub fn wait(&self) { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToAttach(self.pcxt); + } + } +} + +impl Drop for VchordrqLeader { + fn drop(&mut self) { + if !std::thread::panicking() { + unsafe { + pgrx::pg_sys::WaitForParallelWorkersToFinish(self.pcxt); + if is_mvcc_snapshot(self.snapshot) { + pgrx::pg_sys::UnregisterSnapshot(self.snapshot); + } + pgrx::pg_sys::DestroyParallelContext(self.pcxt); + pgrx::pg_sys::ExitParallelMode(); + } + } + } +} + +#[pgrx::pg_guard] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( + _seg: *mut pgrx::pg_sys::dsm_segment, + toc: *mut pgrx::pg_sys::shm_toc, +) { + let _ = rand::rng().reseed(); + let vchordrqshared = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000001, false).cast::() + }; + let tablescandesc = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000002, false) + .cast::() + }; + let vchordrqcached = unsafe { + pgrx::pg_sys::shm_toc_lookup(toc, 0xA000000000000003, false) + .cast::() + .cast_const() + }; + let (heap_lockmode, index_lockmode) = if unsafe { !(*vchordrqshared).isconcurrent } { + ( + pgrx::pg_sys::ShareLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + } else { + ( + pgrx::pg_sys::ShareUpdateExclusiveLock as pgrx::pg_sys::LOCKMODE, + pgrx::pg_sys::RowExclusiveLock as pgrx::pg_sys::LOCKMODE, + ) + }; + let heap = unsafe { pgrx::pg_sys::table_open((*vchordrqshared).heaprelid, heap_lockmode) }; + let index = unsafe { pgrx::pg_sys::index_open((*vchordrqshared).indexrelid, index_lockmode) }; + let index_info = unsafe { pgrx::pg_sys::BuildIndexInfo(index) }; + unsafe { + (*index_info).ii_Concurrent = (*vchordrqshared).isconcurrent; + } + + unsafe { + parallel_build( + index, + heap, + index_info, + tablescandesc, + vchordrqshared, + vchordrqcached, + |_| (), + || { + #[allow(clippy::needless_late_init)] + let order; + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + order = (*shared).barrier_enter_0 as u32; + (*shared).barrier_enter_0 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_0, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_0 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_0, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + order + }, + |_, _| { + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_1 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_1, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_1 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_1, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + }, + || { + // enter the barrier + let shared = vchordrqshared; + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + (*shared).barrier_enter_2 += 1; + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableBroadcast( + &raw mut (*shared).condvar_barrier_enter_2, + ); + // leave the barrier + loop { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); + if (*shared).barrier_leave_2 { + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + break; + } + pgrx::pg_sys::SpinLockRelease(&raw mut (*shared).mutex); + pgrx::pg_sys::ConditionVariableSleep( + &raw mut (*shared).condvar_barrier_leave_2, + pgrx::pg_sys::WaitEventIPC::WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN as _, + ); + } + pgrx::pg_sys::ConditionVariableCancelSleep(); + }, + ); + } + + unsafe { + pgrx::pg_sys::index_close(index, index_lockmode); + pgrx::pg_sys::table_close(heap, heap_lockmode); + } +} + +unsafe fn parallel_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + tablescandesc: *mut pgrx::pg_sys::ParallelTableScanDescData, + vchordrqshared: *mut VchordrqShared, + vchordrqcached: *const u8, + mut callback: impl FnMut(u64), + sync_0: impl FnOnce() -> u32, + sync_1: impl FnOnce(u64, u64), + sync_2: impl FnOnce(), +) { + use vchordrq_cached::VchordrqCachedReader; + + let cached = VchordrqCachedReader::deserialize_ref(unsafe { + let bytes = (vchordrqcached as *const u64).read_unaligned(); + std::slice::from_raw_parts(vchordrqcached.add(8), bytes as _) + }); + + let scan = unsafe { pgrx::pg_sys::table_beginscan_parallel(heap_relation, tablescandesc) }; + let opfamily = unsafe { opfamily(index_relation) }; + let traverser = unsafe { HeapTraverser::new(heap_relation, index_relation, index_info, scan) }; + + struct IdChooser(u32); + impl InsertChooser for IdChooser { + fn choose(&mut self, n: NonZero) -> usize { + self.0 as usize % n.get() + } + } + + struct ChooseSome { + n: usize, + k: usize, + } + impl MaintainChooser for ChooseSome { + fn choose(&mut self, i: usize) -> bool { + i % self.n == self.k + } + } + + let check = || { + pgrx::check_for_interrupts!(); + }; + + let order = sync_0(); + + let index = unsafe { BufferedPostgresRelation::new(index_relation) }; + + match cached { + VchordrqCachedReader::_0(_) => { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + let indexed_vectors = store.len() as u64; + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + let mut chooser = IdChooser(order); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::dispatch::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); + } + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + callback(indtuples); + } + }); + } + VchordrqCachedReader::_1(cached) => { + let index = CachingRelation { + cache: cached, + relation: &index, + }; + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + let indexed_vectors = store.len() as u64; + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + let mut chooser = IdChooser(order); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::dispatch::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); + } + unsafe { + let indtuples; + { + pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); + (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; + indtuples = (*vchordrqshared).indtuples; + pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); + } + callback(indtuples); + } + }); + } + } + + while index.release() { + vchordrq::consume(&index); + } + + drop(index); + + let (indtuples, indexed_vectors) = unsafe { + ( + (*vchordrqshared).indtuples, + (*vchordrqshared).indexed_vectors, + ) + }; + sync_1(indtuples, indexed_vectors); + + let index = unsafe { PostgresRelation::new(index_relation) }; + + let mut chooser = ChooseSome { + n: unsafe { (*vchordrqshared).nparticipants as usize }, + k: order as usize, + }; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_2(); +} + +unsafe fn sequential_build( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + index_info: *mut pgrx::pg_sys::IndexInfo, + vchordrqcached: &[u8], + mut callback: impl FnMut(u64), + sync_0: impl FnOnce(), + sync_1: impl FnOnce(u64, u64), + sync_2: impl FnOnce(), +) { + use vchordrq_cached::VchordrqCachedReader; + + let cached = VchordrqCachedReader::deserialize_ref(vchordrqcached); + + let opfamily = unsafe { opfamily(index_relation) }; + let traverser = unsafe { + HeapTraverser::new( + heap_relation, + index_relation, + index_info, + std::ptr::null_mut(), + ) + }; + + struct ChooseZero; + impl InsertChooser for ChooseZero { + fn choose(&mut self, _: NonZero) -> usize { + 0 + } + } + + struct ChooseAll; + impl MaintainChooser for ChooseAll { + fn choose(&mut self, _: usize) -> bool { + true + } + } + + let check = || { + pgrx::check_for_interrupts!(); + }; + + sync_0(); + + let index = unsafe { BufferedPostgresRelation::new(index_relation) }; + + let mut indtuples = 0; + let mut indexed_vectors = 0_u64; + match cached { + VchordrqCachedReader::_0(_) => { + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + indexed_vectors += store.len() as u64; + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + let mut chooser = ChooseZero; + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::dispatch::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); + } + indtuples += 1; + callback(indtuples); + }); + } + VchordrqCachedReader::_1(cached) => { + let index = CachingRelation { + cache: cached, + relation: &index, + }; + traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { + let ctid = tuple.id(); + let (values, is_nulls) = tuple.build(); + let value = unsafe { (!is_nulls.add(0).read()).then_some(values.add(0).read()) }; + let store = value + .and_then(|x| unsafe { opfamily.store(x) }) + .unwrap_or_default(); + indexed_vectors += store.len() as u64; + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + let mut chooser = ChooseZero; + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::dispatch::insert( + opfamily, + &index, + payload, + vector, + true, + true, + &mut chooser, + &bump, + ); + } + indtuples += 1; + callback(indtuples); + }); + } + } + + while index.release() { + vchordrq::consume(&index); + } + + drop(index); + + sync_1(indtuples, indexed_vectors); + + let index = unsafe { PostgresRelation::new(index_relation) }; + + let mut chooser = ChooseAll; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_2(); +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambuildempty(_index_relation: pgrx::pg_sys::Relation) { + pgrx::error!("Unlogged indexes are not supported."); +} + +unsafe fn options( + index_relation: pgrx::pg_sys::Relation, +) -> (VectorOptions, VchordrqIndexingOptions) { + let att = unsafe { &mut *(*index_relation).rd_att }; + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + let atts = unsafe { att.attrs.as_slice(att.natts as _) }; + #[cfg(feature = "pg18")] + let atts = unsafe { + let ptr = att + .compact_attrs + .as_ptr() + .add(att.natts as _) + .cast::(); + std::slice::from_raw_parts(ptr, att.natts as _) + }; + if atts.is_empty() { + pgrx::error!("indexing on no columns is not supported"); + } + if atts.len() != 1 { + pgrx::error!("multicolumn index is not supported"); + } + // get dim + let typmod = Typmod::new(atts[0].atttypmod).unwrap(); + let dim = if let Some(dim) = typmod.dim() { + dim.get() + } else { + pgrx::error!( + "Dimensions type modifier of a vector column is needed for building the index." + ); + }; + // get v, d + let opfamily = unsafe { opfamily(index_relation) }; + let vector = VectorOptions { + dim, + v: opfamily.vector_kind(), + d: opfamily.distance_kind(), + }; + // get indexing options + let indexing_options = { + let reloption = unsafe { (*index_relation).rd_options as *const Reloption }; + let s = unsafe { Reloption::options(reloption, c"") }.to_string_lossy(); + match toml::from_str::(&s) { + Ok(p) => p, + Err(e) => pgrx::error!("failed to parse options: {}", e), + } + }; + (vector, indexing_options) +} + +fn make_default_build( + vector_options: VectorOptions, + _default_build: VchordrqDefaultBuildOptions, +) -> Vec> { + vec![Structure:: { + centroids: vec![vec![0.0f32; vector_options.dim as usize]], + children: vec![vec![]], + }] +} + +fn make_internal_build( + vector_options: VectorOptions, + opfamily: Opfamily, + internal_build: VchordrqInternalBuildOptions, + sampler: impl Sampler, + reporter: &PostgresReporter, +) -> Vec> { + use humansize::{BINARY, format_size}; + use std::iter::once; + let (reduction, sample_dim) = match internal_build.kmeans_dimension { + None => (None, vector_options.dim as usize), + Some(d) if d < vector_options.dim => (Some(d as usize), d as usize), + Some(d) => { + pgrx::warning!( + "ignoring `kmeans_dimension = {}` because it is less than the vector dimension {}", + d, + vector_options.dim + ); + (None, vector_options.dim as usize) + } + }; + { + let d = sample_dim as u64; + let c = internal_build.lists.last().copied().unwrap_or_default() as u64; + let f = internal_build.sampling_factor as u64; + let t = internal_build.build_threads as u64; + let estimated_memory_usage = 4 * c * d * (1 + t + f); + pgrx::info!( + "clustering: estimated memory usage is {}", + format_size( + estimated_memory_usage, + BINARY.decimal_places(2).decimal_zeroes(2) + ) + ); + } + let max_number_of_samples = internal_build + .lists + .last() + .map(|x| x.saturating_mul(internal_build.sampling_factor)) + .unwrap_or_default(); + let mut sample = sampler.sample(); + let samples = 'samples: { + let mut samples = Square::with_capacity(sample_dim, max_number_of_samples as _); + if samples.len() >= max_number_of_samples as usize { + break 'samples samples; + } + while let Some(mut tuple) = sample.next() { + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + for (vector, _) in vectors { + let mut x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + OwnedVector::Rabitq8(x) => Rabitq8Owned::normalize(x), + OwnedVector::Rabitq4(x) => Rabitq4Owned::normalize(x), + }; + assert_eq!( + vector_options.dim, + x.len() as u32, + "invalid vector dimensions" + ); + if let Some(sample_dim) = reduction { + rabitq::rotate::rotate_inplace(&mut x); + x.truncate(sample_dim); + } + samples.push_slice(x.as_slice()); + if samples.len() >= max_number_of_samples as usize { + break 'samples samples; + } + } + } + } + } + samples + }; + drop(sample); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(internal_build.build_threads as usize) + .build() + .expect("failed to build thread pool"); + pgrx::info!("clustering: using {} threads", pool.current_num_threads()); + let mut result = Vec::>::new(); + let mut samples = Some(samples); + for w in internal_build.lists.iter().rev().copied().chain(once(1)) { + let mut input = if let Some(structure) = result.last() { + let mut input = + Square::with_capacity(vector_options.dim as _, structure.centroids.len()); + for slice in structure.centroids.iter() { + input.push_slice(slice); + } + input + } else if let Some(samples) = samples.take() { + samples + } else { + unreachable!() + }; + let num_points = input.len(); + let num_dim = input.d(); + let num_lists = w as usize; + let num_iterations = internal_build.kmeans_iterations as _; + if result.is_empty() { + let percentage = 0; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = + BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage).unwrap_or(default); + reporter.phase(phase); + } + if num_lists > 1 { + pgrx::info!( + "clustering: starting, clustering {num_points} vectors of {num_dim} dimension into {num_lists} clusters, in {num_iterations} iterations" + ); + } + let mut f = 'f: { + let view = input.as_mut_view(); + if result.last().is_some() { + break 'f k_means::lloyd_k_means( + &pool, + num_dim, + view, + num_lists, + [7; 32], + internal_build.spherical_centroids, + ); + }; + match internal_build.kmeans_algorithm { + KMeansAlgorithm::Lloyd {} => k_means::lloyd_k_means( + &pool, + num_dim, + view, + num_lists, + [7; 32], + internal_build.spherical_centroids, + ), + KMeansAlgorithm::Hierarchical {} => k_means::hierarchical_k_means( + &pool, + num_dim, + view, + num_lists, + [7; 32], + internal_build.spherical_centroids, + ), + } + }; + for i in 0..num_iterations { + pgrx::check_for_interrupts!(); + if result.is_empty() { + let percentage = + ((i as f64 / num_iterations as f64) * 100.0).clamp(0.0, 100.0) as u16; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage) + .unwrap_or(default); + reporter.phase(phase); + } + if num_lists > 1 { + pgrx::info!("clustering: iteration {}", i + 1); + } + f.assign(); + f.update(); + } + let centroids = 'centroids: { + if result.last().is_some() { + break 'centroids f.finish(); + } + if let Some(sample_dim) = reduction { + let mut sample = sampler.sample(); + let index = f.index(); + let index = &index; + let is_spherical = internal_build.spherical_centroids; + let num_threads = { + let config = internal_build.build_threads as f64; + let ratio = sample_dim as f64 / vector_options.dim as f64; + (config * ratio).ceil().max(1.0).min(config) as usize + }; + let (tx, rx) = crossbeam_channel::bounded::>(1024); + let list = std::thread::scope(move |scope| { + assert_ne!(num_threads, 0); + let mut handles = Vec::with_capacity(num_threads); + for _ in 0..num_threads { + let rx = rx.clone(); + let mut sum = Square::from_zeros(vector_options.dim as _, num_lists); + let mut count = vec![0.0f32; num_lists]; + handles.push(scope.spawn(move || { + while let Ok(sample) = rx.recv() { + let mut x = rabitq::rotate::rotate(&sample); + x.truncate(sample_dim); + let (_, id) = index(&x); + f32::vector_add_inplace(&mut sum[id], &sample); + count[id] += 1.0; + } + (sum, count) + })); + } + { + drop(rx); + let tx = tx; + let mut samples = 0_usize; + 'samples: { + if samples >= max_number_of_samples as usize { + break 'samples; + } + while let Some(mut tuple) = sample.next() { + let (values, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(values[0]); + if let Some(datum) = datum { + let vectors = unsafe { opfamily.store(datum) }; + if let Some(vectors) = vectors { + for (vector, _) in vectors { + let x = match vector { + OwnedVector::Vecf32(x) => VectOwned::normalize(x), + OwnedVector::Vecf16(x) => VectOwned::normalize(x), + OwnedVector::Rabitq8(x) => { + Rabitq8Owned::normalize(x) + } + OwnedVector::Rabitq4(x) => { + Rabitq4Owned::normalize(x) + } + }; + assert_eq!( + vector_options.dim, + x.len() as u32, + "invalid vector dimensions" + ); + if tx.send(x).is_err() { + break 'samples; + } + samples += 1; + if samples >= max_number_of_samples as usize { + break 'samples; + } + } + } + } + } + } + } + handles + .into_iter() + .map(|handle| handle.join().expect("failed to spawn threads")) + .collect::>() + }); + let mut sum = Square::from_zeros(vector_options.dim as _, num_lists); + let mut count = vec![0.0f32; num_lists]; + for (sum_1, count_1) in list { + for i in 0..num_lists { + f32::vector_add_inplace(&mut sum[i], &sum_1[i]); + count[i] += count_1[i]; + } + } + pool.install(|| { + use rayon::prelude::*; + + sum.par_iter_mut() + .zip(count.par_iter()) + .for_each(|(sum, count)| f32::vector_mul_scalar_inplace(sum, 1.0 / count)); + + let mut centroids = sum; + + centroids + .par_iter_mut() + .zip(count.par_iter()) + .filter(|&(_, &count)| count == 0.0) + .for_each(|(centroid, _)| { + centroid.fill_with(|| rand::random_range(-1.0..=1.0)) + }); + + if is_spherical { + (&mut centroids).into_par_iter().for_each(|centroid| { + let l = f32::reduce_sum_of_x2(centroid).sqrt(); + f32::vector_mul_scalar_inplace(centroid, 1.0 / l); + }); + } + + centroids + }) + } else { + f.finish() + } + }; + + if result.is_empty() { + let percentage = 100; + let default = BuildPhase::from_code(BuildPhaseCode::InternalBuild); + let phase = + BuildPhase::new(BuildPhaseCode::InternalBuild, 1 + percentage).unwrap_or(default); + reporter.phase(phase); + } + if num_lists > 1 { + pgrx::info!("clustering: finished"); + } + if let Some(structure) = result.last() { + let mut children = vec![Vec::new(); centroids.len()]; + for i in 0..structure.len() as u32 { + let target = k_means_lookup(&structure.centroids[i as usize], ¢roids); + children[target].push(i); + } + let (centroids, children) = std::iter::zip(¢roids, children) + .filter(|(_, children)| !children.is_empty()) + .map(|(centroids, children)| (centroids.to_vec(), children)) + .unzip::<_, _, Vec<_>, Vec<_>>(); + result.push(Structure { + centroids, + children, + }); + } else { + let children = vec![Vec::new(); centroids.len()]; + result.push(Structure { + centroids: centroids.into_iter().map(|x| x.to_vec()).collect(), + children, + }); + } + } + result +} + +#[allow(clippy::collapsible_else_if)] +fn make_external_build( + vector_options: VectorOptions, + _opfamily: Opfamily, + external_build: VchordrqExternalBuildOptions, +) -> Vec> { + use std::collections::BTreeMap; + let VchordrqExternalBuildOptions { table } = external_build; + let mut parents = BTreeMap::new(); + let mut vectors = BTreeMap::new(); + pgrx::spi::Spi::connect(|client| { + use crate::datatype::memory_vector::VectorOutput; + use pgrx::pg_sys::panic::ErrorReportable; + use vector::VectorBorrowed; + let schema_query = "SELECT n.nspname::TEXT + FROM pg_catalog.pg_extension e + LEFT JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vector';"; + let pgvector_schema: String = client + .select(schema_query, None, &[]) + .unwrap_or_report() + .first() + .get_by_name("nspname") + .expect("external build: cannot get schema of pgvector") + .expect("external build: cannot get schema of pgvector"); + let dump_query = + format!("SELECT id, parent, vector::{pgvector_schema}.vector FROM {table};"); + let centroids = client.select(&dump_query, None, &[]).unwrap_or_report(); + for row in centroids { + let id: Option = row.get_by_name("id").unwrap(); + let parent: Option = row.get_by_name("parent").unwrap(); + let vector: Option = row.get_by_name("vector").unwrap(); + let id = id.expect("external build: id could not be NULL"); + let vector = vector.expect("external build: vector could not be NULL"); + let pop = parents.insert(id, parent); + if pop.is_some() { + pgrx::error!( + "external build: there are at least two lines have same id, id = {id}" + ); + } + if vector_options.dim != vector.as_borrowed().dim() { + pgrx::error!("external build: incorrect dimension, id = {id}"); + } + vectors.insert(id, vector.as_borrowed().slice().to_vec()); + } + }); + if parents.len() >= 2 && parents.values().all(|x| x.is_none()) { + // if there are more than one vertex and no edges, + // assume there is an implicit root + let n = parents.len(); + let mut result = Vec::new(); + result.push(Structure { + centroids: vectors.values().cloned().collect::>(), + children: vec![Vec::new(); n], + }); + result.push(Structure { + centroids: vec![{ + // compute the vector on root, without normalizing it + let mut sum = vec![0.0f32; vector_options.dim as _]; + for vector in vectors.values() { + f32::vector_add_inplace(&mut sum, vector); + } + f32::vector_mul_scalar_inplace(&mut sum, 1.0 / n as f32); + sum + }], + children: vec![(0..n as u32).collect()], + }); + return result; + } + let mut children = parents + .keys() + .map(|x| (*x, Vec::new())) + .collect::>(); + let mut root = None; + for (&id, &parent) in parents.iter() { + if let Some(parent) = parent { + if let Some(parent) = children.get_mut(&parent) { + parent.push(id); + } else { + pgrx::error!("external build: parent does not exist, id = {id}, parent = {parent}"); + } + } else { + if let Some(root) = root { + pgrx::error!("external build: two root, id = {root}, id = {id}"); + } else { + root = Some(id); + } + } + } + let Some(root) = root else { + pgrx::error!("external build: there are no root"); + }; + let mut heights = BTreeMap::<_, _>::new(); + fn dfs_for_heights( + heights: &mut BTreeMap>, + children: &BTreeMap>, + u: i32, + ) { + if heights.contains_key(&u) { + pgrx::error!("external build: detect a cycle, id = {u}"); + } + heights.insert(u, None); + let mut height = None; + for &v in children[&u].iter() { + dfs_for_heights(heights, children, v); + let new = heights[&v].unwrap() + 1; + if let Some(height) = height { + if height != new { + pgrx::error!("external build: two heights, id = {u}"); + } + } else { + height = Some(new); + } + } + if height.is_none() { + height = Some(1); + } + heights.insert(u, height); + } + dfs_for_heights(&mut heights, &children, root); + let heights = heights + .into_iter() + .map(|(k, v)| (k, v.expect("not a connected graph"))) + .collect::>(); + if !(1..=8).contains(&(heights[&root] - 1)) { + pgrx::error!( + "external build: unexpected tree height, height = {}", + heights[&root] + ); + } + let mut cursors = vec![0_u32; 1 + heights[&root] as usize]; + let mut labels = BTreeMap::new(); + for id in parents.keys().copied() { + let height = heights[&id]; + let cursor = cursors[height as usize]; + labels.insert(id, (height, cursor)); + cursors[height as usize] += 1; + } + fn extract( + height: u32, + labels: &BTreeMap, + vectors: &BTreeMap, + children: &BTreeMap>, + ) -> (Vec, Vec>) { + labels + .iter() + .filter(|(_, (h, _))| *h == height) + .map(|(id, _)| { + ( + vectors[id].clone(), + children[id].iter().map(|id| labels[id].1).collect(), + ) + }) + .unzip() + } + let mut result = Vec::new(); + for height in 1..=heights[&root] { + let (centroids, children) = extract(height, &labels, &vectors, &children); + result.push(Structure { + centroids, + children, + }); + } + result +} + +struct CachingRelation<'a, R> { + cache: vchordrq_cached::VchordrqCachedReader1<'a>, + relation: &'a R, +} + +enum CachingRelationReadGuard<'a, G: Deref> { + Wrapping(G), + Cached(u32, &'a G::Target), +} + +impl PageGuard for CachingRelationReadGuard<'_, G> { + fn id(&self) -> u32 { + match self { + CachingRelationReadGuard::Wrapping(x) => x.id(), + CachingRelationReadGuard::Cached(id, _) => *id, + } + } +} + +impl Deref for CachingRelationReadGuard<'_, G> { + type Target = G::Target; + + fn deref(&self) -> &Self::Target { + match self { + CachingRelationReadGuard::Wrapping(x) => x, + CachingRelationReadGuard::Cached(_, page) => page, + } + } +} + +impl Relation for CachingRelation<'_, R> { + type Page = R::Page; +} + +impl>> RelationReadTypes + for CachingRelation<'_, R> +{ + type ReadGuard<'a> = CachingRelationReadGuard<'a, R::ReadGuard<'a>>; +} + +impl>> RelationRead + for CachingRelation<'_, R> +{ + fn read(&self, id: u32) -> Self::ReadGuard<'_> { + if let Some(x) = self.cache.get(id) { + CachingRelationReadGuard::Cached(id, x) + } else { + CachingRelationReadGuard::Wrapping(self.relation.read(id)) + } + } +} + +impl>> RelationWriteTypes + for CachingRelation<'_, R> +{ + type WriteGuard<'a> = R::WriteGuard<'a>; +} + +impl>> RelationWrite + for CachingRelation<'_, R> +{ + fn write(&self, id: u32, tracking_freespace: bool) -> Self::WriteGuard<'_> { + self.relation.write(id, tracking_freespace) + } + + fn extend( + &self, + opaque: ::Opaque, + tracking_freespace: bool, + ) -> Self::WriteGuard<'_> { + self.relation.extend(opaque, tracking_freespace) + } + + fn search(&self, freespace: usize) -> Option> { + self.relation.search(freespace) + } +} diff --git a/src/index/vchordrq/am/am_vacuumcleanup.rs b/src/index/vchordrq/am/am_vacuumcleanup.rs new file mode 100644 index 00000000..12e19075 --- /dev/null +++ b/src/index/vchordrq/am/am_vacuumcleanup.rs @@ -0,0 +1,64 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::vchordrq::am::PostgresRelation; +use crate::index::vchordrq::opclass::opfamily; +use vchordrq::MaintainChooser; + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvacuumcleanup( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let index_relation = unsafe { (*info).index }; + unsafe { + sequential_vacuumcleanup(index_relation, || (), || ()); + } + stats +} + +unsafe fn sequential_vacuumcleanup( + index_relation: pgrx::pg_sys::Relation, + sync_0: impl FnOnce(), + sync_1: impl FnOnce(), +) { + struct ChooseAll; + impl MaintainChooser for ChooseAll { + fn choose(&mut self, _: usize) -> bool { + true + } + } + + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let check = || unsafe { + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + + sync_0(); + + let mut chooser = ChooseAll; + crate::index::vchordrq::dispatch::maintain(opfamily, &index, &mut chooser, check); + + sync_1(); +} diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs new file mode 100644 index 00000000..976692b6 --- /dev/null +++ b/src/index/vchordrq/am/mod.rs @@ -0,0 +1,730 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod am_build; +mod am_vacuumcleanup; + +use crate::index::fetcher::*; +use crate::index::gucs; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::*; +use crate::recorder::DefaultRecorder; +use pgrx::datum::Internal; +use pgrx::pg_sys::Datum; +use rand::RngExt; +use std::cell::LazyCell; +use std::ffi::CStr; +use std::num::NonZero; +use std::ops::DerefMut; +use std::ptr::NonNull; +use std::sync::OnceLock; +use vchordrq::InsertChooser; + +#[repr(C)] +pub struct Reloption { + vl_len_: i32, + options: i32, + probes: i32, + epsilon: f64, + maxsim_refine: i32, + maxsim_threshold: i32, +} + +impl Reloption { + unsafe fn options<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { + unsafe { + if this.is_null() { + return default; + } + let count = (&raw const (*this).options).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); + CStr::from_ptr(ptr.cast()) + } + } + pub unsafe fn probes<'a>(this: *const Self, default: &'static CStr) -> &'a CStr { + unsafe { + if this.is_null() { + return default; + } + let count = (&raw const (*this).probes).read(); + if count == 0 { + return default; + } + let ptr = this.cast::().add(count as _); + CStr::from_ptr(ptr.cast()) + } + } + pub unsafe fn epsilon(this: *const Self, default: f64) -> f64 { + unsafe { + if this.is_null() { + return default; + } + (*this).epsilon + } + } + pub unsafe fn maxsim_refine(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).maxsim_refine + } + } + pub unsafe fn maxsim_threshold(this: *const Self, default: i32) -> i32 { + unsafe { + if this.is_null() { + return default; + } + (*this).maxsim_threshold + } + } +} + +const TABLE: &[pgrx::pg_sys::relopt_parse_elt] = &[ + pgrx::pg_sys::relopt_parse_elt { + optname: c"options".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, options) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"probes".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_STRING, + offset: std::mem::offset_of!(Reloption, probes) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"epsilon".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_REAL, + offset: std::mem::offset_of!(Reloption, epsilon) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"maxsim_refine".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, maxsim_refine) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, + pgrx::pg_sys::relopt_parse_elt { + optname: c"maxsim_threshold".as_ptr(), + opttype: pgrx::pg_sys::relopt_type::RELOPT_TYPE_INT, + offset: std::mem::offset_of!(Reloption, maxsim_threshold) as i32, + #[cfg(feature = "pg18")] + isset_offset: 0, + }, +]; + +static RELOPT_KIND: OnceLock = OnceLock::new(); + +pub fn init() { + RELOPT_KIND.get_or_init(|| { + let kind; + unsafe { + kind = pgrx::pg_sys::add_reloption_kind(); + pgrx::pg_sys::add_string_reloption( + kind as _, + c"options".as_ptr(), + c"Vector index options, represented as a TOML string.".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_string_reloption( + kind as _, + c"probes".as_ptr(), + c"Search parameter `vchordrq.probes`".as_ptr(), + c"".as_ptr(), + None, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_real_reloption( + kind as _, + c"epsilon".as_ptr(), + c"Search parameter `vchordrq.epsilon`".as_ptr(), + 1.9, + 0.0, + 4.0, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"maxsim_refine".as_ptr(), + c"Search parameter `vchordrq.maxsim_refine`".as_ptr(), + 0, + 0, + i32::MAX, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + pgrx::pg_sys::add_int_reloption( + kind as _, + c"maxsim_threshold".as_ptr(), + c"Search parameter `vchordrq.maxsim_threshold`".as_ptr(), + 0, + 0, + i32::MAX, + pgrx::pg_sys::AccessExclusiveLock as pgrx::pg_sys::LOCKMODE, + ); + } + kind + }); +} + +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_amhandler(_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> Internal { + type T = pgrx::pg_sys::IndexAmRoutine; + unsafe { + let index_am_routine = pgrx::pg_sys::palloc0(size_of::()) as *mut T; + index_am_routine.write(AM_HANDLER); + Internal::from(Some(Datum::from(index_am_routine))) + } +} + +const AM_HANDLER: pgrx::pg_sys::IndexAmRoutine = const { + let mut am_routine = unsafe { std::mem::zeroed::() }; + + am_routine.type_ = pgrx::pg_sys::NodeTag::T_IndexAmRoutine; + + am_routine.amsupport = 1; + am_routine.amcanorderbyop = true; + + #[cfg(any(feature = "pg17", feature = "pg18"))] + { + am_routine.amcanbuildparallel = true; + } + + // Index access methods that set `amoptionalkey` to `false` + // must index all tuples, even if the first column is `NULL`. + // However, PostgreSQL does not generate a path if there is no + // index clauses, even if there is a `ORDER BY` clause. + // So we have to set it to `true` and set costs of every path + // for vector index scans without `ORDER BY` clauses a large number + // and throw errors if someone really wants such a path. + am_routine.amoptionalkey = true; + + am_routine.amvalidate = Some(amvalidate); + am_routine.amoptions = Some(amoptions); + am_routine.amcostestimate = Some(amcostestimate); + + am_routine.ambuildphasename = Some(am_build::ambuildphasename); + am_routine.ambuild = Some(am_build::ambuild); + am_routine.ambuildempty = Some(am_build::ambuildempty); + am_routine.aminsert = Some(aminsert); + am_routine.ambulkdelete = Some(ambulkdelete); + am_routine.amvacuumcleanup = Some(am_vacuumcleanup::amvacuumcleanup); + + am_routine.ambeginscan = Some(ambeginscan); + am_routine.amrescan = Some(amrescan); + am_routine.amgettuple = Some(amgettuple); + am_routine.amendscan = Some(amendscan); + + am_routine.amparallelvacuumoptions = pgrx::pg_sys::VACUUM_OPTION_PARALLEL_BULKDEL as u8 + | pgrx::pg_sys::VACUUM_OPTION_PARALLEL_CLEANUP as u8; + + am_routine +}; + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amvalidate(_opclass_oid: pgrx::pg_sys::Oid) -> bool { + true +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amoptions( + reloptions: Datum, + validate: bool, +) -> *mut pgrx::pg_sys::bytea { + let relopt_kind = RELOPT_KIND.get().copied().expect("init is not called"); + let rdopts = unsafe { + pgrx::pg_sys::build_reloptions( + reloptions, + validate, + relopt_kind, + size_of::(), + TABLE.as_ptr(), + TABLE.len() as _, + ) + }; + rdopts as *mut pgrx::pg_sys::bytea +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amcostestimate( + root: *mut pgrx::pg_sys::PlannerInfo, + path: *mut pgrx::pg_sys::IndexPath, + _loop_count: f64, + index_startup_cost: *mut pgrx::pg_sys::Cost, + index_total_cost: *mut pgrx::pg_sys::Cost, + index_selectivity: *mut pgrx::pg_sys::Selectivity, + index_correlation: *mut f64, + index_pages: *mut f64, +) { + unsafe { + use pgrx::pg_sys::disable_cost; + let index_opt_info = (*path).indexinfo; + // do not use index, if there are no orderbys or clauses + if ((*path).indexorderbys.is_null() && (*path).indexclauses.is_null()) + || !gucs::vchordrq_enable_scan() + { + *index_startup_cost = disable_cost; + *index_total_cost = disable_cost; + *index_selectivity = 0.0; + *index_correlation = 0.0; + *index_pages = 1.0; + return; + } + // Vchordrq indexes only the vector column, so ordinary filters on + // other columns are heap-side quals rather than index quals. We use + // the planner's filtered row estimate to decide how many distance + // candidates we likely need to produce enough survivors for LIMIT. + // + // Keep this separate from the value returned as `indexSelectivity`: + // `cost_index()` interprets that output as the fraction of parent + // table rows the index scan retrieves, not the fraction surviving all + // filters. With LIMIT, the retrieved fraction is better represented + // by the candidate count computed below. + let (total_rows, filter_selectivity) = { + let baserel = (*index_opt_info).rel; + let total_rows = (*baserel).tuples; + let param_info = (*path).path.param_info; + let filtered_rows = if !param_info.is_null() { + (*param_info).ppi_rows + } else { + (*baserel).rows + }; + let filter_selectivity = if total_rows > 0.0 { + (filtered_rows / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; + (total_rows, filter_selectivity) + }; + // index exists + if !(*index_opt_info).hypothetical { + let relation = Index::open((*index_opt_info).indexoid, pgrx::pg_sys::NoLock as _); + let opfamily = opfamily(relation.raw()); + let is_maxsim = matches!( + opfamily, + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim + ); + let index = PostgresRelation::::new(relation.raw()); + let probes = gucs::vchordrq_probes(relation.raw()); + let cost = vchordrq::cost(&index); + if cost.cells.len() != 1 + probes.len() { + panic!( + "need {} probes, but {} probes provided", + cost.cells.len() - 1, + probes.len() + ); + } + let estimated_index_vectors = if is_maxsim { + cost.indexed_vectors.map_or_else( + || { + (*index_opt_info).tuples.max(total_rows).max(1.0) + * f64::from(gucs::vchordrq_maxsim_planner_document_tokens()) + }, + |indexed_vectors| indexed_vectors as f64, + ) + } else { + (*index_opt_info).tuples.max(0.0) + }; + let node_count = { + let mut count = 0.0; + let r = cost.cells.iter().copied().rev().map(f64::from); + let numerator = std::iter::once(1.0).chain(probes.iter().copied().map(f64::from)); + let denumerator = r.clone(); + let scale = r.skip(1).chain(std::iter::once(estimated_index_vectors)); + for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { + count += scale * 1.0f64.min(numerator / denumerator); + } + count + }; + let page_count = { + let mut pages = 0_f64; + pages += 1.0; + pages += node_count * cost.dim as f64 / 60000.0; + pages += probes.iter().sum::() as f64 * { + let x = (opfamily.vector_kind().number_of_bits_of_an_elements() * cost.dim) + .div_ceil(8); + x.div_ceil(3840 * x.div_ceil(5120).min(2)) as f64 + }; + pages += cost.cells[0] as f64; + pages + }; + if is_maxsim { + let backend = match gucs::vchordrq_maxsim_backend() { + gucs::PostgresMaxsimBackend::CoarseOnly => { + vchordrq::MaxsimCostBackend::CoarseOnly + } + gucs::PostgresMaxsimBackend::CpuExact => vchordrq::MaxsimCostBackend::CpuExact, + gucs::PostgresMaxsimBackend::Gpu => vchordrq::MaxsimCostBackend::Gpu, + gucs::PostgresMaxsimBackend::Auto => vchordrq::MaxsimCostBackend::Auto, + }; + let estimate = vchordrq::estimate_maxsim_cost(vchordrq::MaxsimCostInput { + heap_rows: total_rows, + index_tokens: estimated_index_vectors, + token_nodes_per_query: node_count, + base_index_pages: page_count, + dimension: cost.dim, + element_bits: opfamily.vector_kind().number_of_bits_of_an_elements(), + query_tokens: gucs::vchordrq_maxsim_planner_query_tokens(), + limit_tuples: ((*root).limit_tuples > 0.0).then_some((*root).limit_tuples), + filter_selectivity, + candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + backend, + }); + *index_startup_cost = estimate.startup_cost; + *index_total_cost = estimate.total_cost; + *index_selectivity = estimate.selectivity; + *index_correlation = 0.0; + *index_pages = estimate.index_pages; + return; + } + // `next_count` represents candidates we expect to process to + // surface `limit_tuples` survivors after filter rejection. Clamp + // by `node_count` so the estimate cannot exceed the candidates + // the IVF visits at the configured probe count. + let next_count = if (*root).limit_tuples > 0.0 { + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)).min(node_count) + } else { + node_count + }; + let scan_selectivity = if total_rows > 0.0 { + (next_count / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; + *index_startup_cost = 0.001 * node_count; + *index_total_cost = 0.001 * node_count + next_count; + *index_selectivity = scan_selectivity; + *index_correlation = 0.0; + *index_pages = page_count; + return; + } + *index_startup_cost = 0.0; + *index_total_cost = 0.0; + *index_selectivity = filter_selectivity; + *index_correlation = 0.0; + *index_pages = 1.0; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn aminsert( + index_relation: pgrx::pg_sys::Relation, + values: *mut Datum, + is_null: *mut bool, + heap_tid: pgrx::pg_sys::ItemPointer, + _heap_relation: pgrx::pg_sys::Relation, + _check_unique: pgrx::pg_sys::IndexUniqueCheck::Type, + _index_unchanged: bool, + _index_info: *mut pgrx::pg_sys::IndexInfo, +) -> bool { + struct RngChooser(T); + impl InsertChooser for RngChooser { + fn choose(&mut self, n: NonZero) -> usize { + RngExt::random_range(&mut self.0, 0..n.get()) + } + } + + let opfamily = unsafe { opfamily(index_relation) }; + let index = unsafe { PostgresRelation::new(index_relation) }; + let datum = unsafe { (!is_null.add(0).read()).then_some(values.add(0).read()) }; + let ctid = unsafe { heap_tid.read() }; + if let Some(store) = unsafe { datum.and_then(|x| opfamily.store(x)) } { + for (vector, extra) in store { + let key = ctid_to_key(ctid); + let payload = kv_to_pointer((key, extra)); + let mut chooser = RngChooser(rand::rng()); + let bump = bumpalo::Bump::new(); + crate::index::vchordrq::dispatch::insert( + opfamily, + &index, + payload, + vector, + false, + false, + &mut chooser, + &bump, + ); + } + } + false +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambulkdelete( + info: *mut pgrx::pg_sys::IndexVacuumInfo, + stats: *mut pgrx::pg_sys::IndexBulkDeleteResult, + callback: pgrx::pg_sys::IndexBulkDeleteCallback, + callback_state: *mut std::os::raw::c_void, +) -> *mut pgrx::pg_sys::IndexBulkDeleteResult { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + let mut stats = stats; + if stats.is_null() { + stats = unsafe { + pgrx::pg_sys::palloc0(size_of::()).cast() + }; + } + let opfamily = unsafe { opfamily((*info).index) }; + let index = unsafe { PostgresRelation::new((*info).index) }; + let check = || unsafe { + #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16", feature = "pg17"))] + pgrx::pg_sys::vacuum_delay_point(); + #[cfg(feature = "pg18")] + pgrx::pg_sys::vacuum_delay_point(false); + }; + let callback = callback.expect("null function pointer"); + let callback = |pointer: NonZero| { + let (key, _) = pointer_to_kv(pointer); + let mut ctid = key_to_ctid(key); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + unsafe { + pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) + } + }; + let indexed_vectors = + crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + vchordrq::set_indexed_vectors(&index, indexed_vectors); + stats +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn ambeginscan( + index_relation: pgrx::pg_sys::Relation, + n_keys: std::os::raw::c_int, + n_orderbys: std::os::raw::c_int, +) -> pgrx::pg_sys::IndexScanDesc { + use pgrx::memcxt::PgMemoryContexts::CurrentMemoryContext; + + let scan = unsafe { pgrx::pg_sys::RelationGetIndexScan(index_relation, n_keys, n_orderbys) }; + let scanner: Scanner = Scanner { + hack: None, + scanning: LazyCell::new(Box::new(|| Box::new(std::iter::empty()))), + bump: Box::new(bumpalo::Bump::new()), + }; + unsafe { + (*scan).opaque = CurrentMemoryContext.leak_and_drop_on_delete(scanner).cast(); + } + scan +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amrescan( + scan: pgrx::pg_sys::IndexScanDesc, + keys: pgrx::pg_sys::ScanKey, + _n_keys: std::os::raw::c_int, + orderbys: pgrx::pg_sys::ScanKey, + _n_orderbys: std::os::raw::c_int, +) { + unsafe { + use crate::index::vchordrq::opclass::Opfamily; + if !keys.is_null() && (*scan).numberOfKeys > 0 { + std::ptr::copy(keys, (*scan).keyData, (*scan).numberOfKeys as _); + } + if !orderbys.is_null() && (*scan).numberOfOrderBys > 0 { + std::ptr::copy(orderbys, (*scan).orderByData, (*scan).numberOfOrderBys as _); + } + if (*scan).numberOfOrderBys == 0 && (*scan).numberOfKeys == 0 { + pgrx::error!( + "vector search with no WHERE clause and no ORDER BY clause is not supported" + ); + } + let scanner = &mut *(*scan).opaque.cast::(); + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); + let opfamily = opfamily((*scan).indexRelation); + let index = PostgresRelation::new((*scan).indexRelation); + let options = SearchOptions { + epsilon: gucs::vchordrq_epsilon((*scan).indexRelation), + probes: gucs::vchordrq_probes((*scan).indexRelation), + max_scan_tuples: gucs::vchordrq_max_scan_tuples(), + maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation), + maxsim_candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + maxsim_backend: gucs::vchordrq_maxsim_backend(), + maxsim_gpu_endpoint: gucs::vchordrq_maxsim_gpu_endpoint(), + maxsim_gpu_timeout_ms: gucs::vchordrq_maxsim_gpu_timeout_ms(), + maxsim_gpu_max_batch_tokens: gucs::vchordrq_maxsim_gpu_max_batch_tokens(), + maxsim_gpu_max_batch_bytes: gucs::vchordrq_maxsim_gpu_max_batch_bytes(), + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + prefilter: gucs::vchordrq_prefilter(), + }; + let fetcher = { + let hack = scanner.hack; + LazyCell::new(move || { + HeapFetcher::new( + (*scan).indexRelation, + (*scan).heapRelation, + (*scan).xs_snapshot, + (*scan).xs_heapfetch, + if let Some(hack) = hack { + hack.as_ptr() + } else { + std::ptr::null_mut() + }, + ) + }) + }; + let rate = match gucs::vchordrq_query_sampling_rate() { + 0.0 => None, + rate => Some(rate), + }; + let recorder = DefaultRecorder { + enable: gucs::vchordrq_query_sampling_enable(), + rate, + max_records: gucs::vchordrq_query_sampling_max_records(), + index: (*(*scan).indexRelation).rd_id.to_u32(), + }; + // PAY ATTENTATION: `scanning` references `bump`, so `scanning` must be dropped before `bump`. + let bump = scanner.bump.as_ref(); + scanner.scanning = match opfamily { + Opfamily::VectorL2 + | Opfamily::VectorIp + | Opfamily::VectorCosine + | Opfamily::HalfvecL2 + | Opfamily::HalfvecIp + | Opfamily::HalfvecCosine + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq4L2 + | Opfamily::Rabitq4Ip + | Opfamily::Rabitq4Cosine => { + let mut builder = DefaultBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump, recorder) + })) + } + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim => { + let mut builder = MaxsimBuilder::new(opfamily); + for i in 0..(*scan).numberOfOrderBys { + let data = (*scan).orderByData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + for i in 0..(*scan).numberOfKeys { + let data = (*scan).keyData.add(i as usize); + let value = (*data).sk_argument; + let is_null = ((*data).sk_flags & pgrx::pg_sys::SK_ISNULL as i32) != 0; + builder.add((*data).sk_strategy, (!is_null).then_some(value)); + } + LazyCell::new(Box::new(move || { + // only do this since `PostgresRelation` has no destructor + let index = bump.alloc(index.clone()); + builder.build(index, options, fetcher, bump, recorder) + })) + } + }; + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amgettuple( + scan: pgrx::pg_sys::IndexScanDesc, + direction: pgrx::pg_sys::ScanDirection::Type, +) -> bool { + if direction != pgrx::pg_sys::ScanDirection::ForwardScanDirection { + pgrx::error!("vector search without a forward scan direction is not supported"); + } + // https://www.postgresql.org/docs/current/index-locking.html + // If heap entries referenced physical pointers are deleted before + // they are consumed by PostgreSQL, PostgreSQL will received wrong + // physical pointers: no rows or irreverent rows are referenced. + if unsafe { (*(*scan).xs_snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + pgrx::error!("scanning with a non-MVCC-compliant snapshot is not supported"); + } + let scanner = unsafe { (*scan).opaque.cast::().as_mut().unwrap_unchecked() }; + if let Some((_, key, recheck)) = scanner.scanning.deref_mut().next() { + unsafe { + (*scan).xs_heaptid = key_to_ctid(key); + (*scan).xs_recheck = recheck; + (*scan).xs_recheckorderby = false; + } + true + } else { + false + } +} + +#[pgrx::pg_guard] +pub unsafe extern "C-unwind" fn amendscan(scan: pgrx::pg_sys::IndexScanDesc) { + let scanner = unsafe { &mut *(*scan).opaque.cast::() }; + scanner.scanning = LazyCell::new(Box::new(|| Box::new(std::iter::empty()))); + scanner.bump.reset(); +} + +type Iter = Box>; + +pub struct Scanner { + pub hack: Option>, + scanning: LazyCell Iter>>, + bump: Box, +} + +struct Index { + raw: *mut pgrx::pg_sys::RelationData, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl Index { + fn open(indexrelid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMASK) -> Self { + Self { + raw: unsafe { pgrx::pg_sys::index_open(indexrelid, lockmode) }, + lockmode, + } + } + fn raw(&self) -> *mut pgrx::pg_sys::RelationData { + self.raw + } +} + +impl Drop for Index { + fn drop(&mut self) { + unsafe { + pgrx::pg_sys::index_close(self.raw, self.lockmode); + } + } +} diff --git a/src/index/vchordrq/build.rs b/src/index/vchordrq/build.rs new file mode 100644 index 00000000..bbdc55ae --- /dev/null +++ b/src/index/vchordrq/build.rs @@ -0,0 +1,102 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use simd::{Floating, f16}; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; +use vector::{VectorBorrowed, VectorOwned}; + +pub type Normalized = Vec; + +pub trait Normalize: VectorOwned { + fn normalize(vector: Self) -> Normalized; + fn denormalize(x: Normalized) -> Self; +} + +impl Normalize for VectOwned { + fn normalize(vector: Self) -> Normalized { + vector.into_vec() + } + + fn denormalize(x: Normalized) -> Self { + Self::new(x) + } +} + +impl Normalize for VectOwned { + fn normalize(vector: Self) -> Normalized { + f16::vector_to_f32(vector.slice()) + } + + fn denormalize(x: Normalized) -> Self { + Self::new(f16::vector_from_f32(&x)) + } +} + +impl Normalize for Rabitq8Owned { + fn normalize(vector: Self) -> Normalized { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 8) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + result + } + + fn denormalize(vector: Normalized) -> Self { + let dim = vector.len() as u32; + let (metadata, elements) = rabitq::byte::ugly_code(&vector); + let elements = rabitq::byte::pack_code(&elements); + Rabitq8Owned::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + elements, + ) + } +} + +impl Normalize for Rabitq4Owned { + fn normalize(vector: Self) -> Normalized { + let vector = vector.as_borrowed(); + let scale = vector.sum_of_x2().sqrt() / vector.norm_of_lattice(); + let mut result = Vec::with_capacity(vector.dim() as _); + for c in vector.unpacked_code() { + let base = -0.5 * ((1 << 4) - 1) as f32; + result.push((base + c as f32) * scale); + } + rabitq::rotate::rotate_reversed_inplace(&mut result); + result + } + + fn denormalize(vector: Normalized) -> Self { + let dim = vector.len() as u32; + let (metadata, elements) = rabitq::halfbyte::ugly_code(&vector); + let elements = rabitq::halfbyte::pack_code(&elements); + Rabitq4Owned::new( + dim, + metadata.dis_u_2, + metadata.norm_of_lattice, + metadata.sum_of_code, + f32::reduce_sum_of_abs_x(&vector), + elements, + ) + } +} diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs new file mode 100644 index 00000000..36baee63 --- /dev/null +++ b/src/index/vchordrq/dispatch.rs @@ -0,0 +1,629 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::vchordrq::build::{Normalize, Normalized}; +use crate::index::vchordrq::opclass::Opfamily; +use index::bump::Bump; +use index::fetch::Fetch; +use index::prefetcher::*; +use index::relation::{ + Hints, Page, RelationPrefetch, RelationRead, RelationReadStream, RelationWrite, +}; +use index_accessor::{Dot, L2S}; +use simd::f16; +use std::collections::BinaryHeap; +use std::num::NonZero; +use vchordrq::operator::Op; +use vchordrq::types::*; +use vchordrq::{FastHeap, InsertChooser, MaintainChooser}; +use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::{VectBorrowed, VectOwned}; + +pub fn prewarm(opfamily: Opfamily, index: &R, height: i32) -> String +where + R: RelationRead, + R::Page: Page, +{ + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op, L2S>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op, Dot>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op, L2S>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op, Dot>>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + vchordrq::prewarm::<_, Op>(index, height, make_h0_plain_prefetcher) + } + } +} + +pub fn bulkdelete( + opfamily: Opfamily, + index: &R, + check: impl Fn(), + callback: impl Fn(NonZero) -> bool, +) -> u64 +where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live + } + } +} + +pub fn maintain( + opfamily: Opfamily, + index: &R, + chooser: &mut impl MaintainChooser, + check: impl Fn(), +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let maintain = match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + vchordrq::maintain::<_, Op, L2S>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + vchordrq::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + vchordrq::maintain::<_, Op, L2S>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + vchordrq::maintain::<_, Op, Dot>>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ) + } + (VectorKind::Rabitq8, DistanceKind::L2S) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + (VectorKind::Rabitq8, DistanceKind::Dot) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + (VectorKind::Rabitq4, DistanceKind::L2S) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + (VectorKind::Rabitq4, DistanceKind::Dot) => vchordrq::maintain::<_, Op>( + index, + make_h0_plain_prefetcher, + chooser, + check, + ), + }; + pgrx::debug1!( + "maintain: number_of_formerly_allocated_pages = {}", + maintain.number_of_formerly_allocated_pages + ); + pgrx::debug1!( + "maintain: number_of_freshly_allocated_pages = {}", + maintain.number_of_freshly_allocated_pages + ); + pgrx::debug1!( + "maintain: number_of_freed_pages = {}", + maintain.number_of_freed_pages + ); +} + +pub fn build( + vector_options: VectorOptions, + vchordrq_options: VchordrqIndexOptions, + index: &R, + structures: Vec>, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + match (vector_options.v, vector_options.d) { + (VectorKind::Vecf32, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Vecf32, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Vecf16, DistanceKind::L2S) => vchordrq::build::<_, Op, L2S>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Vecf16, DistanceKind::Dot) => vchordrq::build::<_, Op, Dot>>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq8, DistanceKind::L2S) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq8, DistanceKind::Dot) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq4, DistanceKind::L2S) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + (VectorKind::Rabitq4, DistanceKind::Dot) => vchordrq::build::<_, Op>( + vector_options, + vchordrq_options, + index, + map_structures(structures, Normalize::denormalize), + ), + } +} + +pub fn insert( + opfamily: Opfamily, + index: &R, + payload: NonZero, + vector: OwnedVector, + skip_freespaces: bool, + skip_search: bool, + chooser: &mut impl InsertChooser, + bump: &impl Bump, +) where + R: RelationRead + RelationWrite, + R::Page: Page, +{ + let make_h1_plain_prefetcher = MakeH1PlainPrefetcherForInsertion { index }; + match (vector, opfamily.distance_kind()) { + (OwnedVector::Vecf32(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, L2S>>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op, L2S>>( + index, + payload, + projected.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf32(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf32); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op, Dot>>( + index, + payload, + projected.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, L2S>>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op, L2S>>( + index, + payload, + projected.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Vecf16(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Vecf16); + let projected = RandomProject::project(vector.as_borrowed()); + let key = vchordrq::insert_vector::<_, Op, Dot>>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op, Dot>>( + index, + payload, + projected.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq8(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq8(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq8); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq4(vector), DistanceKind::Dot) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + (OwnedVector::Rabitq4(vector), DistanceKind::L2S) => { + assert!(opfamily.vector_kind() == VectorKind::Rabitq4); + let key = vchordrq::insert_vector::<_, Op>( + index, + payload, + vector.as_borrowed(), + chooser, + skip_search, + ); + vchordrq::insert::<_, Op>( + index, + payload, + vector.as_borrowed(), + key, + bump, + make_h1_plain_prefetcher, + skip_freespaces, + ) + } + } +} + +fn map_structures(x: Vec>, f: impl Fn(T) -> U + Copy) -> Vec> { + x.into_iter() + .map( + |Structure { + centroids, + children, + }| Structure { + centroids: centroids.into_iter().map(f).collect(), + children, + }, + ) + .collect() +} + +pub trait RandomProject { + type Output; + fn project(self) -> Self::Output; +} + +impl RandomProject for VectBorrowed<'_, f32> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + let input = self.slice(); + VectOwned::new(rotate(input)) + } +} + +impl RandomProject for VectBorrowed<'_, f16> { + type Output = VectOwned; + fn project(self) -> VectOwned { + use rabitq::rotate::rotate; + use simd::Floating; + let input = f16::vector_to_f32(self.slice()); + VectOwned::new(f16::vector_from_f32(&rotate(&input))) + } +} + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcherForInsertion<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakeH1PlainPrefetcherForInsertion<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead> PrefetcherHeapFamily<'b, R> for MakeH1PlainPrefetcherForInsertion<'b, R> { + type P + = PlainPrefetcher<'b, R, FastHeap> + where + T: Ord + Fetch<'b>; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch<'b>, + { + PlainPrefetcher::new(self.index, FastHeap::from(seq)) + } + + fn is_not_plain(&self) -> bool { + false + } +} + +#[derive(Debug)] +pub struct MakeH1PlainPrefetcher<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakeH1PlainPrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead> PrefetcherHeapFamily<'b, R> for MakeH1PlainPrefetcher<'b, R> { + type P + = PlainPrefetcher<'b, R, BinaryHeap> + where + T: Ord + Fetch<'b>; + + fn prefetch(&mut self, seq: Vec) -> Self::P + where + T: Ord + Fetch<'b>, + { + PlainPrefetcher::new(self.index, BinaryHeap::from(seq)) + } + + fn is_not_plain(&self) -> bool { + false + } +} + +#[derive(Debug)] +pub struct MakeH0PlainPrefetcher<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakeH0PlainPrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead> PrefetcherSequenceFamily<'b, R> for MakeH0PlainPrefetcher<'b, R> { + type P + = PlainPrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + PlainPrefetcher::new(self.index, seq) + } + + fn is_not_plain(&self) -> bool { + false + } +} + +#[derive(Debug)] +pub struct MakeH0SimplePrefetcher<'b, R> { + pub index: &'b R, +} + +impl<'b, R> Clone for MakeH0SimplePrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { index: self.index } + } +} + +impl<'b, R: RelationRead + RelationPrefetch> PrefetcherSequenceFamily<'b, R> + for MakeH0SimplePrefetcher<'b, R> +{ + type P + = SimplePrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + SimplePrefetcher::new(self.index, seq) + } + + fn is_not_plain(&self) -> bool { + true + } +} + +#[derive(Debug)] +pub struct MakeH0StreamPrefetcher<'b, R> { + pub index: &'b R, + pub hints: Hints, +} + +impl<'b, R> Clone for MakeH0StreamPrefetcher<'b, R> { + fn clone(&self) -> Self { + Self { + index: self.index, + hints: self.hints, + } + } +} + +impl<'b, R: RelationRead + RelationReadStream> PrefetcherSequenceFamily<'b, R> + for MakeH0StreamPrefetcher<'b, R> +{ + type P + = StreamPrefetcher<'b, R, S> + where + S::Item: Fetch<'b>; + + fn prefetch(&mut self, seq: S) -> Self::P + where + S::Item: Fetch<'b>, + { + StreamPrefetcher::new(self.index, seq, self.hints) + } + + fn is_not_plain(&self) -> bool { + true + } +} diff --git a/src/index/vchordrq/filter.rs b/src/index/vchordrq/filter.rs new file mode 100644 index 00000000..89470f5b --- /dev/null +++ b/src/index/vchordrq/filter.rs @@ -0,0 +1,55 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use index::prefetcher::Sequence; + +pub struct Filter { + sequence: S, + predicate: P, +} + +impl Sequence for Filter +where + S: Sequence, + P: FnMut(&S::Item) -> bool, +{ + type Item = S::Item; + + type Inner = S::Inner; + + fn next(&mut self) -> Option { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.next() + } + + fn peek(&mut self) -> Option<&Self::Item> { + while !(self.predicate)(self.sequence.peek()?) { + let _ = self.sequence.next(); + } + self.sequence.peek() + } + + fn into_inner(self) -> Self::Inner { + self.sequence.into_inner() + } +} + +pub fn filter(sequence: S, predicate: P) -> Filter { + Filter { + sequence, + predicate, + } +} diff --git a/src/index/vchordrq/mod.rs b/src/index/vchordrq/mod.rs new file mode 100644 index 00000000..38d75856 --- /dev/null +++ b/src/index/vchordrq/mod.rs @@ -0,0 +1,21 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub mod am; +mod build; +pub mod dispatch; +mod filter; +pub mod opclass; +mod scanners; +pub mod types; diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs new file mode 100644 index 00000000..075b1075 --- /dev/null +++ b/src/index/vchordrq/opclass.rs @@ -0,0 +1,368 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::datatype::memory_halfvec::{HalfvecInput, HalfvecOutput}; +use crate::datatype::memory_rabitq4::{Rabitq4Input, Rabitq4Output}; +use crate::datatype::memory_rabitq8::{Rabitq8Input, Rabitq8Output}; +use crate::datatype::memory_vector::{VectorInput, VectorOutput}; +use crate::index::opclass::Sphere; +use distance::Distance; +use pgrx::datum::FromDatum; +use pgrx::heap_tuple::PgHeapTuple; +use pgrx::pg_sys::Datum; +use std::num::NonZero; +use vchordrq::types::*; +use vector::VectorBorrowed; + +#[derive(Debug, Clone, Copy)] +pub enum Opfamily { + VectorL2, + VectorIp, + VectorCosine, + HalfvecL2, + HalfvecIp, + HalfvecCosine, + Rabitq8L2, + Rabitq8Ip, + Rabitq8Cosine, + Rabitq4L2, + Rabitq4Ip, + Rabitq4Cosine, + VectorMaxsim, + HalfvecMaxsim, + Rabitq8Maxsim, + Rabitq4Maxsim, +} + +impl Opfamily { + fn input(self, vector: BorrowedVector<'_>) -> OwnedVector { + use BorrowedVector as B; + use OwnedVector as O; + match (vector, self) { + (B::Vecf32(x), Self::VectorL2) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorIp | Self::VectorMaxsim) => O::Vecf32(x.own()), + (B::Vecf32(x), Self::VectorCosine) => O::Vecf32(x.function_normalize()), + (B::Vecf32(_), _) => unreachable!(), + (B::Vecf16(x), Self::HalfvecL2) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecIp | Self::HalfvecMaxsim) => O::Vecf16(x.own()), + (B::Vecf16(x), Self::HalfvecCosine) => O::Vecf16(x.function_normalize()), + (B::Vecf16(_), _) => unreachable!(), + (B::Rabitq8(x), Self::Rabitq8L2) => O::Rabitq8(x.own()), + (B::Rabitq8(x), Self::Rabitq8Ip | Self::Rabitq8Maxsim) => O::Rabitq8(x.own()), + (B::Rabitq8(x), Self::Rabitq8Cosine) => O::Rabitq8(x.function_normalize()), + (B::Rabitq8(_), _) => unreachable!(), + (B::Rabitq4(x), Self::Rabitq4L2) => O::Rabitq4(x.own()), + (B::Rabitq4(x), Self::Rabitq4Ip | Self::Rabitq4Maxsim) => O::Rabitq4(x.own()), + (B::Rabitq4(x), Self::Rabitq4Cosine) => O::Rabitq4(x.function_normalize()), + (B::Rabitq4(_), _) => unreachable!(), + } + } + pub unsafe fn store(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let store = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf32(vector.as_borrowed())), 0)] + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Vecf16(vector.as_borrowed())), 0)] + } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), 0)] + } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + vec![(self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), 0)] + } + Self::VectorMaxsim => { + let vectors = + unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf32(vector.as_borrowed())), + i as u16, + )); + } + result + } + Self::HalfvecMaxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Vecf16(vector.as_borrowed())), + i as u16, + )); + } + result + } + Self::Rabitq8Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())), + i as u16, + )); + } + result + } + Self::Rabitq4Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for (i, vector) in vectors.iter_deny_null().enumerate() { + result.push(( + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())), + i as u16, + )); + } + result + } + }; + Some(store) + } + pub unsafe fn input_sphere(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let attno_1 = NonZero::new(1_usize).unwrap(); + let attno_2 = NonZero::new(2_usize).unwrap(); + let tuple = unsafe { PgHeapTuple::from_composite_datum(datum) }; + let center = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vector = tuple.get_by_index::(attno_1).unwrap()?; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } + }; + let radius = tuple.get_by_index::(attno_2).unwrap()?; + Some(Sphere { center, radius }) + } + pub unsafe fn input_vector(self, datum: Datum) -> Option { + if datum.is_null() { + return None; + } + let vector = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vector = unsafe { VectorInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf32(vector.as_borrowed())) + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vector = unsafe { HalfvecInput::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Vecf16(vector.as_borrowed())) + } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vector = unsafe { Rabitq8Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq8(vector.as_borrowed())) + } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vector = unsafe { Rabitq4Input::from_datum(datum, false).unwrap() }; + self.input(BorrowedVector::Rabitq4(vector.as_borrowed())) + } + }; + Some(vector) + } + pub unsafe fn input_vectors(self, datum: Datum) -> Option> { + if datum.is_null() { + return None; + } + let vectors = match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + let vectors = + unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); + } + result + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); + } + result + } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Rabitq8(vector.as_borrowed()))); + } + result + } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + let vectors = unsafe { + pgrx::datum::Array::::from_datum(datum, false).unwrap() + }; + crate::datatype::validate_maxsim_array_len(vectors.len()); + let mut result = Vec::with_capacity(vectors.len()); + for vector in vectors.iter_deny_null() { + result.push(self.input(BorrowedVector::Rabitq4(vector.as_borrowed()))); + } + result + } + }; + Some(vectors) + } + pub fn output(self, x: Distance) -> f32 { + match self { + Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine => x.to_f32() + 1.0f32, + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + x.to_f32().sqrt() + } + Self::VectorIp + | Self::HalfvecIp + | Self::Rabitq8Ip + | Self::Rabitq4Ip + | Self::VectorMaxsim + | Self::HalfvecMaxsim + | Self::Rabitq8Maxsim + | Self::Rabitq4Maxsim => x.to_f32(), + } + } + pub const fn distance_kind(self) -> DistanceKind { + match self { + Self::VectorL2 | Self::HalfvecL2 | Self::Rabitq8L2 | Self::Rabitq4L2 => { + DistanceKind::L2S + } + Self::VectorIp + | Self::HalfvecIp + | Self::Rabitq8Ip + | Self::Rabitq4Ip + | Self::VectorCosine + | Self::HalfvecCosine + | Self::Rabitq8Cosine + | Self::Rabitq4Cosine + | Self::VectorMaxsim + | Self::HalfvecMaxsim + | Self::Rabitq8Maxsim + | Self::Rabitq4Maxsim => DistanceKind::Dot, + } + } + pub const fn vector_kind(self) -> VectorKind { + match self { + Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { + VectorKind::Vecf32 + } + Self::HalfvecL2 | Self::HalfvecIp | Self::HalfvecCosine | Self::HalfvecMaxsim => { + VectorKind::Vecf16 + } + Self::Rabitq8L2 | Self::Rabitq8Ip | Self::Rabitq8Cosine | Self::Rabitq8Maxsim => { + VectorKind::Rabitq8 + } + Self::Rabitq4L2 | Self::Rabitq4Ip | Self::Rabitq4Cosine | Self::Rabitq4Maxsim => { + VectorKind::Rabitq4 + } + } + } +} + +pub unsafe fn opfamily(index_relation: pgrx::pg_sys::Relation) -> Opfamily { + use pgrx::pg_sys::Oid; + + let proc = unsafe { pgrx::pg_sys::index_getprocid(index_relation, 1, 1) }; + + if proc == Oid::INVALID { + pgrx::error!("support function 1 is not found"); + } + + let mut flinfo = pgrx::pg_sys::FmgrInfo::default(); + + unsafe { + pgrx::pg_sys::fmgr_info(proc, &mut flinfo); + } + + let fn_addr = flinfo.fn_addr.expect("null function pointer"); + + let mut fcinfo = unsafe { std::mem::zeroed::() }; + fcinfo.flinfo = &mut flinfo; + fcinfo.fncollation = pgrx::pg_sys::DEFAULT_COLLATION_OID; + fcinfo.context = std::ptr::null_mut(); + fcinfo.resultinfo = std::ptr::null_mut(); + fcinfo.isnull = true; + fcinfo.nargs = 0; + + let result_datum = unsafe { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pgrx::pg_sys::ffi::pg_guard_ffi_boundary(|| fn_addr(&mut fcinfo)) + }; + + let result_option = unsafe { String::from_datum(result_datum, fcinfo.isnull) }; + + let result_string = result_option.expect("null return value"); + + let result = match result_string.as_str() { + "vchordrq_vector_l2_ops" => Opfamily::VectorL2, + "vchordrq_vector_ip_ops" => Opfamily::VectorIp, + "vchordrq_vector_cosine_ops" => Opfamily::VectorCosine, + "vchordrq_halfvec_l2_ops" => Opfamily::HalfvecL2, + "vchordrq_halfvec_ip_ops" => Opfamily::HalfvecIp, + "vchordrq_halfvec_cosine_ops" => Opfamily::HalfvecCosine, + "vchordrq_rabitq8_l2_ops" => Opfamily::Rabitq8L2, + "vchordrq_rabitq8_ip_ops" => Opfamily::Rabitq8Ip, + "vchordrq_rabitq8_cosine_ops" => Opfamily::Rabitq8Cosine, + "vchordrq_rabitq4_l2_ops" => Opfamily::Rabitq4L2, + "vchordrq_rabitq4_ip_ops" => Opfamily::Rabitq4Ip, + "vchordrq_rabitq4_cosine_ops" => Opfamily::Rabitq4Cosine, + "vchordrq_vector_maxsim_ops" => Opfamily::VectorMaxsim, + "vchordrq_halfvec_maxsim_ops" => Opfamily::HalfvecMaxsim, + "vchordrq_rabitq8_maxsim_ops" => Opfamily::Rabitq8Maxsim, + "vchordrq_rabitq4_maxsim_ops" => Opfamily::Rabitq4Maxsim, + _ => pgrx::error!("unknown operator class"), + }; + + unsafe { + pgrx::pg_sys::pfree(result_datum.cast_mut_ptr()); + } + + result +} diff --git a/src/index/vchordrq/scanners/default.rs b/src/index/vchordrq/scanners/default.rs new file mode 100644 index 00000000..d91b4273 --- /dev/null +++ b/src/index/vchordrq/scanners/default.rs @@ -0,0 +1,1274 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::*; +use crate::index::opclass::Sphere; +use crate::index::scanners::{Io, SearchBuilder}; +use crate::index::vchordrq::dispatch::*; +use crate::index::vchordrq::filter::filter; +use crate::index::vchordrq::opclass::Opfamily; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::{Recorder, text}; +use always_equal::AlwaysEqual; +use dary_heap::QuaternaryHeap as Heap; +use index::bump::Bump; +use index::packed::PackedRefMut4; +use index::prefetcher::*; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::{Dot, L2S}; +use simd::f16; +use std::num::NonZero; +use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; +use vchordrq::{RerankMethod, default_search, how, rerank_heap, rerank_index}; +use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; + +pub struct DefaultBuilder { + opfamily: Opfamily, + orderbys: Vec>, + spheres: Vec>>, +} + +impl SearchBuilder for DefaultBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + + type Opaque = vchordrq::Opaque; + + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::HalfvecCosine + | Opfamily::HalfvecIp + | Opfamily::HalfvecL2 + | Opfamily::VectorCosine + | Opfamily::VectorIp + | Opfamily::VectorL2 + | Opfamily::Rabitq8Cosine + | Opfamily::Rabitq8Ip + | Opfamily::Rabitq8L2 + | Opfamily::Rabitq4Cosine + | Opfamily::Rabitq4Ip + | Opfamily::Rabitq4L2 + )); + Self { + opfamily, + orderbys: Vec::new(), + spheres: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 1 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vector(x)) }; + self.orderbys.push(x); + } + 2 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_sphere(x)) }; + self.spheres.push(x); + } + _ => unreachable!(), + } + } + + fn build<'b, R>( + self, + index: &'b R, + options: SearchOptions, + mut fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + recorder: impl Recorder, + ) -> Box + 'b> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, + { + let mut vector = None; + let mut threshold = None; + let mut recheck = false; + for orderby_vector in self.orderbys.into_iter().flatten() { + if vector.is_none() { + vector = Some(orderby_vector); + } else { + pgrx::error!("vector search with multiple vectors is not supported"); + } + } + for Sphere { center, radius } in self.spheres.into_iter().flatten() { + if vector.is_none() { + (vector, threshold) = (Some(center), Some(radius)); + } else { + recheck = true; + } + } + let opfamily = self.opfamily; + let Some(vector) = vector else { + return Box::new(std::iter::empty()) as Box>; + }; + let search_hints = Hints::default().full(true); + let rerank_hints = Hints::default().full(false); + let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; + let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { + index, + hints: search_hints, + }; + let f = move |(distance, payload)| (opfamily.output(distance), payload); + let iter: Box)>> = + match (opfamily.vector_kind(), opfamily.distance_kind()) { + (VectorKind::Vecf32, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let projected = RandomProject::project(unprojected.as_borrowed()); + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Vecf32, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf32(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let projected = RandomProject::project(unprojected.as_borrowed()); + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf32(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Vecf16, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op, L2S>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let projected = RandomProject::project(unprojected.as_borrowed()); + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Vecf16, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = if let OwnedVector::Vecf16(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let projected = RandomProject::project(unprojected.as_borrowed()); + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + projected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = if let OwnedVector::Vecf16(vector) = maybe_vector.unwrap() + { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq8, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq8, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq8(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq8(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq4, DistanceKind::L2S) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + (VectorKind::Rabitq4, DistanceKind::Dot) => { + type Op = vchordrq::operator::Op; + let unprojected = if let OwnedVector::Rabitq4(vector) = vector.clone() { + vector + } else { + unreachable!() + }; + let results = match options.io_search { + Io::Plain => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_plain_prefetcher, + ), + Io::Simple => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_simple_prefetcher, + ), + Io::Stream => default_search::<_, Op>( + index, + unprojected.as_borrowed(), + options.probes, + options.epsilon, + bump, + make_h1_plain_prefetcher, + make_h0_stream_prefetcher, + ), + }; + let method = how(index); + let sequence = Heap::from(results); + match (method, options.io_rerank, options.prefilter) { + (RerankMethod::Index, Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Plain, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Simple, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, false) => { + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Index, Io::Stream, true) => { + let predicate = + id_0(move |(_, AlwaysEqual(PackedRefMut4((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = StreamPrefetcher::new(index, sequence, rerank_hints); + Box::new(rerank_index::(unprojected, prefetcher).map(f)) + } + (RerankMethod::Heap, _, false) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + (RerankMethod::Heap, _, true) => { + let fetch = move |payload| { + let (key, _) = pointer_to_kv(payload); + let mut tuple = fetcher.fetch(key)?; + if !tuple.filter() { + return None; + } + let (datums, is_nulls) = tuple.build(); + let datum = (!is_nulls[0]).then_some(datums[0]); + let maybe_vector = + unsafe { datum.and_then(|x| opfamily.input_vector(x)) }; + let raw = + if let OwnedVector::Rabitq4(vector) = maybe_vector.unwrap() { + vector + } else { + unreachable!() + }; + Some(raw) + }; + let prefetcher = PlainPrefetcher::new(index, sequence); + Box::new( + rerank_heap::(unprojected, prefetcher, fetch).map(f), + ) + } + } + } + }; + let iter = if let Some(threshold) = threshold { + Box::new(iter.take_while(move |(x, _)| *x < threshold)) + } else { + iter + }; + let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { + Box::new(iter.take(max_scan_tuples as _)) + } else { + iter + }; + if recorder.is_enabled() { + match &vector { + OwnedVector::Vecf32(v) => { + recorder.send(&text::vector_out(v.as_borrowed())); + } + OwnedVector::Vecf16(v) => { + recorder.send(&text::halfvec_out(v.as_borrowed())); + } + OwnedVector::Rabitq8(v) => { + recorder.send(&text::rabitq8_out(v.as_borrowed())); + } + OwnedVector::Rabitq4(v) => { + recorder.send(&text::rabitq4_out(v.as_borrowed())); + } + } + } + Box::new(iter.map(move |(distance, pointer)| { + let (key, _) = pointer_to_kv(pointer); + (distance, key, recheck) + })) + } +} + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&(A, AlwaysEqual>)) -> R, +{ + f +} diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs new file mode 100644 index 00000000..1e0fe486 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -0,0 +1,960 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod candidate; +mod exact; +mod external; +mod gpu; +mod profile; +mod rerank; +mod search; + +use self::candidate::{DensePageCandidateGenerator, PageCandidateGenerator}; +use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; +use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; +use crate::index::fetcher::*; +use crate::index::gucs::PostgresMaxsimBackend; +use crate::index::scanners::{Io, SearchBuilder}; +use crate::index::vchordrq::dispatch::*; +use crate::index::vchordrq::filter::filter; +use crate::index::vchordrq::opclass::Opfamily; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::Recorder; +use always_equal::AlwaysEqual; +use dary_heap::QuaternaryHeap as Heap; +use index::bump::Bump; +use index::packed::PackedRefMut8; +use index::prefetcher::*; +use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadStream}; +use index_accessor::Dot; +use simd::f16; +use std::cmp::Reverse; +use std::num::NonZero; +use std::time::Duration; +use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; +use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; +use vector::VectorOwned; +use vector::rabitq4::Rabitq4Owned; +use vector::rabitq8::Rabitq8Owned; +use vector::vect::VectOwned; + +fn retain_top_k(best: &mut std::collections::BinaryHeap, limit: usize, item: T) { + if limit == 0 { + return; + } + if best.len() < limit { + best.push(item); + } else if best.peek().is_some_and(|worst| item < *worst) { + best.pop(); + best.push(item); + } +} + +pub struct MaxsimBuilder { + opfamily: Opfamily, + orderbys: Vec>>, +} + +impl SearchBuilder for MaxsimBuilder { + type Options = SearchOptions; + + type Opfamily = Opfamily; + + type Opaque = vchordrq::Opaque; + + fn new(opfamily: Opfamily) -> Self { + assert!(matches!( + opfamily, + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim + )); + Self { + opfamily, + orderbys: Vec::new(), + } + } + + unsafe fn add(&mut self, strategy: u16, datum: Option) { + match strategy { + 3 => { + let x = unsafe { datum.and_then(|x| self.opfamily.input_vectors(x)) }; + self.orderbys.push(x); + } + _ => unreachable!(), + } + } + + fn build<'b, R>( + self, + index: &'b R, + options: SearchOptions, + mut fetcher: impl Fetcher + 'b, + bump: &'b impl Bump, + _sender: impl Recorder, + ) -> Box + 'b> + where + R: RelationRead + RelationPrefetch + RelationReadStream, + R::Page: Page, + { + let mut vectors = None; + for orderby_vectors in self.orderbys.into_iter().flatten() { + if vectors.is_none() { + vectors = Some(orderby_vectors); + } else { + pgrx::error!("maxsim search with multiple vectors is not supported"); + } + } + if let Some(_max_scan_tuples) = options.max_scan_tuples { + pgrx::error!("maxsim search with max_scan_tuples is not supported"); + } + let maxsim_refine = options.maxsim_refine; + let maxsim_threshold = options.maxsim_threshold; + let maxsim_candidate_limit = options + .maxsim_candidate_limit + .map_or(usize::MAX, |value| value as usize); + let has_maxsim_candidate_limit = options.maxsim_candidate_limit.is_some(); + let maxsim_backend = options.maxsim_backend; + let maxsim_gpu_endpoint = options + .maxsim_gpu_endpoint + .as_deref() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let maxsim_gpu_timeout = Duration::from_millis(options.maxsim_gpu_timeout_ms as u64); + let maxsim_gpu_max_batch_tokens = options.maxsim_gpu_max_batch_tokens as usize; + let maxsim_gpu_max_batch_bytes = options.maxsim_gpu_max_batch_bytes as usize; + if !matches!(maxsim_backend, PostgresMaxsimBackend::CoarseOnly) + && (!has_maxsim_candidate_limit || maxsim_candidate_limit == 0) + { + pgrx::error!("exact MaxSim requires a positive vchordrq.maxsim_candidate_limit"); + } + let opfamily = self.opfamily; + let Some(vectors) = vectors else { + return Box::new(std::iter::empty()) as Box>; + }; + let expected_dim = vchordrq::cost(index).dim; + if vectors.iter().any(|vector| vector.dim() != expected_dim) { + pgrx::error!("dimension is not matched"); + } + let rerank_query = vectors.clone(); + let method = how(index); + if !matches!(method, RerankMethod::Index) { + pgrx::error!("maxsim search with rerank_in_table is not supported"); + } + assert!(matches!(opfamily.distance_kind(), DistanceKind::Dot)); + let search_hints = Hints::default().full(true); + let rerank_hints = Hints::default().full(false); + let make_h1_plain_prefetcher = MakeH1PlainPrefetcher { index }; + let make_h0_plain_prefetcher = MakeH0PlainPrefetcher { index }; + let make_h0_simple_prefetcher = MakeH0SimplePrefetcher { index }; + let make_h0_stream_prefetcher = MakeH0StreamPrefetcher { + index, + hints: search_hints, + }; + let n = vectors.len(); + let accu_map = |(Reverse(distance), AlwaysEqual(payload))| (distance, payload); + let rough_map = + |((_, AlwaysEqual(rough)), AlwaysEqual(PackedRefMut8(&mut (payload, ..)))): ( + _, + AlwaysEqual, _, _)>>, + )| (rough, payload); + let mut token_searches = { + let fetcher = &mut fetcher; + let iter: Box> = match opfamily.vector_kind() { + VectorKind::Vecf32 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Vecf16 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq8 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq8(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq4 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq4(vector) = vector { + vector + } else { + unreachable!() + } + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let token_search_timer = profile::ProfileTimer::start(); + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let refine_active = maxsim_refine != 0 && !results.is_empty(); + let search_results = results.len() as u64; + let token_search_elapsed = token_search_timer.elapsed(); + profile::update(|profile| { + profile.token_search_calls += 1; + profile.token_search_us += profile::duration_us(token_search_elapsed); + profile.token_search_results += search_results; + }); + let token_refine_timer = profile::ProfileTimer::start(); + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); + } + let token_refine_elapsed = token_refine_timer.elapsed(); + profile::update(|profile| { + profile.token_refine_calls += u64::from(refine_active); + profile.token_refine_us += profile::duration_us(token_refine_elapsed); + profile.accurate_token_hits += accu_set.len() as u64; + profile.rough_token_hits += rough_set.len() as u64; + }); + (accu_set, rough_set, estimation_by_threshold) + })) + } + }; + iter + }; + let mut candidates = + DensePageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + drop(token_searches); + let iter: Box> = match maxsim_backend { + PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { + let distance = candidate.approximate_distance.to_f32(); + let recheck = false; + (distance, candidate.heap_key, recheck) + })), + PostgresMaxsimBackend::CpuExact => { + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let exact = + CpuExactMaxsimBackend.rerank(&rerank_query, &mut candidates, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) + } + PostgresMaxsimBackend::Gpu => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) + } + PostgresMaxsimBackend::Auto => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = match exact { + Ok(exact) => exact, + Err(error) => { + report_gpu_fallback(&error); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + CpuExactMaxsimBackend + .rerank(&rerank_query, &mut candidate_iter, &mut source) + .unwrap_or_else(|error| pgrx::error!("{error}")) + } + }; + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) + } + }; + #[allow(clippy::let_and_return)] + iter + } +} + +#[inline(always)] +pub fn id_0(f: F) -> F +where + F: for<'a> FnMut(&(A, AlwaysEqual>)) -> R, +{ + f +} + +#[cfg(test)] +mod tests { + use super::retain_top_k; + use std::collections::BinaryHeap; + + #[test] + fn bounded_top_k_keeps_global_order_and_ties() { + let mut best = BinaryHeap::new(); + for item in [(3, 30), (1, 20), (1, 10), (2, 40), (0, 50)] { + retain_top_k(&mut best, 3, item); + } + let mut rows = best.into_vec(); + rows.sort_unstable(); + assert_eq!(rows, vec![(0, 50), (1, 10), (1, 20)]); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs new file mode 100644 index 00000000..af9b45b0 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -0,0 +1,254 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::profile; +use crate::index::fetcher::pointer_to_kv; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::hash_map::Entry; +use std::collections::{BinaryHeap, HashMap}; +use std::num::NonZero; + +pub(super) type HeapKey = [u16; 3]; +pub(super) type TokenCandidate = (Distance, NonZero); +pub(super) type TokenSearchResult = (Vec, Vec, Distance); + +#[derive(Clone, Copy, Debug)] +pub(super) struct PageCandidate { + pub approximate_distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) trait PageCandidateGenerator { + type Candidates: Iterator; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates; +} + +#[derive(Default)] +pub(super) struct DensePageCandidateGenerator; + +impl PageCandidateGenerator for DensePageCandidateGenerator { + type Candidates = PageCandidates; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates { + let mut page_lookup = HashMap::new(); + let mut page_keys = Vec::new(); + let mut best_by_page_query = Vec::new(); + let mut estimations = Vec::with_capacity(query_count); + let mut hit_updates = 0u64; + let mut page_token_updates = 0u64; + for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() + { + debug_assert!(query_id < query_count); + let hit_collect_timer = profile::ProfileTimer::start(); + hit_updates += (accu_set.len() + rough_set.len()) as u64; + let is_empty = accu_set.is_empty() && rough_set.is_empty(); + let mut estimation_by_scope = Distance::NEG_INFINITY; + for (distance, payload) in accu_set { + estimation_by_scope = std::cmp::max(estimation_by_scope, distance); + let (key, _) = pointer_to_kv(payload); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); + } + for (distance, payload) in rough_set { + let (key, _) = pointer_to_kv(payload); + page_token_updates += u64::from(record_page_token_distance( + &mut page_lookup, + &mut page_keys, + &mut best_by_page_query, + query_count, + query_id, + key, + distance, + )); + } + estimations.push(if !is_empty { + std::cmp::max(estimation_by_scope, estimation_by_threshold) + } else { + Distance::ZERO + }); + let hit_collect_elapsed = hit_collect_timer.elapsed(); + profile::update(|profile| { + profile.hit_collect_us += profile::duration_us(hit_collect_elapsed); + }); + } + debug_assert_eq!(estimations.len(), query_count); + profile::update(|profile| { + profile.hit_updates += hit_updates; + profile.page_token_updates += page_token_updates; + }); + let page_aggregate_timer = profile::ProfileTimer::start(); + let mut page_order = (0..page_keys.len()).collect::>(); + page_order.sort_unstable_by_key(|&page_index| page_keys[page_index]); + let inner = page_order + .into_iter() + .map(|page_index| { + let key = page_keys[page_index]; + let start = page_index * query_count; + let values = &best_by_page_query[start..start + query_count]; + let mut maxsim = 0.0f32; + for (query_id, distance) in values.iter().copied().enumerate() { + let distance = distance.unwrap_or(estimations[query_id]); + maxsim += distance.to_f32(); + } + (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) + }) + .collect::>(); + let page_aggregate_elapsed = page_aggregate_timer.elapsed(); + profile::update(|profile| { + profile.page_aggregate_us += profile::duration_us(page_aggregate_elapsed); + profile.aggregated_pages += page_keys.len() as u64; + }); + PageCandidates { + inner, + remaining: candidate_limit, + } + } +} + +fn record_page_token_distance( + page_lookup: &mut HashMap, + page_keys: &mut Vec, + best_by_page_query: &mut Vec>, + query_count: usize, + query_id: usize, + key: HeapKey, + distance: Distance, +) -> bool { + let page_index = match page_lookup.entry(key) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + let page_index = page_keys.len(); + entry.insert(page_index); + page_keys.push(key); + best_by_page_query.resize(best_by_page_query.len() + query_count, None); + page_index + } + }; + let slot = &mut best_by_page_query[page_index * query_count + query_id]; + let first_for_page_token = slot.is_none(); + let best = slot.get_or_insert(Distance::INFINITY); + *best = std::cmp::min(*best, distance); + first_for_page_token +} + +pub(super) struct PageCandidates { + inner: BinaryHeap<(Reverse, AlwaysEqual)>, + remaining: usize, +} + +impl Iterator for PageCandidates { + type Item = PageCandidate; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + let (Reverse(approximate_distance), AlwaysEqual(heap_key)) = self.inner.pop()?; + self.remaining -= 1; + Some(PageCandidate { + approximate_distance, + heap_key, + }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len().min(self.remaining); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for PageCandidates {} +impl std::iter::FusedIterator for PageCandidates {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::kv_to_pointer; + + fn distance(value: f32) -> Distance { + Distance::from_f32(value) + } + + #[test] + fn aggregates_tokens_by_page_and_preserves_distance_order() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![ + ( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-0.5), kv_to_pointer((page_2, 0))), + ], + vec![(distance(-0.4), kv_to_pointer((page_2, 1)))], + distance(-0.75), + ), + ( + vec![(distance(-2.0), kv_to_pointer((page_1, 1)))], + vec![], + distance(-1.0), + ), + ] + .into_iter(); + let candidates = DensePageCandidateGenerator + .generate(2, &mut token_searches, usize::MAX) + .collect::>(); + + assert_eq!(candidates.len(), 2); + assert_eq!(candidates[0].heap_key, page_1); + assert_eq!(candidates[0].approximate_distance.to_f32(), -3.0); + assert_eq!(candidates[1].heap_key, page_2); + assert_eq!(candidates[1].approximate_distance.to_f32(), -1.5); + } + + #[test] + fn candidate_limit_is_applied_after_global_ordering() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-2.0), kv_to_pointer((page_2, 0))), + ], + vec![], + Distance::ZERO, + )] + .into_iter(); + let candidates = DensePageCandidateGenerator + .generate(1, &mut token_searches, 1) + .collect::>(); + + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].heap_key, page_2); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/exact.rs b/src/index/vchordrq/scanners/maxsim/exact.rs new file mode 100644 index 00000000..2d7e750d --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/exact.rs @@ -0,0 +1,505 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorStorage, + TileMaxsimSourceBinding, resolve_tilemaxsim_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::profile; +use super::rerank::RerankError; +use crate::datatype::memory_halfvec::HalfvecInput; +use crate::datatype::memory_vector::VectorInput; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use distance::Distance; +use pgrx::datum::{Array, DatumWithOid, FromDatum, IntoDatum}; +use pgrx::iter::TableIterator; +use pgrx::{name, pg_extern}; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; +use std::time::Duration; +use vchordrq::types::OwnedVector; +use vector::VectorBorrowed; + +/// Exact TileMaxSim for a caller-supplied candidate set of external tensors. +/// Candidate selection, tenancy, ACLs, graph traversal, and clustering are +/// intentionally outside this function. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_vector( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, VectorInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_vector_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +/// Half-precision overload of exact caller-scoped TileMaxSim. +#[pg_extern(sql = "")] +fn _vchordrq_tilemaxsim_rerank_halfvec( + source_oid: pgrx::pg_sys::Oid, + query: Array<'_, HalfvecInput<'_>>, + candidate_ids: Array<'_, i64>, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = collect_halfvec_query(query) + .and_then(|query| { + collect_candidate_ids(candidate_ids) + .and_then(|candidate_ids| execute_rerank(source_oid, query, candidate_ids, top_k)) + }) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn collect_vector_query( + query: Array<'_, VectorInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf32(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn collect_halfvec_query( + query: Array<'_, HalfvecInput<'_>>, +) -> Result, RerankError> { + let mut vectors = Vec::with_capacity(query.len()); + for vector in query.iter() { + let vector = vector.ok_or(RerankError::TensorMismatch)?; + vectors.push(OwnedVector::Vecf16(vector.as_borrowed().own())); + } + validate_query(&vectors)?; + Ok(vectors) +} + +fn validate_query(query: &[OwnedVector]) -> Result<(), RerankError> { + if query.is_empty() { + return Err(RerankError::Configuration( + "query tensor must contain at least one token", + )); + } + Ok(()) +} + +fn collect_candidate_ids(candidate_ids: Array<'_, i64>) -> Result, RerankError> { + let mut result = Vec::with_capacity(candidate_ids.len()); + let mut unique = BTreeSet::new(); + for public_id in candidate_ids.iter() { + let public_id = public_id.ok_or(RerankError::Configuration( + "candidate_ids must not contain NULL", + ))?; + if !unique.insert(public_id) { + return Err(RerankError::Configuration( + "candidate_ids must not contain duplicates", + )); + } + result.push(public_id); + } + if result.is_empty() { + return Err(RerankError::Configuration( + "candidate_ids must contain at least one ID", + )); + } + Ok(result) +} + +fn execute_rerank( + source_oid: pgrx::pg_sys::Oid, + query: Vec, + candidate_ids: Vec, + top_k: i32, +) -> Result, RerankError> { + validate_top_k(candidate_ids.len(), top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires vchordrq.maxsim_backend = 'gpu'", + )); + } + require_mvcc_snapshot()?; + + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); + profile::update(|profile| { + profile.query_tokens = query.len() as u64; + profile.generated_candidates = candidate_ids.len() as u64; + }); + + let preflight_timer = profile::ProfileTimer::start(); + let binding = resolve_tilemaxsim_source(source_oid)?; + let source_lock = RelationLock::open(source_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if source_lock.oid() != binding.source_oid + || descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid + { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source changed during execution".into(), + )); + } + preflight_descriptor_access(&binding)?; + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_timer.elapsed()); + }); + + let descriptor_timer = profile::ProfileTimer::start(); + let (mut source, public_ids) = load_visible_descriptors(&binding, &candidate_ids)?; + let mut candidates = public_ids + .keys() + .copied() + .map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }) + .collect::>() + .into_iter(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_timer.elapsed()); + profile.visible_candidates = public_ids.len() as u64; + profile.descriptors = public_ids.len() as u64; + }); + + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let mut backend = GpuExternalTileMaxsimBackend::new( + UnixSocketTransport::new(endpoint), + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + if let Some(tenant) = gucs::vchordrq_maxsim_tenant() { + backend = backend.with_scheduling(tenant, gucs::vchordrq_maxsim_priority()); + } + let sidecar_timer = profile::ProfileTimer::start(); + let result_limit = top_k as usize; + let mut best = BinaryHeap::with_capacity(result_limit.saturating_add(1)); + backend.rerank_batches(&query, &mut candidates, &mut source, |batch| { + for result in batch { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + let row = (result.distance, public_id); + super::retain_top_k(&mut best, result_limit, row); + } + Ok(()) + })?; + let mut rows = best.into_vec(); + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_timer.elapsed()); + }); + + let result_timer = profile::ProfileTimer::start(); + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + let output = rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_timer.elapsed()); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) +} + +fn validate_top_k(candidate_count: usize, top_k: i32) -> Result<(), RerankError> { + if top_k <= 0 || top_k as usize > candidate_count { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_ids length", + )); + } + Ok(()) +} + +fn require_mvcc_snapshot() -> Result<(), RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() + || unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC + { + return Err(RerankError::Configuration( + "external TileMaxSim rerank requires an active MVCC snapshot", + )); + } + Ok(()) +} + +fn preflight_descriptor_access(binding: &TileMaxsimSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.source_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every TileMaxSim relation".into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &TileMaxsimSourceBinding, + candidate_ids: &[i64], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + let mut candidates_by_id = BTreeMap::new(); + for (ordinal, public_id) in candidate_ids.iter().copied().enumerate() { + candidates_by_id.insert(public_id, candidate_for_ordinal(ordinal)?); + } + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let int8_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::INT8OID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(int8_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ + candidate_ids.to_vec().into(), + binding.model_contract_id.clone().into(), + ]; + let rows = client + .select(&prepared, Some(candidate_ids.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + for row in rows { + let public_id = required_heap_column::(&row, "public_id")?; + let candidate = candidates_by_id.get(&public_id).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown public ID".into()) + })?; + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() + || public_ids.insert(candidate.heap_key, public_id).is_some() + { + return Err(RerankError::Registry( + "descriptor query returned a duplicate public ID".into(), + )); + } + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn candidate_for_ordinal(ordinal: usize) -> Result { + let ordinal = u32::try_from(ordinal).map_err(|_| RerankError::RequestTooLarge)?; + Ok(PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, (ordinal >> 16) as u16, ordinal as u16], + }) +} + +fn descriptor_query( + binding: &TileMaxsimSourceBinding, + preflight: bool, +) -> Result { + let source_relation = relation_name(binding.source_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.{public_id}::bigint = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.{public_id}::bigint AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {source_relation} AS h + JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id}::bigint = h.{public_id}::bigint + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("query returned NULL {name}"))) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_candidates_support_full_corpus_ordinals() { + assert!(validate_top_k(256, 10).is_ok()); + assert!(validate_top_k(10, 0).is_err()); + assert!(validate_top_k(10, 11).is_err()); + let first = candidate_for_ordinal(0).unwrap().heap_key; + let last = candidate_for_ordinal(1_000_000).unwrap().heap_key; + assert_ne!(first, last); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs new file mode 100644 index 00000000..713fdd7f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -0,0 +1,673 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use super::candidate::PageCandidate; +use super::rerank::RerankError; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use pgrx::datum::{FromDatum, IntoDatum}; +use std::ffi::CString; + +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENSOR_ROWS: u32 = 65_536; +const MAX_TENSOR_DIMENSION: u32 = 60_000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorDtype { + F32, + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorStorage { + SameHeap, + DescriptorRelation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct ExternalTensorColumns { + pub model_contract: i16, + pub public_id: i16, + pub tensor_ref: i16, + pub tensor_rows: i16, + pub tensor_dimension: i16, + pub tensor_dtype: i16, + pub tensor_checksum: i16, +} + +impl ExternalTensorColumns { + pub(super) fn validate(self) -> Result { + let columns = [ + self.model_contract, + self.public_id, + self.tensor_ref, + self.tensor_rows, + self.tensor_dimension, + self.tensor_dtype, + self.tensor_checksum, + ]; + if columns.iter().any(|column| *column <= 0) { + return Err(RerankError::InvalidDescriptor( + "registered attribute number is invalid", + )); + } + for (index, column) in columns.iter().enumerate() { + if columns[..index].contains(column) { + return Err(RerankError::InvalidDescriptor( + "registered descriptor columns are not distinct", + )); + } + } + Ok(self) + } +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorDescriptor { + pub candidate: PageCandidate, + /// Stable application identifier. It stays inside PostgreSQL and is never + /// encoded into the sidecar request. + #[allow( + dead_code, + reason = "returned by the score API through a separate visible-row map" + )] + pub public_id: i64, + pub tensor_ref: String, + pub rows: u32, + pub dimension: u32, + pub dtype: ExternalTensorDtype, + pub checksum: String, +} + +pub(super) trait CandidateTensorDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError>; +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub(super) struct HeapTensorRefSource<'a, F> { + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorSourceBinding { + pub index_oid: pgrx::pg_sys::Oid, + pub heap_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + #[allow( + dead_code, + reason = "physical attnums are consumed by the optional Phase 3C heap source" + )] + pub columns: Option, + pub column_names: ExternalTensorColumnNames, +} + +/// External tensor metadata bound directly to an application source relation. +/// +/// Unlike [`ExternalTensorSourceBinding`], this binding has no ANN index: the +/// caller supplies an already-authorized candidate ID set and VectorChord only +/// performs exact TileMaxSim over those candidates. +#[derive(Clone, Debug)] +pub(super) struct TileMaxsimSourceBinding { + pub source_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + pub column_names: ExternalTensorColumnNames, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorColumnNames { + pub model_contract: String, + pub public_id: String, + pub descriptor_public_id: Option, + pub tensor_ref: String, + pub tensor_rows: String, + pub tensor_dimension: String, + pub tensor_dtype: String, + pub tensor_checksum: String, +} + +/// Resolve the privilege-aware SQL registry boundary. Same-heap bindings also +/// resolve the physical attribute numbers consumed by [`HeapTensorRefSource`]; +/// independent descriptor relations are projected later through SPI. +/// +/// The SECURITY DEFINER SQL function performs ownership/SELECT checks and +/// revalidates the live index, heap relation, opclass, column types, and NOT +/// NULL constraints. This Rust layer deliberately calls that function rather +/// than reading the private registry table with extension privileges. +pub(super) fn resolve_external_tensor_source( + index_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_maxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_index::oid AS index_oid, + heap_relation::oid AS heap_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[index_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_index_oid = required_column::(&row, "index_oid")?; + let heap_oid = required_column::(&row, "heap_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_index_oid != index_oid { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolved a different index".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + "heap_array" => { + return Err(RerankError::Registry( + "registered MaxSim tensor source is not external".into(), + )); + } + _ => { + return Err(RerankError::Registry( + "registered external MaxSim tensor source is inconsistent".into(), + )); + } + }; + + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + let columns = match storage { + ExternalTensorStorage::SameHeap => Some( + ExternalTensorColumns { + model_contract: resolve_attnum(heap_oid, &column_names.model_contract)?, + public_id: resolve_attnum(heap_oid, &column_names.public_id)?, + tensor_ref: resolve_attnum(heap_oid, &column_names.tensor_ref)?, + tensor_rows: resolve_attnum(heap_oid, &column_names.tensor_rows)?, + tensor_dimension: resolve_attnum(heap_oid, &column_names.tensor_dimension)?, + tensor_dtype: resolve_attnum(heap_oid, &column_names.tensor_dtype)?, + tensor_checksum: resolve_attnum(heap_oid, &column_names.tensor_checksum)?, + } + .validate()?, + ), + ExternalTensorStorage::DescriptorRelation => { + if column_names.descriptor_public_id.is_none() { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + None + } + }; + validate_model_contract(&model_contract_id)?; + Ok(ExternalTensorSourceBinding { + index_oid, + heap_oid, + descriptor_oid, + storage, + model_contract_id, + columns, + column_names, + }) + }) +} + +/// Resolve a source-relation binding for exact candidate-only TileMaxSim. +/// +/// The SECURITY DEFINER SQL resolver revalidates relation ownership, caller +/// SELECT privileges, column types, uniqueness, and descriptor metadata. This +/// Rust layer deliberately sees only the validated projection. +pub(super) fn resolve_tilemaxsim_source( + source_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_tilemaxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_source::oid AS source_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[source_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_source_oid = required_column::(&row, "source_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_source_oid != source_oid { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source resolved a different relation".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + _ => { + return Err(RerankError::Registry( + "registered TileMaxSim tensor source is inconsistent".into(), + )); + } + }; + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + if matches!(storage, ExternalTensorStorage::DescriptorRelation) + && column_names.descriptor_public_id.is_none() + { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + validate_model_contract(&model_contract_id)?; + Ok(TileMaxsimSourceBinding { + source_oid, + descriptor_oid, + storage, + model_contract_id, + column_names, + }) + }) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("registry resolver returned NULL {name}"))) +} + +fn optional_column( + row: &pgrx::spi::SpiTupleTable<'_>, + name: &str, +) -> Result, RerankError> +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name).map_err(registry_error) +} + +fn resolve_attnum(heap_oid: pgrx::pg_sys::Oid, column_name: &str) -> Result { + let column_name = CString::new(column_name) + .map_err(|_| RerankError::Registry("registry column contains a NUL byte".into()))?; + let attnum = unsafe { pgrx::pg_sys::get_attnum(heap_oid, column_name.as_ptr()) }; + if attnum <= 0 { + return Err(RerankError::Registry( + "registered descriptor column disappeared during resolution".into(), + )); + } + Ok(attnum) +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +impl<'a, F> HeapTensorRefSource<'a, F> { + pub(super) fn new( + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, + ) -> Result { + validate_model_contract(&model_contract_id)?; + Ok(Self { + fetcher, + model_contract_id, + columns: columns.validate()?, + }) + } +} + +impl CandidateTensorDescriptorSource for HeapTensorRefSource<'_, F> { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + + // This is the data-exfiltration boundary: no descriptor value may be + // read before PostgreSQL has accepted the active base-relation qual. + if !tuple.filter() { + return Ok(None); + } + + let model_contract = read_attribute::(&mut tuple, self.columns.model_contract)?; + if model_contract != self.model_contract_id { + return Err(RerankError::ModelContractMismatch); + } + let public_id = read_attribute::(&mut tuple, self.columns.public_id)?; + let tensor_ref = read_attribute::(&mut tuple, self.columns.tensor_ref)?; + let rows = read_attribute::(&mut tuple, self.columns.tensor_rows)?; + let dimension = read_attribute::(&mut tuple, self.columns.tensor_dimension)?; + let dtype = read_attribute::(&mut tuple, self.columns.tensor_dtype)?; + let checksum = read_attribute::(&mut tuple, self.columns.tensor_checksum)?; + + validate_descriptor( + candidate, public_id, tensor_ref, rows, dimension, dtype, checksum, + ) + .map(Some) + } +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +fn read_attribute(tuple: &mut impl Tuple, attnum: i16) -> Result { + let attribute = tuple + .attribute(attnum) + .ok_or(RerankError::InvalidDescriptor( + "registered attribute is unavailable", + ))?; + unsafe { T::from_datum(attribute.datum, attribute.is_null) }.ok_or( + RerankError::InvalidDescriptor("registered descriptor value is NULL or malformed"), + ) +} + +fn validate_model_contract(value: &str) -> Result<(), RerankError> { + if value.is_empty() || value.len() > MAX_MODEL_CONTRACT_BYTES || contains_control(value) { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn validate_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: String, + rows: i32, + dimension: i32, + dtype: String, + checksum: String, +) -> Result { + if tensor_ref.is_empty() + || tensor_ref.len() > MAX_TENSOR_REF_BYTES + || contains_control(&tensor_ref) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + let rows = u32::try_from(rows) + .ok() + .filter(|rows| (1..=MAX_TENSOR_ROWS).contains(rows)) + .ok_or(RerankError::InvalidDescriptor( + "tensor row count is invalid", + ))?; + let dimension = u32::try_from(dimension) + .ok() + .filter(|dimension| (1..=MAX_TENSOR_DIMENSION).contains(dimension)) + .ok_or(RerankError::InvalidDescriptor( + "tensor dimension is invalid", + ))?; + let dtype = match dtype.as_str() { + "float32" => ExternalTensorDtype::F32, + "float16" => ExternalTensorDtype::F16, + _ => { + return Err(RerankError::InvalidDescriptor( + "tensor dtype must be float16 or float32", + )); + } + }; + if checksum.len() > MAX_CHECKSUM_BYTES || !is_sha256_checksum(&checksum) { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref, + rows, + dimension, + dtype, + checksum, + }) +} + +fn contains_control(value: &str) -> bool { + value.chars().any(char::is_control) +} + +fn is_sha256_checksum(value: &str) -> bool { + value.strip_prefix("sha256:").is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::TupleAttribute; + use distance::Distance; + + fn candidate() -> PageCandidate { + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + } + } + + fn columns() -> ExternalTensorColumns { + ExternalTensorColumns { + model_contract: 1, + public_id: 2, + tensor_ref: 3, + tensor_rows: 4, + tensor_dimension: 5, + tensor_dtype: 6, + tensor_checksum: 7, + } + } + + struct RejectingFetcher; + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("external descriptor sources do not build index expressions") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose descriptor attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: [u16; 3]) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn source_filters_before_reading_descriptor_attributes() { + let mut fetcher = RejectingFetcher; + let mut source = + HeapTensorRefSource::new(&mut fetcher, "contract@1".into(), columns()).unwrap(); + assert!(source.fetch(candidate()).unwrap().is_none()); + } + + #[test] + fn descriptor_validation_accepts_bounded_immutable_metadata() { + let descriptor = validate_descriptor( + candidate(), + 42, + "s3://immutable-bucket/page-42.tensor".into(), + 747, + 320, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + assert_eq!(descriptor.public_id, 42); + assert_eq!(descriptor.rows, 747); + assert_eq!(descriptor.dimension, 320); + assert_eq!(descriptor.dtype, ExternalTensorDtype::F16); + } + + #[test] + fn descriptor_validation_rejects_ambiguous_or_unbounded_metadata() { + assert!(matches!( + columns_with_duplicate().validate(), + Err(RerankError::InvalidDescriptor(_)) + )); + for (rows, dimension, dtype, checksum) in [ + (0, 320, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 0, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "bf16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "float16", "sha256:short".into()), + ] { + assert!(matches!( + validate_descriptor( + candidate(), + 1, + "s3://immutable/tensor".into(), + rows, + dimension, + dtype.into(), + checksum, + ), + Err(RerankError::InvalidDescriptor(_)) + )); + } + } + + fn columns_with_duplicate() -> ExternalTensorColumns { + ExternalTensorColumns { + public_id: 1, + ..columns() + } + } +} diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs new file mode 100644 index 00000000..f530ca50 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -0,0 +1,2094 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, +}; +use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, RerankResults}; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::io::{Read, Write}; +use std::mem::size_of_val; +use std::net::{TcpStream, ToSocketAddrs}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; +use vchordrq::types::OwnedVector; + +const MAGIC: &[u8; 4] = b"VCTM"; +const VERSION: u16 = 1; +const EXTERNAL_VERSION: u16 = 2; +const SCHEDULED_EXTERNAL_VERSION: u16 = 3; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; +const HEADER_LEN: usize = 24; +const MAX_REMOTE_ERROR_BYTES: usize = 64 * 1024; +const MAX_EXTERNAL_CANDIDATES_PER_BATCH: usize = 65_536; +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENANT_BYTES: usize = 256; + +static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); +static LAST_FALLBACK_WARNING_SECONDS: AtomicU64 = AtomicU64::new(0); + +pub(super) fn report_gpu_fallback(error: &RerankError) { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + let should_warn = LAST_FALLBACK_WARNING_SECONDS + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |last| { + (last == 0 || now.saturating_sub(last) >= 60).then_some(now) + }) + .is_ok(); + if should_warn { + pgrx::warning!("GPU MaxSim failed; using cpu_exact: {error}"); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TensorDtype { + F32 = 1, + F16 = 2, +} + +pub(super) trait TileMaxsimTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError>; +} + +pub(super) struct GpuTileMaxsimBackend { + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuTileMaxsimBackend { + pub fn new( + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl ExactMaxsimBackend for GpuTileMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_request( + request_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response(&response, request_id, &encoded.heap_keys) + } +} + +/// GPU backend for Phase 3B external tensor descriptors. +/// +/// This codec is intentionally separate from [`ExactMaxsimBackend`]: an +/// external full tensor may have different logical values from the indexed +/// sketch, so it cannot be substituted into ordinary `@#` execution. +pub(super) struct GpuExternalTileMaxsimBackend { + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option, +} + +#[derive(Clone, Debug)] +pub(super) struct TileMaxsimScheduling { + pub tenant: String, + pub priority: i32, +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn new( + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + model_contract_id, + timeout, + max_batch_tokens, + max_batch_bytes, + scheduling: None, + } + } + + pub(super) fn with_scheduling(mut self, tenant: String, priority: i32) -> Self { + self.scheduling = Some(TileMaxsimScheduling { tenant, priority }); + self + } +} + +impl GpuExternalTileMaxsimBackend { + #[cfg(test)] + pub(super) fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let mut inner = BinaryHeap::new(); + self.rerank_batches(query, candidates, source, |batch| { + inner.extend(batch.inner); + Ok(()) + })?; + Ok(RerankResults { inner }) + } + + /// Execute one logical rerank as bounded sidecar requests. + /// + /// The callback is invoked after every complete batch so callers that only + /// need a global top-k do not have to retain every score. `self.timeout` is + /// a deadline for the whole logical query, rather than a fresh timeout for + /// each IPC request. + pub(super) fn rerank_batches( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + mut consume: F, + ) -> Result<(), RerankError> + where + S: CandidateTensorDescriptorSource, + F: FnMut(RerankResults) -> Result<(), RerankError>, + { + let deadline = Instant::now() + .checked_add(self.timeout) + .ok_or_else(|| RerankError::Transport("logical request deadline overflow".into()))?; + let mut pending = None; + + loop { + let descriptors = collect_external_batch( + &self.model_contract_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + self.scheduling.as_ref(), + &mut pending, + )?; + if descriptors.is_empty() { + return Ok(()); + } + + let remaining = remaining_logical_timeout(deadline)?; + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_external_descriptors( + request_id, + &self.model_contract_id, + query, + &descriptors, + self.max_batch_tokens, + self.max_batch_bytes, + self.scheduling.as_ref(), + remaining, + )?; + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = self.transport.round_trip( + &encoded.frame, + remaining_logical_timeout(deadline)?, + max_response_bytes, + )?; + consume(decode_response_for_version( + &response, + encoded.version, + request_id, + &encoded.heap_keys, + )?)?; + } + } +} + +fn remaining_logical_timeout(deadline: Instant) -> Result { + deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("logical request timed out".into())) +} + +struct EncodedRequest { + frame: Vec, + heap_keys: Vec, + version: u16, +} + +#[cfg(test)] +fn encode_external_request( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + timeout: Duration, +) -> Result { + let mut descriptors = Vec::new(); + for candidate in candidates { + if let Some(descriptor) = source.fetch(candidate)? { + descriptors.push(descriptor); + } + } + encode_external_descriptors( + request_id, + model_contract_id, + query, + &descriptors, + max_batch_tokens, + max_batch_bytes, + scheduling, + timeout, + ) +} + +fn encode_external_descriptors( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + descriptors: &[ExternalTensorDescriptor], + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + timeout: Duration, +) -> Result { + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + if let Some(scheduling) = scheduling { + if scheduling.tenant.is_empty() + || scheduling.tenant.len() > MAX_TENANT_BYTES + || scheduling.tenant.chars().any(char::is_control) + || !(-100..=100).contains(&scheduling.priority) + { + return Err(RerankError::Configuration( + "TileMaxSim scheduler tenant or priority is invalid", + )); + } + } + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let mut declared_tensor_bytes = tensor_bytes(query_rows, dimension, dtype)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + writer + .u32(u32::try_from(model_contract_id.len()).map_err(|_| RerankError::RequestTooLarge)?)?; + let version = if let Some(scheduling) = scheduling { + writer.i32(scheduling.priority)?; + writer.u32( + u32::try_from(timeout.as_millis().clamp(1, 600_000)) + .map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(scheduling.tenant.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + SCHEDULED_EXTERNAL_VERSION + } else { + EXTERNAL_VERSION + }; + writer.bytes(model_contract_id.as_bytes())?; + if let Some(scheduling) = scheduling { + writer.bytes(scheduling.tenant.as_bytes())?; + } + encode_tensor_values(&mut writer, query, dtype)?; + + if descriptors.len() > MAX_EXTERNAL_CANDIDATES_PER_BATCH { + return Err(RerankError::RequestTooLarge); + } + let mut heap_keys = Vec::with_capacity(descriptors.len()); + for descriptor in descriptors { + validate_external_for_request( + descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + total_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + declared_tensor_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(descriptor.rows)?; + writer.u32( + u32::try_from(descriptor.tensor_ref.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(descriptor.checksum.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.bytes(descriptor.tensor_ref.as_bytes())?; + writer.bytes(descriptor.checksum.as_bytes())?; + heap_keys.push(descriptor.candidate.heap_key); + } + + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, version); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64( + 16, + u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?, + ); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + version, + }) +} + +#[allow(clippy::too_many_arguments)] +fn collect_external_batch( + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, + scheduling: Option<&TileMaxsimScheduling>, + pending: &mut Option, +) -> Result, RerankError> { + validate_external_request_identity(model_contract_id, scheduling)?; + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let query_bytes = tensor_bytes(query_rows, dimension, dtype)?; + let mut total_tokens = query.len(); + let mut declared_tensor_bytes = query_bytes; + let mut frame_bytes = external_request_base_bytes(model_contract_id, scheduling, query_bytes)?; + if total_tokens > max_batch_tokens + || declared_tensor_bytes > max_batch_bytes + || frame_bytes > max_batch_bytes + { + return Err(RerankError::RequestTooLarge); + } + + let mut batch = Vec::new(); + while batch.len() < MAX_EXTERNAL_CANDIDATES_PER_BATCH { + let descriptor = if let Some(descriptor) = pending.take() { + descriptor + } else { + let mut fetched = None; + for candidate in &mut *candidates { + if let Some(descriptor) = source.fetch(candidate)? { + fetched = Some(descriptor); + break; + } + } + let Some(descriptor) = fetched else { + break; + }; + descriptor + }; + + validate_external_for_request( + &descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + let next_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + let next_declared_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + let descriptor_frame_bytes = 16usize + .checked_add(descriptor.tensor_ref.len()) + .and_then(|size| size.checked_add(descriptor.checksum.len())) + .ok_or(RerankError::RequestTooLarge)?; + let next_frame_bytes = frame_bytes + .checked_add(descriptor_frame_bytes) + .ok_or(RerankError::RequestTooLarge)?; + if next_tokens > max_batch_tokens + || next_declared_bytes > max_batch_bytes + || next_frame_bytes > max_batch_bytes + { + if batch.is_empty() { + return Err(RerankError::RequestTooLarge); + } + *pending = Some(descriptor); + break; + } + + total_tokens = next_tokens; + declared_tensor_bytes = next_declared_bytes; + frame_bytes = next_frame_bytes; + batch.push(descriptor); + } + Ok(batch) +} + +fn validate_external_request_identity( + model_contract_id: &str, + scheduling: Option<&TileMaxsimScheduling>, +) -> Result<(), RerankError> { + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + if let Some(scheduling) = scheduling { + if scheduling.tenant.is_empty() + || scheduling.tenant.len() > MAX_TENANT_BYTES + || scheduling.tenant.chars().any(char::is_control) + || !(-100..=100).contains(&scheduling.priority) + { + return Err(RerankError::Configuration( + "TileMaxSim scheduler tenant or priority is invalid", + )); + } + } + Ok(()) +} + +fn external_request_base_bytes( + model_contract_id: &str, + scheduling: Option<&TileMaxsimScheduling>, + query_bytes: usize, +) -> Result { + let fixed = if scheduling.is_some() { + 56usize + } else { + 44usize + }; + fixed + .checked_add(model_contract_id.len()) + .and_then(|size| { + size.checked_add(scheduling.map_or(0, |scheduling| scheduling.tenant.len())) + }) + .and_then(|size| size.checked_add(query_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: TensorDtype) -> Result { + let scalar_bytes = match dtype { + TensorDtype::F32 => 4usize, + TensorDtype::F16 => 2usize, + }; + usize::try_from(rows) + .ok() + .and_then(|rows| rows.checked_mul(dimension as usize)) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + +fn validate_external_for_request( + descriptor: &ExternalTensorDescriptor, + dimension: u32, + dtype: ExternalTensorDtype, + max_tensor_ref_bytes: usize, + max_checksum_bytes: usize, +) -> Result<(), RerankError> { + if descriptor.rows == 0 + || descriptor.rows > 65_536 + || dimension > 60_000 + || descriptor.dimension != dimension + || descriptor.dtype != dtype + { + return Err(RerankError::TensorMismatch); + } + if descriptor.tensor_ref.is_empty() + || descriptor.tensor_ref.len() > max_tensor_ref_bytes + || descriptor.tensor_ref.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + if descriptor.checksum.len() > max_checksum_bytes + || !descriptor + .checksum + .strip_prefix("sha256:") + .is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) + { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(()) +} + +fn encode_request( + request_id: u64, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + let (dtype, dimension) = tensor_metadata(query)?; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let (candidate_dtype, candidate_dimension) = tensor_metadata(&tensor.vectors)?; + if candidate_dtype != dtype || candidate_dimension != dimension { + return Err(RerankError::TensorMismatch); + } + total_tokens = total_tokens + .checked_add(tensor.vectors.len()) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + let rows = u32::try_from(tensor.vectors.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(rows)?; + encode_tensor_values(&mut writer, &tensor.vectors, dtype)?; + heap_keys.push(tensor.candidate.heap_key); + } + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + let body_len = u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64(16, body_len); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + version: VERSION, + }) +} + +fn tensor_metadata(vectors: &[OwnedVector]) -> Result<(TensorDtype, u32), RerankError> { + let Some(first) = vectors.first() else { + return Err(RerankError::TensorMismatch); + }; + let dtype = match first { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + let dimension = first.dim(); + for vector in vectors { + let this_dtype = match vector { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + if this_dtype != dtype || vector.dim() != dimension { + return Err(RerankError::TensorMismatch); + } + } + Ok((dtype, dimension)) +} + +fn encode_tensor_values( + writer: &mut BoundedWriter, + vectors: &[OwnedVector], + dtype: TensorDtype, +) -> Result<(), RerankError> { + for vector in vectors { + match (dtype, vector) { + (TensorDtype::F32, OwnedVector::Vecf32(vector)) => { + for value in vector.slice() { + writer.bytes(&value.to_le_bytes())?; + } + } + (TensorDtype::F16, OwnedVector::Vecf16(vector)) => { + for value in vector.slice() { + writer.u16(value.to_bits())?; + } + } + _ => return Err(RerankError::TensorMismatch), + } + } + Ok(()) +} + +fn decode_response( + frame: &[u8], + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + decode_response_for_version(frame, VERSION, request_id, heap_keys) +} + +fn decode_response_for_version( + frame: &[u8], + expected_version: u16, + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + let mut cursor = Cursor::new(frame); + if cursor.bytes(4)? != MAGIC { + return Err(RerankError::Protocol("invalid magic".into())); + } + if cursor.u16()? != expected_version { + return Err(RerankError::Protocol("unsupported version".into())); + } + if cursor.u16()? != RESPONSE_KIND { + return Err(RerankError::Protocol("unexpected message kind".into())); + } + if cursor.u64()? != request_id { + return Err(RerankError::Protocol("request ID mismatch".into())); + } + let body_len = usize::try_from(cursor.u64()?) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + if body_len != frame.len().saturating_sub(HEADER_LEN) { + return Err(RerankError::Protocol("response length mismatch".into())); + } + let status = cursor.u32()?; + if status != 0 { + let length = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("remote error is too large".into()))?; + if length > MAX_REMOTE_ERROR_BYTES { + return Err(RerankError::Protocol("remote error is too large".into())); + } + let message = std::str::from_utf8(cursor.bytes(length)?) + .map_err(|_| RerankError::Protocol("remote error is not UTF-8".into()))?; + cursor.finish()?; + return Err(RerankError::Remote(message.into())); + } + let result_count = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("result count is too large".into()))?; + if result_count != heap_keys.len() { + return Err(RerankError::Protocol("partial result set".into())); + } + let mut seen = vec![false; heap_keys.len()]; + let mut results = BinaryHeap::new(); + for _ in 0..result_count { + let candidate_id = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("candidate ID is too large".into()))?; + let Some(heap_key) = heap_keys.get(candidate_id).copied() else { + return Err(RerankError::Protocol("unknown candidate ID".into())); + }; + if std::mem::replace(&mut seen[candidate_id], true) { + return Err(RerankError::Protocol("duplicate candidate ID".into())); + } + let similarity = f32::from_bits(cursor.u32()?); + if !similarity.is_finite() { + return Err(RerankError::Protocol("non-finite similarity".into())); + } + let distance = Distance::from_f32(-similarity); + results.push((Reverse(distance), Reverse(heap_key))); + } + cursor.finish()?; + if seen.iter().any(|seen| !seen) { + return Err(RerankError::Protocol("partial result set".into())); + } + Ok(RerankResults { inner: results }) +} + +struct BoundedWriter { + bytes: Vec, + limit: usize, +} + +impl BoundedWriter { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn len(&self) -> usize { + self.bytes.len() + } + + fn ensure(&self, additional: usize) -> Result<(), RerankError> { + let size = self + .bytes + .len() + .checked_add(additional) + .ok_or(RerankError::RequestTooLarge)?; + if size > self.limit { + return Err(RerankError::RequestTooLarge); + } + Ok(()) + } + + fn zeros(&mut self, count: usize) -> Result<(), RerankError> { + self.ensure(count)?; + self.bytes.resize(self.bytes.len() + count, 0); + Ok(()) + } + + fn bytes(&mut self, bytes: &[u8]) -> Result<(), RerankError> { + self.ensure(bytes.len())?; + self.bytes.extend_from_slice(bytes); + Ok(()) + } + + fn u8(&mut self, value: u8) -> Result<(), RerankError> { + self.bytes(&[value]) + } + + fn u16(&mut self, value: u16) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn u32(&mut self, value: u32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn i32(&mut self, value: i32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn patch_bytes(&mut self, offset: usize, bytes: &[u8]) { + self.bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + + fn patch_u16(&mut self, offset: usize, value: u16) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u32(&mut self, offset: usize, value: u32) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u64(&mut self, offset: usize, value: u64) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn finish(self) -> Vec { + self.bytes + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn bytes(&mut self, count: usize) -> Result<&'a [u8], RerankError> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| RerankError::Protocol("message offset overflow".into()))?; + let bytes = self + .bytes + .get(self.offset..end) + .ok_or_else(|| RerankError::Protocol("truncated message".into()))?; + self.offset = end; + Ok(bytes) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.bytes(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.bytes(4)?.try_into().unwrap())) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.bytes(8)?.try_into().unwrap())) + } + + fn finish(self) -> Result<(), RerankError> { + if self.offset != self.bytes.len() { + return Err(RerankError::Protocol("trailing response bytes".into())); + } + Ok(()) + } +} + +pub(super) struct UnixSocketTransport { + endpoint: String, +} + +impl UnixSocketTransport { + pub fn new(endpoint: String) -> Self { + Self { endpoint } + } +} + +#[cfg(unix)] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError> { + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + let deadline = Instant::now() + .checked_add(timeout) + .ok_or_else(|| RerankError::Transport("invalid timeout".into()))?; + let mut stream = connect_endpoint_interruptible(&self.endpoint, deadline)?; + let poll = remaining_until(deadline)?.min(Duration::from_millis(50)); + stream + .set_read_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + stream + .set_write_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + + write_interruptible(&mut stream, request, deadline)?; + let mut header = [0u8; HEADER_LEN]; + read_interruptible(&mut stream, &mut header, deadline)?; + let body_len = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + let response_len = HEADER_LEN + .checked_add(body_len) + .ok_or_else(|| RerankError::Protocol("response length overflow".into()))?; + if response_len > max_response_bytes { + return Err(RerankError::Protocol( + "response exceeds configured limit".into(), + )); + } + let mut response = Vec::with_capacity(response_len); + response.extend_from_slice(&header); + response.resize(response_len, 0); + read_interruptible(&mut stream, &mut response[HEADER_LEN..], deadline)?; + Ok(response) + } +} + +#[cfg(unix)] +enum TransportStream { + Unix(std::os::unix::net::UnixStream), + Tcp(TcpStream), +} + +#[cfg(unix)] +impl TransportStream { + fn set_read_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_read_timeout(timeout), + Self::Tcp(stream) => stream.set_read_timeout(timeout), + } + } + + fn set_write_timeout(&self, timeout: Option) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.set_write_timeout(timeout), + Self::Tcp(stream) => stream.set_write_timeout(timeout), + } + } +} + +#[cfg(unix)] +impl Read for TransportStream { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.read(buffer), + Self::Tcp(stream) => stream.read(buffer), + } + } +} + +#[cfg(unix)] +impl Write for TransportStream { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + match self { + Self::Unix(stream) => stream.write(buffer), + Self::Tcp(stream) => stream.write(buffer), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match self { + Self::Unix(stream) => stream.flush(), + Self::Tcp(stream) => stream.flush(), + } + } +} + +#[cfg(unix)] +fn connect_endpoint_interruptible( + endpoint: &str, + deadline: Instant, +) -> Result { + let Some(authority) = endpoint.strip_prefix("tcp://") else { + return connect_interruptible(endpoint, deadline).map(TransportStream::Unix); + }; + if authority.is_empty() + || authority.contains('/') + || authority.contains('@') + || !authority.contains(':') + { + return Err(RerankError::Transport( + "TCP endpoint must be tcp://HOST:PORT without credentials or a path".into(), + )); + } + let remaining = remaining_until(deadline)?; + let addresses = authority + .to_socket_addrs() + .map_err(|error| RerankError::Transport(error.to_string()))? + .collect::>(); + if addresses.is_empty() { + return Err(RerankError::Transport( + "TCP endpoint resolved to no addresses".into(), + )); + } + let mut last_error = None; + for address in addresses { + pgrx::check_for_interrupts!(); + let attempt = remaining_until(deadline)?.min(remaining); + match TcpStream::connect_timeout(&address, attempt) { + Ok(stream) => { + stream.set_nodelay(true).ok(); + return Ok(TransportStream::Tcp(stream)); + } + Err(error) => last_error = Some(error), + } + } + Err(RerankError::Transport( + last_error + .map(|error| error.to_string()) + .unwrap_or_else(|| "TCP connection failed".to_owned()), + )) +} + +#[cfg(unix)] +fn connect_interruptible( + endpoint: &str, + deadline: Instant, +) -> Result { + use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; + + let (address, address_len) = unix_socket_address(endpoint)?; + let raw_fd = unsafe { libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0) }; + if raw_fd < 0 { + return Err(last_transport_error()); + } + let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) }; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFD, + libc::F_SETFD, + libc::FD_CLOEXEC, + true, + )?; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + true, + )?; + + let connected = unsafe { + libc::connect( + fd.as_raw_fd(), + (&raw const address).cast::(), + address_len, + ) + } == 0; + if !connected { + let error = std::io::Error::last_os_error(); + let raw_error = error.raw_os_error(); + if raw_error != Some(libc::EINPROGRESS) + && raw_error != Some(libc::EAGAIN) + && raw_error != Some(libc::EWOULDBLOCK) + { + return Err(RerankError::Transport(error.to_string())); + } + wait_for_connect(fd.as_raw_fd(), deadline)?; + } + + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + false, + )?; + Ok(std::os::unix::net::UnixStream::from(fd)) +} + +#[cfg(unix)] +fn unix_socket_address( + endpoint: &str, +) -> Result<(libc::sockaddr_un, libc::socklen_t), RerankError> { + let path = endpoint.as_bytes(); + let mut address = unsafe { std::mem::zeroed::() }; + if path.contains(&0) { + return Err(RerankError::Transport( + "endpoint contains a NUL byte".into(), + )); + } + if path.len() >= address.sun_path.len() { + return Err(RerankError::Transport("endpoint path is too long".into())); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + unsafe { + std::ptr::copy_nonoverlapping( + path.as_ptr().cast::(), + address.sun_path.as_mut_ptr(), + path.len(), + ); + } + let length = std::mem::offset_of!(libc::sockaddr_un, sun_path) + .checked_add(path.len()) + .and_then(|length| length.checked_add(1)) + .and_then(|length| libc::socklen_t::try_from(length).ok()) + .ok_or_else(|| RerankError::Transport("endpoint path is too long".into()))?; + #[cfg(any( + target_os = "aix", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "haiku", + target_os = "hurd", + target_os = "ios", + target_os = "macos", + target_os = "netbsd", + target_os = "openbsd", + target_os = "tvos", + target_os = "visionos", + target_os = "watchos" + ))] + { + address.sun_len = u8::try_from(length) + .map_err(|_| RerankError::Transport("endpoint path is too long".into()))?; + } + Ok((address, length)) +} + +#[cfg(unix)] +fn update_fd_flag( + fd: std::os::fd::RawFd, + get_command: libc::c_int, + set_command: libc::c_int, + flag: libc::c_int, + enabled: bool, +) -> Result<(), RerankError> { + let current = unsafe { libc::fcntl(fd, get_command) }; + if current < 0 { + return Err(last_transport_error()); + } + let updated = if enabled { + current | flag + } else { + current & !flag + }; + if unsafe { libc::fcntl(fd, set_command, updated) } < 0 { + return Err(last_transport_error()); + } + Ok(()) +} + +#[cfg(unix)] +fn wait_for_connect(fd: std::os::fd::RawFd, deadline: Instant) -> Result<(), RerankError> { + loop { + pgrx::check_for_interrupts!(); + let remaining = remaining_until(deadline)?; + let timeout_ms = remaining.min(Duration::from_millis(50)).as_millis().max(1) as libc::c_int; + let mut poll_fd = libc::pollfd { + fd, + events: libc::POLLOUT, + revents: 0, + }; + let result = unsafe { libc::poll(&mut poll_fd, 1, timeout_ms) }; + if result == 0 { + continue; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() == std::io::ErrorKind::Interrupted { + continue; + } + return Err(RerankError::Transport(error.to_string())); + } + let mut socket_error = 0; + let mut socket_error_len = size_of_val(&socket_error) as libc::socklen_t; + if unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&raw mut socket_error).cast(), + &raw mut socket_error_len, + ) + } < 0 + { + return Err(last_transport_error()); + } + if socket_error != 0 { + return Err(RerankError::Transport( + std::io::Error::from_raw_os_error(socket_error).to_string(), + )); + } + return Ok(()); + } +} + +#[cfg(unix)] +fn remaining_until(deadline: Instant) -> Result { + deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("request timed out".into())) +} + +#[cfg(unix)] +fn last_transport_error() -> RerankError { + RerankError::Transport(std::io::Error::last_os_error().to_string()) +} + +#[cfg(unix)] +fn write_interruptible( + stream: &mut TransportStream, + mut bytes: &[u8], + deadline: Instant, +) -> Result<(), RerankError> { + while !bytes.is_empty() { + match stream.write(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(unix)] +fn read_interruptible( + stream: &mut TransportStream, + mut bytes: &mut [u8], + deadline: Instant, +) -> Result<(), RerankError> { + while !bytes.is_empty() { + match stream.read(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &mut bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(not(unix))] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + _request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + Err(RerankError::Transport( + "Unix sockets are not supported on this platform".into(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, + }; + use super::super::rerank::CandidateTensor; + use super::*; + use std::cell::RefCell; + use std::collections::BTreeMap; + use std::rc::Rc; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + struct MockDescriptorSource(BTreeMap); + + impl CandidateTensorDescriptorSource for MockDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some( + self.0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + )) + } + } + + struct MockTransport { + similarities: Vec<(u32, f32)>, + } + + impl TileMaxsimTransport for MockTransport { + fn round_trip( + &mut self, + request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + Ok(success_response_with_version( + version, + request_id, + &self.similarities, + )) + } + } + + #[derive(Default)] + struct BatchObservations { + candidate_counts: Vec, + transport_timeouts: Vec, + scheduled_timeouts_ms: Vec, + } + + struct BatchingTransport { + observations: Rc>, + delay: Duration, + } + + impl TileMaxsimTransport for BatchingTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + let candidate_count = u32::from_le_bytes(request[32..36].try_into().unwrap()); + let call = { + let mut observations = self.observations.borrow_mut(); + let call = observations.candidate_counts.len(); + observations.candidate_counts.push(candidate_count); + observations.transport_timeouts.push(timeout); + if version == SCHEDULED_EXTERNAL_VERSION { + observations + .scheduled_timeouts_ms + .push(u32::from_le_bytes(request[48..52].try_into().unwrap())); + } + call + }; + std::thread::sleep(self.delay); + let similarities = (0..candidate_count) + .map(|candidate_id| (candidate_id, call as f32 * 10.0 + candidate_id as f32)) + .collect::>(); + Ok(success_response_with_version( + version, + request_id, + &similarities, + )) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + fn half_vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf16(VectOwned::new( + values.iter().copied().map(simd::f16::from_f32).collect(), + )) + } + + fn success_response(request_id: u64, similarities: &[(u32, f32)]) -> Vec { + success_response_with_version(VERSION, request_id, similarities) + } + + fn success_response_with_version( + version: u16, + request_id: u64, + similarities: &[(u32, f32)], + ) -> Vec { + let body_len = 8 + similarities.len() * 8; + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&version.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&0u32.to_le_bytes()); + response.extend_from_slice(&(similarities.len() as u32).to_le_bytes()); + for (candidate_id, similarity) in similarities { + response.extend_from_slice(&candidate_id.to_le_bytes()); + response.extend_from_slice(&similarity.to_bits().to_le_bytes()); + } + response + } + + fn external_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: &str, + rows: u32, + dimension: u32, + dtype: ExternalTensorDtype, + ) -> ExternalTensorDescriptor { + ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref: tensor_ref.into(), + rows, + dimension, + dtype, + checksum: format!("sha256:{}", "a".repeat(64)), + } + } + + fn error_response(request_id: u64, message: &str) -> Vec { + let body_len = 8 + message.len(); + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&VERSION.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&1u32.to_le_bytes()); + response.extend_from_slice(&(message.len() as u32).to_le_bytes()); + response.extend_from_slice(message.as_bytes()); + response + } + + #[test] + fn gpu_backend_maps_positive_similarity_to_ascending_distance() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_1, + }, + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_2, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[0.5, 0.0])]), + ])); + let transport = MockTransport { + similarities: vec![(0, 1.0), (1, 2.0)], + }; + let results = GpuTileMaxsimBackend::new(transport, Duration::from_secs(1), 100, 4096) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, page_2); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_1); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn response_rejects_partial_and_duplicate_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let partial = success_response(7, &[(0, 1.0)]); + assert!(matches!( + decode_response(&partial, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let duplicate = success_response(7, &[(0, 1.0), (0, 2.0)]); + assert!(matches!( + decode_response(&duplicate, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_ids_may_arrive_out_of_order() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let response = success_response(7, &[(1, 0.5), (0, 1.0)]); + let results = decode_response(&response, 7, &keys) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, keys[0]); + assert_eq!(results[0].distance.to_f32(), -1.0); + assert_eq!(results[1].heap_key, keys[1]); + assert_eq!(results[1].distance.to_f32(), -0.5); + } + + #[test] + fn response_rejects_unknown_non_finite_and_trailing_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let unknown = success_response(7, &[(0, 1.0), (2, 2.0)]); + assert!(matches!( + decode_response(&unknown, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let non_finite = success_response(7, &[(0, f32::NAN), (1, 2.0)]); + assert!(matches!( + decode_response(&non_finite, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let mut trailing = success_response(7, &[(0, 1.0), (1, 2.0)]); + trailing.push(0); + assert!(matches!( + decode_response(&trailing, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_rejects_invalid_header_fields() { + let keys = [[0, 0, 1]]; + let valid = success_response(7, &[(0, 1.0)]); + + for (offset, replacement) in [(0, 0u8), (4, 2), (6, 1), (8, 8), (16, 0)] { + let mut invalid = valid.clone(); + invalid[offset] = replacement; + assert!(matches!( + decode_response(&invalid, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + } + + #[test] + fn response_surfaces_remote_error() { + let response = error_response(7, "CUDA queue is unavailable"); + assert!(matches!( + decode_response(&response, 7, &[]), + Err(RerankError::Remote(message)) if message == "CUDA queue is unavailable" + )); + } + + #[test] + fn request_frame_is_versioned_and_length_prefixed() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[0.5, 0.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + VERSION + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[6..8].try_into().unwrap()), + REQUEST_KIND + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[8..16].try_into().unwrap()), + 9 + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[16..24].try_into().unwrap()) as usize, + encoded.frame.len() - HEADER_LEN + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_request_encodes_contract_and_opaque_descriptor_ids() { + let page = [0, 0, 7]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let tensor_ref = "s3://immutable/page-9001.tensor"; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor(candidate, 9001, tensor_ref, 2, 2, ExternalTensorDtype::F16), + )])); + let contract = "colqwen@immutable-revision"; + let encoded = encode_external_request( + 19, + contract, + &query, + &mut candidates, + &mut source, + 100, + 4096, + None, + Duration::from_secs(2), + ) + .unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + EXTERNAL_VERSION + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!(&encoded.frame[44..44 + contract_len], contract.as_bytes()); + let candidate_offset = 44 + contract_len + 4; // one 2-D f16 query row + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset..candidate_offset + 4] + .try_into() + .unwrap() + ), + 0 + ); + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset + 4..candidate_offset + 8] + .try_into() + .unwrap() + ), + 2 + ); + let reference_len = u32::from_le_bytes( + encoded.frame[candidate_offset + 8..candidate_offset + 12] + .try_into() + .unwrap(), + ) as usize; + let reference_offset = candidate_offset + 16; + assert_eq!( + &encoded.frame[reference_offset..reference_offset + reference_len], + tensor_ref.as_bytes() + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn scheduled_external_request_encodes_tenant_priority_and_timeout() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9002, + "sha256://opaque", + 2, + 2, + ExternalTensorDtype::F16, + ), + )])); + let scheduling = TileMaxsimScheduling { + tenant: "tenant-a".to_owned(), + priority: 17, + }; + let encoded = encode_external_request( + 20, + "contract@1", + &query, + &mut candidates, + &mut source, + 100, + 4096, + Some(&scheduling), + Duration::from_millis(4_000), + ) + .unwrap(); + + assert_eq!(encoded.version, SCHEDULED_EXTERNAL_VERSION); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + SCHEDULED_EXTERNAL_VERSION + ); + assert_eq!( + i32::from_le_bytes(encoded.frame[44..48].try_into().unwrap()), + 17 + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[48..52].try_into().unwrap()), + 4_000 + ); + let tenant_len = u32::from_le_bytes(encoded.frame[52..56].try_into().unwrap()) as usize; + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!( + &encoded.frame[56 + contract_len..56 + contract_len + tenant_len], + b"tenant-a" + ); + } + + #[test] + fn external_backend_maps_scores_without_exposing_public_ids() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + i64::MAX, + "object://immutable/tensor", + 4, + 2, + ExternalTensorDtype::F32, + ), + )])); + let transport = MockTransport { + similarities: vec![(0, 3.5)], + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_secs(1), + 100, + 4096, + ) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].heap_key, page); + assert_eq!(results[0].distance.to_f32(), -3.5); + } + + #[test] + fn external_backend_splits_one_logical_query_and_shares_its_deadline() { + let pages = [[0, 0, 1], [0, 0, 2], [0, 0, 3]]; + let candidates = pages.map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }); + let query = vec![vector(&[1.0, 0.0])]; + let mut candidate_iter = candidates.into_iter(); + let mut source = + MockDescriptorSource(BTreeMap::from_iter(candidates.into_iter().enumerate().map( + |(index, candidate)| { + ( + candidate.heap_key, + external_descriptor( + candidate, + index as i64, + &format!("object://immutable/tensor-{index}"), + 2, + 2, + ExternalTensorDtype::F32, + ), + ) + }, + ))); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::from_millis(5), + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(100), + 5, + 4096, + ) + .with_scheduling("tenant-a".into(), 9) + .rerank(&query, &mut candidate_iter, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 3); + let observations = observations.borrow(); + assert_eq!(observations.candidate_counts, vec![2, 1]); + assert!(observations.transport_timeouts[1] < observations.transport_timeouts[0]); + assert!(observations.scheduled_timeouts_ms[1] < observations.scheduled_timeouts_ms[0]); + } + + #[test] + fn external_backend_rejects_a_single_unsplittable_tensor() { + let page = [0, 0, 1]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 1, + "object://immutable/tensor", + 5, + 2, + ExternalTensorDtype::F32, + ), + )])); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::ZERO, + }; + let result = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(100), + 5, + 4096, + ) + .rerank(&query, &mut candidates, &mut source); + + assert!(matches!(result, Err(RerankError::RequestTooLarge))); + assert!(observations.borrow().candidate_counts.is_empty()); + } + + #[test] + fn external_backend_does_not_refresh_timeout_for_later_batches() { + let pages = [[0, 0, 1], [0, 0, 2], [0, 0, 3]]; + let candidates = pages.map(|heap_key| PageCandidate { + approximate_distance: Distance::ZERO, + heap_key, + }); + let query = vec![vector(&[1.0, 0.0])]; + let mut candidate_iter = candidates.into_iter(); + let mut source = + MockDescriptorSource(BTreeMap::from_iter(candidates.into_iter().enumerate().map( + |(index, candidate)| { + ( + candidate.heap_key, + external_descriptor( + candidate, + index as i64, + &format!("object://immutable/tensor-{index}"), + 2, + 2, + ExternalTensorDtype::F32, + ), + ) + }, + ))); + let observations = Rc::new(RefCell::new(BatchObservations::default())); + let transport = BatchingTransport { + observations: Rc::clone(&observations), + delay: Duration::from_millis(20), + }; + let result = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_millis(5), + 5, + 4096, + ) + .rerank(&query, &mut candidate_iter, &mut source); + + assert!(matches!( + result, + Err(RerankError::Transport(message)) if message == "logical request timed out" + )); + assert_eq!(observations.borrow().candidate_counts, vec![2]); + } + + #[test] + fn external_request_rejects_shape_and_declared_payload_overflow() { + let page = [0, 0, 9]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, 0.0])]; + let make_source = |dimension, rows| { + MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9, + "object://immutable/tensor", + rows, + dimension, + ExternalTensorDtype::F16, + ), + )])) + }; + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(3, 1), + 100, + 4096, + None, + Duration::from_secs(2), + ), + Err(RerankError::TensorMismatch) + )); + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(2, 100), + 1000, + 64, + None, + Duration::from_secs(2), + ), + Err(RerankError::RequestTooLarge) + )); + } + + #[test] + fn request_frame_encodes_f16_tensor_bits() { + let page = [0, 0, 1]; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = + MockTensorSource(BTreeMap::from([(page, vec![half_vector(&[0.25, 2.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(encoded.frame[36], TensorDtype::F16 as u8); + assert_eq!( + u16::from_le_bytes(encoded.frame[40..42].try_into().unwrap()), + simd::f16::from_f32(1.0).to_bits() + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[42..44].try_into().unwrap()), + simd::f16::from_f32(-0.5).to_bits() + ); + } + + #[test] + fn request_limits_are_enforced_before_transport() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[1.0, 0.0])])])); + assert!(matches!( + encode_request(1, &query, &mut candidates, &mut source, 1, 4096), + Err(RerankError::RequestTooLarge) + )); + } + + #[cfg(unix)] + #[test] + fn unix_socket_address_is_length_bounded_and_nul_terminated() { + let (address, length) = unix_socket_address("/tmp/vectorchord.sock").unwrap(); + let path_offset = std::mem::offset_of!(libc::sockaddr_un, sun_path); + + assert_eq!(address.sun_family, libc::AF_UNIX as libc::sa_family_t); + assert_eq!( + length as usize, + path_offset + "/tmp/vectorchord.sock".len() + 1 + ); + assert_eq!( + &address.sun_path[.."/tmp/vectorchord.sock".len()], + "/tmp/vectorchord.sock" + .as_bytes() + .iter() + .map(|byte| *byte as libc::c_char) + .collect::>() + ); + assert_eq!(address.sun_path["/tmp/vectorchord.sock".len()], 0); + + let too_long = "x".repeat(address.sun_path.len()); + assert!(matches!( + unix_socket_address(&too_long), + Err(RerankError::Transport(message)) if message == "endpoint path is too long" + )); + assert!(matches!( + unix_socket_address("invalid\0path"), + Err(RerankError::Transport(message)) if message == "endpoint contains a NUL byte" + )); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/profile.rs b/src/index/vchordrq/scanners/maxsim/profile.rs new file mode 100644 index 00000000..da6bfd7c --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/profile.rs @@ -0,0 +1,177 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// Copyright (c) 2026 Hu Xinjing + +use std::cell::RefCell; +use std::time::{Duration, Instant}; + +#[derive(Default)] +pub(super) struct MaxsimProfile { + pub query_tokens: u64, + pub token_search_calls: u64, + pub token_search_us: u64, + pub token_search_results: u64, + pub token_refine_calls: u64, + pub token_refine_us: u64, + pub accurate_token_hits: u64, + pub rough_token_hits: u64, + pub hit_collect_us: u64, + pub hit_updates: u64, + pub page_token_updates: u64, + pub hit_sort_us: u64, + pub page_aggregate_us: u64, + pub aggregated_pages: u64, + pub preflight_us: u64, + pub candidate_generation_us: u64, + pub generated_candidates: u64, + pub visibility_us: u64, + pub visible_candidates: u64, + pub descriptor_us: u64, + pub descriptors: u64, + pub sidecar_us: u64, + pub result_finalize_us: u64, + pub returned_rows: u64, + pub total_us: u64, +} + +impl MaxsimProfile { + fn to_json(&self) -> String { + format!( + concat!( + "{{", + "\"schema_version\":1,", + "\"query_tokens\":{},", + "\"token_search_calls\":{},", + "\"token_search_us\":{},", + "\"token_search_results\":{},", + "\"token_refine_calls\":{},", + "\"token_refine_us\":{},", + "\"accurate_token_hits\":{},", + "\"rough_token_hits\":{},", + "\"hit_collect_us\":{},", + "\"hit_updates\":{},", + "\"page_token_updates\":{},", + "\"hit_sort_us\":{},", + "\"page_aggregate_us\":{},", + "\"aggregated_pages\":{},", + "\"preflight_us\":{},", + "\"candidate_generation_us\":{},", + "\"generated_candidates\":{},", + "\"visibility_us\":{},", + "\"visible_candidates\":{},", + "\"descriptor_us\":{},", + "\"descriptors\":{},", + "\"sidecar_us\":{},", + "\"result_finalize_us\":{},", + "\"returned_rows\":{},", + "\"total_us\":{}", + "}}" + ), + self.query_tokens, + self.token_search_calls, + self.token_search_us, + self.token_search_results, + self.token_refine_calls, + self.token_refine_us, + self.accurate_token_hits, + self.rough_token_hits, + self.hit_collect_us, + self.hit_updates, + self.page_token_updates, + self.hit_sort_us, + self.page_aggregate_us, + self.aggregated_pages, + self.preflight_us, + self.candidate_generation_us, + self.generated_candidates, + self.visibility_us, + self.visible_candidates, + self.descriptor_us, + self.descriptors, + self.sidecar_us, + self.result_finalize_us, + self.returned_rows, + self.total_us, + ) + } +} + +thread_local! { + static ACTIVE_PROFILE: RefCell> = const { RefCell::new(None) }; +} + +pub(super) struct ProfileGuard { + enabled: bool, + finished: bool, +} + +impl ProfileGuard { + pub fn start(enabled: bool) -> Self { + if enabled { + ACTIVE_PROFILE.with(|slot| { + *slot.borrow_mut() = Some(MaxsimProfile::default()); + }); + } + Self { + enabled, + finished: false, + } + } + + pub fn finish(mut self, total: Duration) { + if self.enabled { + let profile = ACTIVE_PROFILE.with(|slot| { + let mut profile = slot.borrow_mut().take(); + if let Some(profile) = profile.as_mut() { + profile.total_us = duration_us(total); + } + profile + }); + if let Some(profile) = profile { + pgrx::notice!("vchordrq_maxsim_profile {}", profile.to_json()); + } + } + self.finished = true; + } +} + +impl Drop for ProfileGuard { + fn drop(&mut self) { + if self.enabled && !self.finished { + ACTIVE_PROFILE.with(|slot| { + slot.borrow_mut().take(); + }); + } + } +} + +pub(super) struct ProfileTimer(Option); + +impl ProfileTimer { + pub fn start() -> Self { + let enabled = ACTIVE_PROFILE.with(|slot| slot.borrow().is_some()); + Self(enabled.then(Instant::now)) + } + + pub fn elapsed(self) -> Duration { + self.0.map_or(Duration::ZERO, |started| started.elapsed()) + } +} + +pub(super) fn update(f: impl FnOnce(&mut MaxsimProfile)) { + ACTIVE_PROFILE.with(|slot| { + if let Some(profile) = slot.borrow_mut().as_mut() { + f(profile); + } + }); +} + +pub(super) fn duration_us(duration: Duration) -> u64 { + duration.as_micros().min(u128::from(u64::MAX)) as u64 +} diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs new file mode 100644 index 00000000..151ae4ea --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -0,0 +1,315 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use crate::index::vchordrq::opclass::Opfamily; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use vchordrq::types::OwnedVector; + +pub(super) struct CandidateTensor { + pub candidate: PageCandidate, + pub vectors: Vec, +} + +pub(super) trait CandidateTensorSource { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError>; +} + +pub(super) struct HeapArrayTensorSource<'a, F> { + fetcher: &'a mut F, + opfamily: Opfamily, +} + +impl<'a, F> HeapArrayTensorSource<'a, F> { + pub fn new(fetcher: &'a mut F, opfamily: Opfamily) -> Self { + Self { fetcher, opfamily } + } +} + +impl CandidateTensorSource for HeapArrayTensorSource<'_, F> { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + // Exact sources are the last boundary before a tensor may leave the + // executor process. Re-evaluate the active base-relation scan qual + // here even if token reranking already prefiltered some hits. This + // keeps same-relation quals ahead of CPU/GPU tensor access and + // also covers configurations with token-level refine disabled. + if !tuple.filter() { + return Ok(None); + } + let (values, is_nulls) = tuple.build(); + if is_nulls[0] { + return Err(RerankError::TensorMismatch); + } + let vectors = + unsafe { self.opfamily.input_vectors(values[0]) }.ok_or(RerankError::TensorMismatch)?; + Ok(Some(CandidateTensor { candidate, vectors })) + } +} + +pub(super) trait ExactMaxsimBackend { + type Results: Iterator; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result; +} + +#[derive(Debug)] +pub(super) enum RerankError { + TensorMismatch, + ModelContractMismatch, + InvalidDescriptor(&'static str), + Registry(String), + Configuration(&'static str), + UnsupportedTensorKind, + RequestTooLarge, + Transport(String), + Protocol(String), + Remote(String), +} + +impl std::fmt::Display for RerankError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TensorMismatch => write!(f, "MaxSim tensor kind or dimension is not matched"), + Self::ModelContractMismatch => write!(f, "MaxSim model contract is not matched"), + Self::InvalidDescriptor(message) => { + write!(f, "invalid external MaxSim tensor descriptor: {message}") + } + Self::Registry(message) => { + write!(f, "MaxSim tensor-source registry error: {message}") + } + Self::Configuration(message) => write!(f, "MaxSim configuration error: {message}"), + Self::UnsupportedTensorKind => { + write!(f, "GPU MaxSim supports only vector and halfvec tensors") + } + Self::RequestTooLarge => write!(f, "GPU MaxSim request exceeds configured limits"), + Self::Transport(message) => write!(f, "GPU MaxSim transport error: {message}"), + Self::Protocol(message) => write!(f, "GPU MaxSim protocol error: {message}"), + Self::Remote(message) => write!(f, "GPU MaxSim sidecar error: {message}"), + } + } +} + +impl std::error::Error for RerankError {} + +#[derive(Default)] +pub(super) struct CpuExactMaxsimBackend; + +impl ExactMaxsimBackend for CpuExactMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let mut results = BinaryHeap::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let Some(distance) = exact_maxsim_distance(query, &tensor.vectors) else { + return Err(RerankError::TensorMismatch); + }; + results.push((Reverse(distance), Reverse(tensor.candidate.heap_key))); + } + Ok(RerankResults { inner: results }) + } +} + +fn exact_maxsim_distance(query: &[OwnedVector], document: &[OwnedVector]) -> Option { + if query.is_empty() || document.is_empty() { + return None; + } + let mut maxsim = 0.0f32; + for query_vector in query { + let mut best = Distance::INFINITY; + for document_vector in document { + let distance = document_vector.operator_dot(query_vector)?; + best = std::cmp::min(best, distance); + } + maxsim += best.to_f32(); + } + Some(Distance::from_f32(maxsim)) +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct RerankedPage { + pub distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) struct RerankResults { + pub(super) inner: BinaryHeap<(Reverse, Reverse)>, +} + +impl Iterator for RerankResults { + type Item = RerankedPage; + + fn next(&mut self) -> Option { + let (Reverse(distance), Reverse(heap_key)) = self.inner.pop()?; + Some(RerankedPage { distance, heap_key }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for RerankResults {} +impl std::iter::FusedIterator for RerankResults {} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + struct RejectingFetcher; + + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("a rejected tuple must not materialize its tensor") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose heap attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn heap_source_applies_scan_qual_before_materializing_tensor() { + let mut fetcher = RejectingFetcher; + let mut source = HeapArrayTensorSource::new(&mut fetcher, Opfamily::VectorMaxsim); + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + + assert!(source.fetch(candidate).unwrap().is_none()); + } + + #[test] + fn cpu_backend_reorders_candidates_by_exact_page_maxsim() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(-2.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(-1.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]), + (page_2, vec![vector(&[0.5, 0.5])]), + ])); + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].heap_key, page_1); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_2); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn exact_ties_are_ordered_by_heap_key() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[1.0, 0.0])]), + ])); + + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!( + results.iter().map(|page| page.heap_key).collect::>(), + vec![page_1, page_2] + ); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs new file mode 100644 index 00000000..da9aca34 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -0,0 +1,625 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, + ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::rerank::RerankError; +use super::{MaxsimBuilder, profile}; +use crate::index::fetcher::{ + Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, +}; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::DefaultRecorder; +use distance::Distance; +use pgrx::datum::{DatumWithOid, FromDatum}; +use pgrx::iter::TableIterator; +use pgrx::{AnyArray, IntoDatum, name}; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; +use std::time::Duration; + +const MAX_EXPLICIT_CANDIDATES: i32 = 65_536; + +/// Restricted Phase 3B search surface. Candidate generation reads only the +/// named index. Descriptor projection happens later through SPI, under the +/// caller's normal SELECT privileges and active MVCC snapshot. +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_maxsim_search_external( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = execute_external_search(index_oid, query, candidate_limit, top_k) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn execute_external_search( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> Result, RerankError> { + let profile_guard = profile::ProfileGuard::start(gucs::vchordrq_maxsim_profile()); + let total_timer = profile::ProfileTimer::start(); + validate_search_limits(candidate_limit, top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external MaxSim search currently requires vchordrq.maxsim_backend = 'gpu'", + )); + } + + let preflight_timer = profile::ProfileTimer::start(); + let binding = resolve_external_tensor_source(index_oid)?; + let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; + let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if binding.index_oid != index_lock.oid() || binding.heap_oid != heap_lock.oid() { + return Err(RerankError::Registry( + "registered MaxSim tensor source changed during execution".into(), + )); + } + if descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid { + return Err(RerankError::Registry( + "registered descriptor relation changed during execution".into(), + )); + } + + let opfamily = unsafe { opfamily(index_lock.raw()) }; + if !matches!(opfamily, Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim) { + return Err(RerankError::UnsupportedTensorKind); + } + let indexed_type = unsafe { pgrx::pg_sys::get_atttype(index_oid, 1) }; + if indexed_type != query.oid() { + return Err(RerankError::TensorMismatch); + } + let query_vectors = + unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + profile::update(|profile| { + profile.query_tokens = query_vectors.len() as u64; + }); + + // The registry resolution happens before the index read, and this + // privilege-only SELECT is planned/executed before any candidate CTID is + // generated. It fails early when the caller cannot project the registered + // descriptor columns. The same query shape is used for the actual fetch. + preflight_descriptor_access(&binding)?; + let preflight_elapsed = preflight_timer.elapsed(); + profile::update(|profile| { + profile.preflight_us += profile::duration_us(preflight_elapsed); + }); + + let candidate_timer = profile::ProfileTimer::start(); + let candidates = generate_candidates( + index_lock.raw(), + opfamily, + query.datum(), + candidate_limit as u32, + )?; + let candidate_elapsed = candidate_timer.elapsed(); + profile::update(|profile| { + profile.candidate_generation_us += profile::duration_us(candidate_elapsed); + profile.generated_candidates = candidates.len() as u64; + }); + let visibility_timer = profile::ProfileTimer::start(); + let resolved = + resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let visibility_elapsed = visibility_timer.elapsed(); + profile::update(|profile| { + profile.visibility_us += profile::duration_us(visibility_elapsed); + profile.visible_candidates = resolved.len() as u64; + }); + let descriptor_timer = profile::ProfileTimer::start(); + let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; + let visible_candidates = resolved + .into_iter() + .filter_map(|resolved| { + public_ids + .contains_key(&resolved.candidate.heap_key) + .then_some(resolved.candidate) + }) + .collect::>(); + let descriptor_elapsed = descriptor_timer.elapsed(); + profile::update(|profile| { + profile.descriptor_us += profile::duration_us(descriptor_elapsed); + profile.descriptors = public_ids.len() as u64; + }); + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let transport = UnixSocketTransport::new(endpoint); + let mut backend = GpuExternalTileMaxsimBackend::new( + transport, + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + if let Some(tenant) = gucs::vchordrq_maxsim_tenant() { + backend = backend.with_scheduling(tenant, gucs::vchordrq_maxsim_priority()); + } + let mut candidate_iter = visible_candidates.into_iter(); + let sidecar_timer = profile::ProfileTimer::start(); + let result_limit = top_k as usize; + let mut best = BinaryHeap::with_capacity(result_limit.saturating_add(1)); + backend.rerank_batches(&query_vectors, &mut candidate_iter, &mut source, |batch| { + for result in batch { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + let row = (result.distance, public_id); + super::retain_top_k(&mut best, result_limit, row); + } + Ok(()) + })?; + let mut rows = best.into_vec(); + let sidecar_elapsed = sidecar_timer.elapsed(); + profile::update(|profile| { + profile.sidecar_us += profile::duration_us(sidecar_elapsed); + }); + let result_finalize_timer = profile::ProfileTimer::start(); + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + let output = rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>(); + let result_finalize_elapsed = result_finalize_timer.elapsed(); + profile::update(|profile| { + profile.result_finalize_us += profile::duration_us(result_finalize_elapsed); + profile.returned_rows = output.len() as u64; + }); + profile_guard.finish(total_timer.elapsed()); + Ok(output.into_iter()) +} + +#[derive(Clone, Copy)] +struct ResolvedCandidate { + candidate: PageCandidate, + current_ctid: pgrx::pg_sys::ItemPointerData, +} + +fn resolve_visible_candidates( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + candidates: impl Iterator, +) -> Result, RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() { + return Err(RerankError::Configuration( + "external MaxSim search requires an active MVCC snapshot", + )); + } + if unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC { + return Err(RerankError::Configuration( + "external MaxSim search requires an MVCC snapshot", + )); + } + + let mut fetcher = + unsafe { HeapFetcher::new_standalone(index_relation, heap_relation, snapshot) }; + let mut resolved = Vec::new(); + for candidate in candidates { + pgrx::check_for_interrupts!(); + if let Some(tuple) = fetcher.fetch(candidate.heap_key) { + resolved.push(ResolvedCandidate { + candidate, + current_ctid: tuple.ctid(), + }); + } + } + Ok(resolved) +} + +fn validate_search_limits(candidate_limit: i32, top_k: i32) -> Result<(), RerankError> { + if !(1..=MAX_EXPLICIT_CANDIDATES).contains(&candidate_limit) { + return Err(RerankError::Configuration( + "candidate_limit must be between 1 and 65536", + )); + } + if top_k <= 0 || top_k > candidate_limit { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_limit", + )); + } + Ok(()) +} + +fn generate_candidates( + index_relation: pgrx::pg_sys::Relation, + opfamily: Opfamily, + query: pgrx::pg_sys::Datum, + candidate_limit: u32, +) -> Result, RerankError> { + let index = unsafe { PostgresRelation::::new(index_relation) }; + let mut builder = MaxsimBuilder::new(opfamily); + unsafe { builder.add(3, Some(query)) }; + let options = SearchOptions { + epsilon: unsafe { gucs::vchordrq_epsilon(index_relation) }, + probes: unsafe { gucs::vchordrq_probes(index_relation) }, + max_scan_tuples: None, + maxsim_refine: gucs::vchordrq_maxsim_refine(index_relation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold(index_relation), + maxsim_candidate_limit: Some(candidate_limit), + maxsim_backend: PostgresMaxsimBackend::CoarseOnly, + maxsim_gpu_endpoint: None, + maxsim_gpu_timeout_ms: 1, + maxsim_gpu_max_batch_tokens: 1, + maxsim_gpu_max_batch_bytes: 1, + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + // General same-relation quals would require the optional Phase 3C + // CustomScan. The restricted function applies PostgreSQL row + // visibility during the following SPI descriptor fetch. + prefilter: false, + }; + let bump = bumpalo::Bump::new(); + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: unsafe { (*index_relation).rd_id.to_u32() }, + }; + let candidates = builder + .build(&index, options, NoHeapFetcher, &bump, recorder) + .map(|(distance, heap_key, _)| PageCandidate { + approximate_distance: Distance::from_f32(distance), + heap_key, + }) + .collect(); + Ok(candidates) +} + +fn preflight_descriptor_access(binding: &ExternalTensorSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.heap_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every external MaxSim relation" + .into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &ExternalTensorSourceBinding, + candidates: &[ResolvedCandidate], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + if candidates.is_empty() { + return Ok((MaterializedDescriptorSource::default(), BTreeMap::new())); + } + let mut candidates_by_key = BTreeMap::new(); + for resolved in candidates { + if candidates_by_key + .insert(ctid_to_key(resolved.current_ctid), resolved.candidate) + .is_some() + { + return Err(RerankError::Registry( + "multiple index candidates resolved to the same visible CTID".into(), + )); + } + } + let ctids = candidates + .iter() + .map(|resolved| resolved.current_ctid) + .collect::>(); + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let tid_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::TIDOID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(tid_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ctids.into(), binding.model_contract_id.clone().into()]; + let rows = client + .select(&prepared, Some(candidates.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + let mut unique_public_ids = BTreeSet::new(); + for row in rows { + let ctid = required_heap_column::(&row, "heap_tid")?; + let heap_key = ctid_to_key(ctid); + let candidate = candidates_by_key.get(&heap_key).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown CTID".into()) + })?; + let public_id = required_heap_column::(&row, "public_id")?; + if !unique_public_ids.insert(public_id) { + return Err(RerankError::InvalidDescriptor( + "public IDs are not unique in the visible candidate batch", + )); + } + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() { + return Err(RerankError::Registry( + "descriptor query returned a duplicate CTID".into(), + )); + } + public_ids.insert(candidate.heap_key, public_id); + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn descriptor_query( + binding: &ExternalTensorSourceBinding, + preflight: bool, +) -> Result { + let heap_relation = relation_name(binding.heap_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.ctid = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + LEFT JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id} = h.{public_id} + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn raw(&self) -> pgrx::pg_sys::Relation { + self.raw + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +struct NoHeapFetcher; +struct NoHeapTuple; + +impl Tuple for NoHeapTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + unreachable!("restricted external candidate generation must not read the heap") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl FilterableTuple for NoHeapTuple { + fn filter(&mut self) -> bool { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl Fetcher for NoHeapFetcher { + type Tuple<'a> = NoHeapTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_search_limits_are_bounded() { + assert!(validate_search_limits(256, 10).is_ok()); + for (candidates, top_k) in [(0, 1), (65_537, 1), (1, 0), (10, 11)] { + assert!(matches!( + validate_search_limits(candidates, top_k), + Err(RerankError::Configuration(_)) + )); + } + } + + #[test] + fn materialized_source_only_returns_visible_descriptors_once() { + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + let descriptor = validate_descriptor( + candidate, + 42, + "s3://immutable/tensor".into(), + 1, + 2, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + let mut source = + MaterializedDescriptorSource(BTreeMap::from([(candidate.heap_key, descriptor)])); + assert_eq!(source.fetch(candidate).unwrap().unwrap().public_id, 42); + assert!(source.fetch(candidate).unwrap().is_none()); + } +} diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs new file mode 100644 index 00000000..c5e392bc --- /dev/null +++ b/src/index/vchordrq/scanners/mod.rs @@ -0,0 +1,41 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +mod default; +mod maxsim; + +use crate::index::gucs::PostgresMaxsimBackend; +use crate::index::scanners::Io; +use std::ffi::CString; + +pub use default::DefaultBuilder; +pub use maxsim::MaxsimBuilder; + +#[derive(Debug)] +pub struct SearchOptions { + pub epsilon: f32, + pub probes: Vec, + pub max_scan_tuples: Option, + pub maxsim_refine: u32, + pub maxsim_threshold: u32, + pub maxsim_candidate_limit: Option, + pub maxsim_backend: PostgresMaxsimBackend, + pub maxsim_gpu_endpoint: Option, + pub maxsim_gpu_timeout_ms: u32, + pub maxsim_gpu_max_batch_tokens: u32, + pub maxsim_gpu_max_batch_bytes: u32, + pub io_search: Io, + pub io_rerank: Io, + pub prefilter: bool, +} diff --git a/src/index/vchordrq/types.rs b/src/index/vchordrq/types.rs new file mode 100644 index 00000000..cff94858 --- /dev/null +++ b/src/index/vchordrq/types.rs @@ -0,0 +1,217 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use serde::{Deserialize, Serialize}; +use validator::{Validate, ValidationError, ValidationErrors}; +use vchordrq::types::VchordrqIndexOptions; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqDefaultBuildOptions {} + +#[allow(clippy::derivable_impls)] +impl Default for VchordrqDefaultBuildOptions { + fn default() -> Self { + Self {} + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum KMeansAlgorithm { + Lloyd {}, + Hierarchical {}, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqInternalBuildOptions { + #[serde(default = "VchordrqInternalBuildOptions::default_lists")] + #[validate(length(min = 0, max = 8), custom(function = VchordrqInternalBuildOptions::validate_lists))] + pub lists: Vec, + #[serde(default = "VchordrqInternalBuildOptions::default_spherical_centroids")] + pub spherical_centroids: bool, + #[serde(default = "VchordrqInternalBuildOptions::default_sampling_factor")] + #[validate(range(min = 1, max = 1024))] + pub sampling_factor: u32, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_iterations")] + #[validate(range(min = 0, max = 1024))] + pub kmeans_iterations: u32, + #[serde(default = "VchordrqInternalBuildOptions::default_build_threads")] + #[validate(range(min = 1, max = 255))] + pub build_threads: u16, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_algorithm")] + pub kmeans_algorithm: KMeansAlgorithm, + #[serde(default = "VchordrqInternalBuildOptions::default_kmeans_dimension")] + #[validate(range(min = 1, max = 16000))] + pub kmeans_dimension: Option, +} + +impl VchordrqInternalBuildOptions { + fn default_lists() -> Vec { + Vec::new() + } + fn validate_lists(lists: &[u32]) -> Result<(), ValidationError> { + if !lists.is_sorted() { + return Err(ValidationError::new("`lists` should be in ascending order")); + } + if !lists.iter().all(|x| (1..=1 << 24).contains(x)) { + return Err(ValidationError::new("list is too long or too short")); + } + Ok(()) + } + fn default_spherical_centroids() -> bool { + false + } + fn default_sampling_factor() -> u32 { + 256 + } + fn default_kmeans_iterations() -> u32 { + 10 + } + fn default_build_threads() -> u16 { + 1 + } + fn default_kmeans_algorithm() -> KMeansAlgorithm { + KMeansAlgorithm::Lloyd {} + } + fn default_kmeans_dimension() -> Option { + None + } +} + +impl Default for VchordrqInternalBuildOptions { + fn default() -> Self { + Self { + lists: Self::default_lists(), + spherical_centroids: Self::default_spherical_centroids(), + sampling_factor: Self::default_sampling_factor(), + kmeans_iterations: Self::default_kmeans_iterations(), + build_threads: Self::default_build_threads(), + kmeans_algorithm: Self::default_kmeans_algorithm(), + kmeans_dimension: Self::default_kmeans_dimension(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqExternalBuildOptions { + #[validate(custom(function = VchordrqExternalBuildOptions::validate_table))] + pub table: String, +} + +impl VchordrqExternalBuildOptions { + fn validate_table(table: &str) -> Result<(), ValidationError> { + let (schema_name, table_name) = if let Some((left, right)) = table.split_once(".") { + (Some(left), right) + } else { + (None, table) + }; + fn check(s: &str) -> bool { + if s.is_empty() { + return false; + } + if !matches!(s.as_bytes()[0], b'A'..=b'Z' | b'a'..=b'z' | b'_') { + return false; + } + for c in s.as_bytes().iter().copied() { + if !matches!(c, b'0'..=b'9' | b'A'..=b'Z' | b'a'..=b'z' | b'_' | b'$') { + return false; + } + } + true + } + if let Some(schema_name) = schema_name { + if !check(schema_name) { + return Err(ValidationError::new("table name is not well-formed")); + } + } + if !check(table_name) { + return Err(ValidationError::new("table name is not well-formed")); + } + Ok(()) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[serde(rename_all = "snake_case")] +pub enum VchordrqBuildSourceOptions { + Default(VchordrqDefaultBuildOptions), + Internal(VchordrqInternalBuildOptions), + External(VchordrqExternalBuildOptions), +} + +impl Default for VchordrqBuildSourceOptions { + fn default() -> Self { + Self::Default(Default::default()) + } +} + +impl Validate for VchordrqBuildSourceOptions { + fn validate(&self) -> Result<(), ValidationErrors> { + use VchordrqBuildSourceOptions::*; + match self { + Default(default_build) => default_build.validate(), + Internal(internal_build) => internal_build.validate(), + External(external_build) => external_build.validate(), + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqBuildOptions { + #[serde(flatten)] + #[validate(nested)] + pub source: VchordrqBuildSourceOptions, + #[serde(deserialize_with = "VchordrqBuildOptions::deserialize_pin")] + #[serde(default = "VchordrqBuildOptions::default_pin")] + #[validate(range(min = -1, max = 2))] + pub pin: i32, +} + +impl VchordrqBuildOptions { + pub fn deserialize_pin<'de, D: serde::Deserializer<'de>>( + deserializer: D, + ) -> Result { + #[derive(Deserialize)] + #[serde(untagged)] + enum Untagged { + Bool(bool), + I32(i32), + } + + match Untagged::deserialize(deserializer)? { + Untagged::Bool(b) => Ok(if b { 1 } else { -1 }), + Untagged::I32(i) => Ok(i), + } + } + pub fn default_pin() -> i32 { + -1 + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct VchordrqIndexingOptions { + #[serde(flatten)] + #[validate(nested)] + pub index: VchordrqIndexOptions, + #[serde(default)] + #[validate(nested)] + pub build: VchordrqBuildOptions, +} diff --git a/src/lib.rs b/src/lib.rs index 3de782c0..93fcc081 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,29 +1,83 @@ -#![allow(clippy::too_many_arguments)] -#![allow(clippy::needless_range_loop)] -#![allow(clippy::type_complexity)] -#![allow(clippy::identity_op)] +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#![allow(unsafe_code)] +#![deny(ffi_unwind_calls)] -mod algorithm; mod datatype; -mod gucs; mod index; -mod postgres; -mod types; +mod recorder; mod upgrade; -mod utils; -pgrx::pg_module_magic!(); +pgrx::pg_module_magic!( + name = c"vchord", + version = { + const RAW: &str = env!("VCHORD_VERSION"); + const BUFFER: [u8; RAW.len() + 1] = { + let mut buffer = [0u8; RAW.len() + 1]; + let mut i = 0_usize; + while i < RAW.len() { + buffer[i] = RAW.as_bytes()[i]; + i += 1; + } + buffer + }; + const STR: &::core::ffi::CStr = + if let Ok(s) = ::core::ffi::CStr::from_bytes_with_nul(&BUFFER) { + s + } else { + panic!("there are null characters in VCHORD_VERSION") + }; + const { STR } + } +); +const _: &str = include_str!("./sql/bootstrap.sql"); +const _: &str = include_str!("./sql/finalize.sql"); pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap); pgrx::extension_sql_file!("./sql/finalize.sql", finalize); #[pgrx::pg_guard] -unsafe extern "C" fn _PG_init() { - detect::init(); +#[unsafe(export_name = "_PG_init")] +unsafe extern "C-unwind" fn _pg_init() { + if !unsafe { pgrx::pg_sys::process_shared_preload_libraries_in_progress } { + pgrx::error!("vchord must be loaded via shared_preload_libraries."); + } + IS_MAIN.set(true); + index::init(); + recorder::init(); unsafe { - index::init(); - gucs::init(); + #[cfg(feature = "pg14")] + pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr()); + #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17", feature = "pg18"))] + pgrx::pg_sys::MarkGUCPrefixReserved(c"vchord".as_ptr()); } } -#[cfg(not(target_endian = "little"))] -compile_error!("Target architecture is not supported."); +std::thread_local! { + static IS_MAIN: core::cell::Cell = const { core::cell::Cell::new(false) }; +} + +#[must_use] +fn is_main() -> bool { + IS_MAIN.get() +} + +#[cfg(not(panic = "unwind"))] +compile_error!("This crate must be compiled with `-Cpanic=unwind`."); + +#[cfg(not(miri))] +#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] +#[cfg(any(target_os = "linux", target_os = "macos"))] +#[global_allocator] +static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; diff --git a/src/postgres.rs b/src/postgres.rs deleted file mode 100644 index 4597931b..00000000 --- a/src/postgres.rs +++ /dev/null @@ -1,295 +0,0 @@ -use std::mem::{offset_of, MaybeUninit}; -use std::ptr::NonNull; - -const _: () = assert!( - offset_of!(pgrx::pg_sys::PageHeaderData, pd_linp) % pgrx::pg_sys::MAXIMUM_ALIGNOF as usize == 0 -); - -const fn size_of_contents() -> usize { - use pgrx::pg_sys::{PageHeaderData, BLCKSZ}; - let size_of_page = BLCKSZ as usize; - let size_of_header = offset_of!(PageHeaderData, pd_linp); - let size_of_opaque = size_of::(); - size_of_page - size_of_header - size_of_opaque -} - -#[repr(C, align(8))] -pub struct Page { - header: pgrx::pg_sys::PageHeaderData, - content: [u8; size_of_contents()], - opaque: Opaque, -} - -const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); -const _: () = assert!(size_of::() == pgrx::pg_sys::BLCKSZ as usize); - -impl Page { - pub fn init_mut(this: &mut MaybeUninit) -> &mut Self { - unsafe { - pgrx::pg_sys::PageInit( - this.as_mut_ptr() as pgrx::pg_sys::Page, - pgrx::pg_sys::BLCKSZ as usize, - size_of::(), - ); - (&raw mut (*this.as_mut_ptr()).opaque).write(Opaque::default()); - } - let this = unsafe { MaybeUninit::assume_init_mut(this) }; - assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); - this - } - #[allow(dead_code)] - pub unsafe fn assume_init_mut(this: &mut MaybeUninit) -> &mut Self { - let this = unsafe { MaybeUninit::assume_init_mut(this) }; - assert_eq!(offset_of!(Self, opaque), this.header.pd_special as usize); - this - } - pub fn get_opaque(&self) -> &Opaque { - &self.opaque - } - pub fn get_opaque_mut(&mut self) -> &mut Opaque { - &mut self.opaque - } - pub fn len(&self) -> u16 { - use pgrx::pg_sys::{ItemIdData, PageHeaderData}; - assert!(self.header.pd_lower as usize <= size_of::()); - assert!(self.header.pd_upper as usize <= size_of::()); - let lower = self.header.pd_lower as usize; - let upper = self.header.pd_upper as usize; - assert!(lower < upper); - ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16 - } - pub fn get(&self, i: u16) -> Option<&[u8]> { - use pgrx::pg_sys::{ItemIdData, PageHeaderData}; - if i == 0 { - return None; - } - assert!(self.header.pd_lower as usize <= size_of::()); - let lower = self.header.pd_lower as usize; - let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; - if i > n { - return None; - } - let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; - let lp_off = iid.lp_off() as usize; - let lp_len = iid.lp_len() as usize; - match iid.lp_flags() { - pgrx::pg_sys::LP_UNUSED => return None, - pgrx::pg_sys::LP_NORMAL => (), - pgrx::pg_sys::LP_REDIRECT => unimplemented!(), - pgrx::pg_sys::LP_DEAD => unimplemented!(), - _ => unreachable!(), - } - assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off && lp_off <= size_of::()); - assert!(lp_len <= size_of::()); - assert!(lp_off + lp_len <= size_of::()); - unsafe { - let ptr = (self as *const Self).cast::().add(lp_off as _); - Some(std::slice::from_raw_parts(ptr, lp_len as _)) - } - } - pub fn get_mut(&mut self, i: u16) -> Option<&mut [u8]> { - use pgrx::pg_sys::{ItemIdData, PageHeaderData}; - if i == 0 { - return None; - } - assert!(self.header.pd_lower as usize <= size_of::()); - let lower = self.header.pd_lower as usize; - let n = ((lower - offset_of!(PageHeaderData, pd_linp)) / size_of::()) as u16; - if i > n { - return None; - } - let iid = unsafe { self.header.pd_linp.as_ptr().add((i - 1) as _).read() }; - let lp_off = iid.lp_off() as usize; - let lp_len = iid.lp_len() as usize; - assert!(offset_of!(PageHeaderData, pd_linp) <= lp_off && lp_off <= size_of::()); - assert!(lp_len <= size_of::()); - assert!(lp_off + lp_len <= size_of::()); - unsafe { - let ptr = (self as *mut Self).cast::().add(lp_off as _); - Some(std::slice::from_raw_parts_mut(ptr, lp_len as _)) - } - } - pub fn alloc(&mut self, data: &[u8]) -> Option { - unsafe { - let i = pgrx::pg_sys::PageAddItemExtended( - (self as *const Self).cast_mut().cast(), - data.as_ptr().cast_mut().cast(), - data.len(), - 0, - 0, - ); - if i == 0 { - None - } else { - Some(i) - } - } - } - pub fn free(&mut self, i: u16) { - unsafe { - pgrx::pg_sys::PageIndexTupleDeleteNoCompact((self as *mut Self).cast(), i); - } - } - pub fn freespace(&self) -> u16 { - unsafe { pgrx::pg_sys::PageGetFreeSpace((self as *const Self).cast_mut().cast()) as u16 } - } -} - -#[repr(C, align(8))] -pub struct Opaque { - pub next: u32, -} - -impl Default for Opaque { - fn default() -> Self { - Self { next: u32::MAX } - } -} - -const _: () = assert!(align_of::() == pgrx::pg_sys::MAXIMUM_ALIGNOF as usize); - -pub struct BufferReadGuard { - buf: i32, - page: NonNull, - id: u32, -} - -impl BufferReadGuard { - pub fn get(&self) -> &Page { - unsafe { self.page.as_ref() } - } - #[allow(dead_code)] - pub fn id(&self) -> u32 { - self.id - } -} - -impl Drop for BufferReadGuard { - fn drop(&mut self) { - unsafe { - pgrx::pg_sys::UnlockReleaseBuffer(self.buf); - } - } -} - -pub struct BufferWriteGuard { - buf: i32, - page: NonNull, - state: *mut pgrx::pg_sys::GenericXLogState, - id: u32, -} - -impl BufferWriteGuard { - pub fn get(&self) -> &Page { - unsafe { self.page.as_ref() } - } - pub fn get_mut(&mut self) -> &mut Page { - unsafe { self.page.as_mut() } - } - pub fn id(&self) -> u32 { - self.id - } -} - -impl Drop for BufferWriteGuard { - fn drop(&mut self) { - unsafe { - if std::thread::panicking() { - pgrx::pg_sys::GenericXLogAbort(self.state); - } else { - pgrx::pg_sys::GenericXLogFinish(self.state); - } - pgrx::pg_sys::UnlockReleaseBuffer(self.buf); - } - } -} - -#[derive(Debug)] -pub struct Relation { - raw: pgrx::pg_sys::Relation, -} - -impl Relation { - pub unsafe fn new(raw: pgrx::pg_sys::Relation) -> Self { - Self { raw } - } - pub fn read(&self, id: u32) -> BufferReadGuard { - assert!(id != u32::MAX, "no such page"); - unsafe { - use pgrx::pg_sys::{ - BufferGetPage, LockBuffer, ReadBufferExtended, ReadBufferMode, BUFFER_LOCK_SHARE, - }; - let buf = ReadBufferExtended( - self.raw, - 0, - id, - ReadBufferMode::RBM_NORMAL, - std::ptr::null_mut(), - ); - LockBuffer(buf, BUFFER_LOCK_SHARE as _); - let page = NonNull::new(BufferGetPage(buf).cast()).expect("failed to get page"); - BufferReadGuard { buf, page, id } - } - } - pub fn write(&self, id: u32) -> BufferWriteGuard { - assert!(id != u32::MAX, "no such page"); - unsafe { - use pgrx::pg_sys::{ - ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, - ReadBufferExtended, ReadBufferMode, BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, - }; - let buf = ReadBufferExtended( - self.raw, - ForkNumber::MAIN_FORKNUM, - id, - ReadBufferMode::RBM_NORMAL, - std::ptr::null_mut(), - ); - LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); - let state = GenericXLogStart(self.raw); - let page = NonNull::new( - GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), - ) - .expect("failed to get page"); - BufferWriteGuard { - buf, - page: page.cast(), - state, - id, - } - } - } - pub fn extend(&self) -> BufferWriteGuard { - unsafe { - use pgrx::pg_sys::{ - ExclusiveLock, ForkNumber, GenericXLogRegisterBuffer, GenericXLogStart, LockBuffer, - LockRelationForExtension, ReadBufferExtended, ReadBufferMode, - UnlockRelationForExtension, BUFFER_LOCK_EXCLUSIVE, GENERIC_XLOG_FULL_IMAGE, - }; - LockRelationForExtension(self.raw, ExclusiveLock as _); - let buf = ReadBufferExtended( - self.raw, - ForkNumber::MAIN_FORKNUM, - u32::MAX, - ReadBufferMode::RBM_NORMAL, - std::ptr::null_mut(), - ); - UnlockRelationForExtension(self.raw, ExclusiveLock as _); - LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE as _); - let state = GenericXLogStart(self.raw); - let mut page = NonNull::new( - GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE as _) - .cast::>(), - ) - .expect("failed to get page"); - Page::init_mut(page.as_mut()); - BufferWriteGuard { - buf, - page: page.cast(), - state, - id: pgrx::pg_sys::BufferGetBlockNumber(buf), - } - } - } -} diff --git a/src/recorder/hook.rs b/src/recorder/hook.rs new file mode 100644 index 00000000..c0e5e786 --- /dev/null +++ b/src/recorder/hook.rs @@ -0,0 +1,53 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::recorder::worker::{delete_database, delete_index}; + +static mut PREV_OBJECT_ACCESS: pgrx::pg_sys::object_access_hook_type = None; + +#[pgrx::pg_guard] +unsafe extern "C-unwind" fn recorder_object_access( + access: pgrx::pg_sys::ObjectAccessType::Type, + class_id: pgrx::pg_sys::Oid, + object_id: pgrx::pg_sys::Oid, + sub_id: ::std::os::raw::c_int, + arg: *mut ::std::os::raw::c_void, +) { + unsafe { + use pgrx::pg_sys::submodules::ffi::pg_guard_ffi_boundary; + if let Some(prev_object_access_hook) = PREV_OBJECT_ACCESS { + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + prev_object_access_hook(access, class_id, object_id, sub_id, arg) + }); + } + if access == pgrx::pg_sys::ObjectAccessType::OAT_DROP + && class_id == pgrx::pg_sys::DatabaseRelationId + { + delete_database(object_id.to_u32()); + } else if access == pgrx::pg_sys::ObjectAccessType::OAT_DROP + && class_id == pgrx::pg_sys::RelationRelationId + { + delete_index(object_id.to_u32()); + } + } +} + +pub fn init() { + assert!(crate::is_main()); + unsafe { + PREV_OBJECT_ACCESS = pgrx::pg_sys::object_access_hook; + pgrx::pg_sys::object_access_hook = Some(recorder_object_access); + } +} diff --git a/src/recorder/mod.rs b/src/recorder/mod.rs new file mode 100644 index 00000000..1488a4ac --- /dev/null +++ b/src/recorder/mod.rs @@ -0,0 +1,26 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +pub use types::{DefaultRecorder, Recorder}; +pub use worker::dump; + +mod hook; +mod types; +mod worker; + +pub mod text; + +pub fn init() { + hook::init(); +} diff --git a/src/recorder/text.rs b/src/recorder/text.rs new file mode 100644 index 00000000..19ed4546 --- /dev/null +++ b/src/recorder/text.rs @@ -0,0 +1,86 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use simd::f16; +use vector::rabitq4::Rabitq4Borrowed; +use vector::rabitq8::Rabitq8Borrowed; +use vector::vect::VectBorrowed; + +pub fn vector_out(vector: VectBorrowed<'_, f32>) -> String { + let mut result = String::from("["); + for x in vector.slice() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} + +pub fn halfvec_out(vector: VectBorrowed<'_, f16>) -> String { + let mut result = String::from("["); + for x in vector.slice() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} + +pub fn rabitq8_out(vector: Rabitq8Borrowed<'_>) -> String { + let mut result = String::new(); + result.push('('); + result.push_str(&vector.sum_of_x2().to_string()); + result.push(','); + result.push_str(&vector.norm_of_lattice().to_string()); + result.push(','); + result.push_str(&vector.sum_of_code().to_string()); + result.push(','); + result.push_str(&vector.sum_of_abs_x().to_string()); + result.push(')'); + result.push('['); + for x in vector.unpacked_code() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} + +pub fn rabitq4_out(vector: Rabitq4Borrowed<'_>) -> String { + let mut result = String::new(); + result.push('('); + result.push_str(&vector.sum_of_x2().to_string()); + result.push(','); + result.push_str(&vector.norm_of_lattice().to_string()); + result.push(','); + result.push_str(&vector.sum_of_code().to_string()); + result.push(','); + result.push_str(&vector.sum_of_abs_x().to_string()); + result.push(')'); + result.push('['); + for x in vector.unpacked_code() { + if !result.ends_with('[') { + result.push(','); + } + result.push_str(&x.to_string()); + } + result.push(']'); + result +} diff --git a/src/recorder/types.rs b/src/recorder/types.rs new file mode 100644 index 00000000..ab40ab02 --- /dev/null +++ b/src/recorder/types.rs @@ -0,0 +1,62 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::recorder::worker::push; +use rand::RngExt; +use std::cell::RefMut; + +pub trait Recorder { + fn is_enabled(&self) -> bool; + fn send(&self, sample: &str); +} + +#[derive(Debug)] +pub struct DefaultRecorder { + pub enable: bool, + pub rate: Option, + pub max_records: u32, + pub index: u32, +} + +pub struct PgRefCell(std::cell::RefCell); + +unsafe impl Send for PgRefCell {} +unsafe impl Sync for PgRefCell {} + +impl PgRefCell { + pub const fn new(x: T) -> Self { + Self(std::cell::RefCell::new(x)) + } + pub fn borrow_mut(&self) -> RefMut<'_, T> { + assert!( + crate::is_main(), + "cannot borrow the value outside main thread" + ); + self.0.borrow_mut() + } +} + +impl Recorder for DefaultRecorder { + fn is_enabled(&self) -> bool { + self.enable + } + fn send(&self, sample: &str) { + if let Some(rate) = self.rate { + let mut rng = rand::rng(); + if rng.random_bool(rate) { + push(self.index, sample, self.max_records); + } + } + } +} diff --git a/src/recorder/worker.rs b/src/recorder/worker.rs new file mode 100644 index 00000000..b56be6d8 --- /dev/null +++ b/src/recorder/worker.rs @@ -0,0 +1,153 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::recorder::types::PgRefCell; +use std::cell::RefMut; +use std::fs; +use std::path::Path; + +// The directory name must start with "pgsql_tmp" to be excluded by pg_basebackup +const RECORDER_DIR: &str = "pgsql_tmp_vchord_sampling"; +const RECORDER_VERSION: u32 = 1; + +static CONNECTION: PgRefCell> = + PgRefCell::>::new(None); + +fn get<'a>(create: bool) -> Option> { + if unsafe { !pgrx::pg_sys::IsBackendPid(pgrx::pg_sys::MyProcPid) } { + return None; + } + let database_oid = unsafe { pgrx::pg_sys::MyDatabaseId.to_u32() }; + if database_oid == 0 { + return None; + } + let mut connection = CONNECTION.borrow_mut(); + if connection.is_none() + && let Err(err) = || -> rusqlite::Result<()> { + if !Path::new(RECORDER_DIR).exists() { + if create { + let _ = fs::create_dir_all(RECORDER_DIR); + } + } + let p = format!("{RECORDER_DIR}/database_{database_oid}.sqlite"); + let mut flags = rusqlite::OpenFlags::default(); + if !create { + flags.remove(rusqlite::OpenFlags::SQLITE_OPEN_CREATE); + } + let mut conn = rusqlite::Connection::open_with_flags(&p, flags)?; + conn.pragma_update(Some("main"), "journal_mode", "WAL")?; + conn.pragma_update(Some("main"), "synchronous", "NORMAL")?; + let tx = conn.transaction()?; + let version: u32 = tx + .pragma_query_value(Some("main"), "user_version", |row| row.get(0)) + .unwrap_or(RECORDER_VERSION); + if version != RECORDER_VERSION && version != 0 { + let mut statement = tx.prepare( + "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'index_%';", + )?; + let tables = statement.query_map((), |row| row.get::(0))?; + for name in tables.into_iter().flatten() { + let drop_statement = format!("DROP TABLE IF EXISTS {name}"); + tx.execute(&drop_statement, ())?; + } + } + tx.pragma_update(Some("main"), "user_version", RECORDER_VERSION)?; + tx.commit()?; + let _ = connection.insert(conn); + Ok(()) + }() + { + if err.sqlite_error_code() == Some(rusqlite::ErrorCode::DatabaseCorrupt) { + delete_database(database_oid); + } + pgrx::debug1!("Recorder: Error initializing database: {}", err); + return None; + } + RefMut::filter_map(connection, |c| c.as_mut()).ok() +} + +pub fn push(index: u32, sample: &str, max_records: u32) { + let mut connection = match get(true) { + Some(c) => c, + None => return, + }; + let init_statement = format!( + " + CREATE TABLE IF NOT EXISTS index_{index} (sample TEXT, create_at REAL); + CREATE INDEX IF NOT EXISTS i ON index_{index} (create_at); + " + ); + let insert_statement = + format!("INSERT INTO index_{index} (sample, create_at) VALUES (?1, unixepoch('subsec'))"); + let count_statement = format!("SELECT COUNT(create_at) FROM index_{index}"); + let maintain_statement = format!( + "DELETE FROM index_{index} WHERE rowid = ( + SELECT rowid FROM index_{index} ORDER BY create_at ASC LIMIT ?1);" + ); + if let Err(err) = || -> rusqlite::Result<()> { + let tx = connection.transaction()?; + tx.execute_batch(&init_statement)?; + tx.prepare_cached(&insert_statement)?.execute((sample,))?; + let records = tx.query_one(&count_statement, (), |row| row.get::(0))?; + if records > max_records { + tx.execute(&maintain_statement, (records - max_records,))?; + } + tx.commit()?; + Ok(()) + }() { + pgrx::debug1!("Recorder: Error pushing sample: {}", err); + } +} + +pub fn delete_index(index: u32) { + let connection = match get(false) { + Some(c) => c, + None => return, + }; + let drop_statement = format!("DROP TABLE IF EXISTS index_{index}"); + if let Err(e) = connection.execute(&drop_statement, ()) { + pgrx::debug1!("Recorder: Error deleting index table: {}", e); + }; +} + +pub fn delete_database(database_oid: u32) { + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite")); + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite-shm")); + let _ = fs::remove_file(format!("{RECORDER_DIR}/database_{database_oid}.sqlite-wal")); +} + +pub fn dump(index: u32) -> Vec { + let connection = match get(false) { + Some(c) => c, + None => return Vec::new(), + }; + let load_statement = format!("SELECT sample FROM index_{index} ORDER BY create_at DESC"); + match || -> rusqlite::Result> { + let mut stmt = connection.prepare(&load_statement)?; + let mut rows = stmt.query(())?; + let mut result = Vec::new(); + while let Some(row) = rows.next()? { + if let Ok(sample) = row.get::(0) { + result.push(sample); + } + } + Ok(result) + }() { + Ok(v) => v, + Err(e) => { + pgrx::debug1!("Recorder: Error loading samples: {}", e); + Vec::new() + } + } +} diff --git a/src/sql/bootstrap.sql b/src/sql/bootstrap.sql index e69de29b..c9f28449 100644 --- a/src/sql/bootstrap.sql +++ b/src/sql/bootstrap.sql @@ -0,0 +1,8 @@ +-- List of shell types + +CREATE TYPE sphere_vector; +CREATE TYPE sphere_halfvec; +CREATE TYPE rabitq8; +CREATE TYPE sphere_rabitq8; +CREATE TYPE rabitq4; +CREATE TYPE sphere_rabitq4; diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 8fdb4add..932660c2 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -1,15 +1,2186 @@ +-- List of types + +CREATE TYPE rabitq8 ( + INPUT = _vchord_rabitq8_in, + OUTPUT = _vchord_rabitq8_out, + TYPMOD_IN = _vchord_rabitq8_typmod_in, + RECEIVE = _vchord_rabitq8_recv, + SEND = _vchord_rabitq8_send, + STORAGE = external +); + +CREATE TYPE rabitq4 ( + INPUT = _vchord_rabitq4_in, + OUTPUT = _vchord_rabitq4_out, + TYPMOD_IN = _vchord_rabitq4_typmod_in, + RECEIVE = _vchord_rabitq4_recv, + SEND = _vchord_rabitq4_send, + STORAGE = external +); + +CREATE TYPE sphere_vector AS ( + center vector, + radius REAL +); + +CREATE TYPE sphere_halfvec AS ( + center halfvec, + radius REAL +); + +CREATE TYPE sphere_rabitq8 AS ( + center rabitq8, + radius REAL +); + +CREATE TYPE sphere_rabitq4 AS ( + center rabitq4, + radius REAL +); + +-- List of internal functions + +CREATE FUNCTION _vchord_rabitq8_operator_maxsim(rabitq8[], rabitq8[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_operator_maxsim_wrapper'; + +CREATE FUNCTION _vchord_rabitq4_operator_maxsim(rabitq4[], rabitq4[]) RETURNS real +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_operator_maxsim_wrapper'; + +-- List of operators + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq8_operator_l2, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq8_operator_ip, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq8_operator_cosine, + LEFTARG = rabitq8, + RIGHTARG = rabitq8, + COMMUTATOR = <=> +); + +CREATE OPERATOR <-> ( + PROCEDURE = _vchord_rabitq4_operator_l2, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <-> +); + +CREATE OPERATOR <#> ( + PROCEDURE = _vchord_rabitq4_operator_ip, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <#> +); + +CREATE OPERATOR <=> ( + PROCEDURE = _vchord_rabitq4_operator_cosine, + LEFTARG = rabitq4, + RIGHTARG = rabitq4, + COMMUTATOR = <=> +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_vector_sphere_l2_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_halfvec_sphere_l2_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq8_sphere_l2_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<->> ( + PROCEDURE = _vchord_rabitq4_sphere_l2_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_vector_sphere_ip_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_halfvec_sphere_ip_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq8_sphere_ip_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<#>> ( + PROCEDURE = _vchord_rabitq4_sphere_ip_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_vector_sphere_cosine_in, + LEFTARG = vector, + RIGHTARG = sphere_vector +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_halfvec_sphere_cosine_in, + LEFTARG = halfvec, + RIGHTARG = sphere_halfvec +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq8_sphere_cosine_in, + LEFTARG = rabitq8, + RIGHTARG = sphere_rabitq8 +); + +CREATE OPERATOR <<=>> ( + PROCEDURE = _vchord_rabitq4_sphere_cosine_in, + LEFTARG = rabitq4, + RIGHTARG = sphere_rabitq4 +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_vector_operator_maxsim, + LEFTARG = vector[], + RIGHTARG = vector[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_halfvec_operator_maxsim, + LEFTARG = halfvec[], + RIGHTARG = halfvec[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq8_operator_maxsim, + LEFTARG = rabitq8[], + RIGHTARG = rabitq8[] +); + +CREATE OPERATOR @# ( + PROCEDURE = _vchord_rabitq4_operator_maxsim, + LEFTARG = rabitq4[], + RIGHTARG = rabitq4[] +); + +-- List of functions + +CREATE FUNCTION sphere(vector, real) RETURNS sphere_vector +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_vector'; + +CREATE FUNCTION sphere(halfvec, real) RETURNS sphere_halfvec +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_halfvec'; + +CREATE FUNCTION sphere(rabitq8, real) RETURNS sphere_rabitq8 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq8'; + +CREATE FUNCTION sphere(rabitq4, real) RETURNS sphere_rabitq4 +IMMUTABLE PARALLEL SAFE LANGUAGE sql AS 'SELECT ROW($1, $2)::sphere_rabitq4'; + +CREATE FUNCTION quantize_to_rabitq8(vector) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION quantize_to_rabitq8(halfvec) RETURNS rabitq8 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq8_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq8) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq8) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq8_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(vector) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_vector_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION quantize_to_rabitq4(halfvec) RETURNS rabitq4 +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_halfvec_quantize_to_rabitq4_wrapper'; + +CREATE FUNCTION dequantize_to_vector(rabitq4) RETURNS vector +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_vector_wrapper'; + +CREATE FUNCTION dequantize_to_halfvec(rabitq4) RETURNS halfvec +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchord_rabitq4_dequantize_to_halfvec_wrapper'; + +CREATE FUNCTION vchordrq_sampled_values(regclass) RETURNS SETOF TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_sampled_values_wrapper'; + +CREATE FUNCTION vchordrq_sampled_queries(regclass) +RETURNS TABLE( + schema_name NAME, + index_name NAME, + table_name NAME, + column_name NAME, + operator NAME, + value TEXT +) +STRICT LANGUAGE plpgsql AS $$ +DECLARE + ext_schema TEXT; + query_text TEXT; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension e + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + query_text := format( + $q$ + WITH index_metadata AS ( + SELECT + NS.nspname AS schema_name, + I.relname AS index_name, + C.relname AS table_name, + PA.attname AS column_name, + OP.oprname AS operator + FROM + pg_catalog.pg_index X + JOIN + pg_catalog.pg_class C ON C.oid = X.indrelid + JOIN + pg_catalog.pg_namespace NS ON C.relnamespace = NS.oid + JOIN + pg_catalog.pg_class I ON I.oid = X.indexrelid + JOIN + pg_catalog.pg_am A ON A.oid = I.relam + LEFT JOIN + pg_catalog.pg_opclass AS OPC ON OPC.oid = X.indclass[0] + LEFT JOIN + pg_catalog.pg_amop AO ON OPC.opcfamily = AO.amopfamily + LEFT JOIN + pg_catalog.pg_operator OP ON OP.oid = AO.amopopr + LEFT JOIN + pg_catalog.pg_attribute PA ON PA.attrelid = X.indrelid AND PA.attnum = X.indkey[0] + WHERE + A.amname = 'vchordrq' + AND AO.amopstrategy = 1 + AND C.relkind = 'r' + AND X.indnatts = 1 + AND X.indexrelid = %1$s + ) + SELECT + im.schema_name, + im.index_name, + im.table_name, + im.column_name, + im.operator, + s.value + FROM + index_metadata im, + LATERAL %2$I.vchordrq_sampled_values(%1$s) AS s(value); + $q$, + $1::oid, + ext_schema + ); + RETURN QUERY EXECUTE query_text; +END; +$$; + +CREATE FUNCTION vchordrq_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhandler_wrapper'; + +CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper'; + +CREATE FUNCTION vchordrq_evaluate_query_recall( + query text, + exact_search boolean default false, + accu_probes TEXT default NULL, + accu_epsilon real default 1.9 +) +RETURNS real +LANGUAGE plpgsql +AS $$ +DECLARE + rough tid[]; + accu tid[]; + match_count integer := 0; + accu_k integer; + recall real; + rough_probes text; +BEGIN + IF query IS NULL OR exact_search IS NULL OR accu_epsilon IS NULL THEN + RETURN NULL; + END IF; + IF query LIKE '%@#%' AND NOT exact_search THEN + RAISE EXCEPTION 'MaxSim operator cannot be used for estimated recall evaluation. Please use exact_search => true.'; + END IF; + IF NOT exact_search THEN + BEGIN + rough_probes := current_setting('vchordrq.probes'); + END; + END IF; + + BEGIN + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + rough; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing ANN query "%": %', query, SQLERRM; + END; + + BEGIN + IF exact_search THEN + SET LOCAL vchordrq.enable_scan = off; + ELSE + IF accu_probes IS NULL THEN + IF rough_probes = '' THEN + accu_probes := ''; + ELSIF position(',' in rough_probes) > 0 THEN + accu_probes := '65535,65535'; + ELSE + accu_probes := '65535'; + END IF; + END IF; + EXECUTE format('SET LOCAL "vchordrq.probes" = %L', accu_probes); + EXECUTE format('SET LOCAL "vchordrq.epsilon" = %L', accu_epsilon); + SET LOCAL vchordrq.max_scan_tuples = -1; + END IF; + EXECUTE + format('SELECT coalesce(array_agg(id), array[]::tid[]) FROM (%s) AS result(id)', query) + INTO + accu; + EXCEPTION WHEN OTHERS THEN + RAISE EXCEPTION 'Error executing Ground Truth query "%": %', query, SQLERRM; + END; + accu_k := cardinality(accu); + IF accu_k = 0 THEN + RAISE WARNING 'Query "%": No results found, returning NaN for recall.', query; + RETURN 'NaN'; + END IF; + SELECT COUNT(*) INTO match_count FROM (SELECT unnest(rough) INTERSECT SELECT unnest(accu)) AS tids; + recall := match_count::real / accu_k::real; + RETURN recall; +END; +$$; + +CREATE FUNCTION vchordg_amhandler(internal) RETURNS index_am_handler +IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_amhandler_wrapper'; + +CREATE FUNCTION vchordg_prewarm(regclass) RETURNS TEXT +STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordg_prewarm_wrapper'; + -- List of access methods -CREATE ACCESS METHOD rabbithole TYPE INDEX HANDLER _rabbithole_amhandler; -COMMENT ON ACCESS METHOD rabbithole IS 'rabbithole index access method'; +CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler; +CREATE ACCESS METHOD vchordg TYPE INDEX HANDLER vchordg_amhandler; -- List of operator families -CREATE OPERATOR FAMILY vector_l2_ops USING rabbithole; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY halfvec_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq8_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY rabitq4_maxsim_ops USING vchordrq; +CREATE OPERATOR FAMILY vector_l2_ops USING vchordg; +CREATE OPERATOR FAMILY vector_ip_ops USING vchordg; +CREATE OPERATOR FAMILY vector_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_l2_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_ip_ops USING vchordg; +CREATE OPERATOR FAMILY halfvec_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq8_cosine_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_l2_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_ip_ops USING vchordg; +CREATE OPERATOR FAMILY rabitq4_cosine_ops USING vchordg; -- List of operator classes CREATE OPERATOR CLASS vector_l2_ops - FOR TYPE vector USING rabbithole FAMILY vector_l2_ops AS + FOR TYPE vector USING vchordrq FAMILY vector_l2_ops AS + OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordrq FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordrq FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordrq_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordrq FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordrq_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordrq FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordrq FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordrq_support_rabitq4_cosine_ops(); + +CREATE OPERATOR CLASS vector_maxsim_ops + FOR TYPE vector[] USING vchordrq FAMILY vector_maxsim_ops AS + OPERATOR 3 @# (vector[], vector[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_vector_maxsim_ops(); + +CREATE OPERATOR CLASS halfvec_maxsim_ops + FOR TYPE halfvec[] USING vchordrq FAMILY halfvec_maxsim_ops AS + OPERATOR 3 @# (halfvec[], halfvec[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_halfvec_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq8_maxsim_ops + FOR TYPE rabitq8[] USING vchordrq FAMILY rabitq8_maxsim_ops AS + OPERATOR 3 @# (rabitq8[], rabitq8[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq8_maxsim_ops(); + +CREATE OPERATOR CLASS rabitq4_maxsim_ops + FOR TYPE rabitq4[] USING vchordrq FAMILY rabitq4_maxsim_ops AS + OPERATOR 3 @# (rabitq4[], rabitq4[]) FOR ORDER BY float_ops, + FUNCTION 1 _vchordrq_support_rabitq4_maxsim_ops(); + +CREATE OPERATOR CLASS vector_l2_ops + FOR TYPE vector USING vchordg FAMILY vector_l2_ops AS OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops, - OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH; + OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_l2_ops(); + +CREATE OPERATOR CLASS vector_ip_ops + FOR TYPE vector USING vchordg FAMILY vector_ip_ops AS + OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_ip_ops(); + +CREATE OPERATOR CLASS vector_cosine_ops + FOR TYPE vector USING vchordg FAMILY vector_cosine_ops AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH, + FUNCTION 1 _vchordg_support_vector_cosine_ops(); + +CREATE OPERATOR CLASS halfvec_l2_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_l2_ops AS + OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_l2_ops(); + +CREATE OPERATOR CLASS halfvec_ip_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_ip_ops AS + OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_ip_ops(); + +CREATE OPERATOR CLASS halfvec_cosine_ops + FOR TYPE halfvec USING vchordg FAMILY halfvec_cosine_ops AS + OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (halfvec, sphere_halfvec) FOR SEARCH, + FUNCTION 1 _vchordg_support_halfvec_cosine_ops(); + +CREATE OPERATOR CLASS rabitq8_l2_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_l2_ops AS + OPERATOR 1 <-> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_l2_ops(); + +CREATE OPERATOR CLASS rabitq8_ip_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_ip_ops AS + OPERATOR 1 <#> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_ip_ops(); + +CREATE OPERATOR CLASS rabitq8_cosine_ops + FOR TYPE rabitq8 USING vchordg FAMILY rabitq8_cosine_ops AS + OPERATOR 1 <=> (rabitq8, rabitq8) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq8, sphere_rabitq8) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq8_cosine_ops(); + +CREATE OPERATOR CLASS rabitq4_l2_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_l2_ops AS + OPERATOR 1 <-> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<->> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_l2_ops(); + +CREATE OPERATOR CLASS rabitq4_ip_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_ip_ops AS + OPERATOR 1 <#> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<#>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_ip_ops(); + +CREATE OPERATOR CLASS rabitq4_cosine_ops + FOR TYPE rabitq4 USING vchordg FAMILY rabitq4_cosine_ops AS + OPERATOR 1 <=> (rabitq4, rabitq4) FOR ORDER BY float_ops, + OPERATOR 2 <<=>> (rabitq4, sphere_rabitq4) FOR SEARCH, + FUNCTION 1 _vchordg_support_rabitq4_cosine_ops(); + +-- List of views + +CREATE VIEW vchordrq_sampled_queries AS +SELECT + record.schema_name, + record.index_name, + record.table_name, + record.column_name, + record.operator, + record.value +FROM + ( + SELECT i.oid + FROM pg_catalog.pg_class AS i + JOIN pg_catalog.pg_index AS ix ON i.oid = ix.indexrelid + JOIN pg_catalog.pg_opclass AS opc ON ix.indclass[0] = opc.oid + JOIN pg_catalog.pg_am AS am ON opc.opcmethod = am.oid + WHERE am.amname = 'vchordrq' + ) AS index_oids +CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); + +-- Exact TileMaxSim over a caller-scoped tensor set. This source registry is +-- deliberately relation-keyed rather than index-keyed: the caller owns ACL, +-- graph, and application filtering decisions. VectorChord accepts the full +-- visible ID set without an artificial count cap, loads its tensor pages, and +-- applies the configured GPU cache policy. + +CREATE TABLE _vchordrq_tilemaxsim_sources ( + source_oid oid PRIMARY KEY, + model_contract_id text NOT NULL CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint NOT NULL CHECK ( + tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_rows_attnum smallint NOT NULL CHECK ( + tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dim_attnum smallint NOT NULL CHECK ( + tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_dtype_attnum smallint NOT NULL CHECK ( + tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + tensor_checksum_attnum smallint NOT NULL CHECK ( + tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_tilemaxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_tilemaxsim_source( + source_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + public_id_is_unique boolean; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor descriptor columns must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be external_ref or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.oid, c.relowner + INTO source_oid, source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid + AND c.relkind IN ('r', 'm'); + IF source_oid IS NULL THEN + RAISE EXCEPTION 'source_relation % must be a table or materialized view', source_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may register a TileMaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL integer or bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'public ID column % must have a non-partial single-key unique index', + public_id_column; + END IF; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a TileMaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL + OR atttypid NOT IN ('integer'::regtype, 'bigint'::regtype) + OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL integer or bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO public_id_is_unique; + IF NOT public_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION 'external_ref sources must not specify a descriptor relation'; + END IF; + tensor_relation_oid := source_oid; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'TileMaxSim tensor source columns must be distinct'; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_tilemaxsim_sources ( + source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + ON CONFLICT (source_oid) DO UPDATE SET + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + source_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_tilemaxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_tilemaxsim_source(source_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_owner oid; + removed_count bigint; +BEGIN + IF source_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_relation::oid AND c.relkind IN ('r', 'm'); + IF ext_schema IS NULL + OR source_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the source relation owner may unregister its TileMaxSim source'; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources WHERE source_oid = $1', + ext_schema + ) USING source_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_tilemaxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_source_info(source_relation regclass) +RETURNS TABLE( + registered_source regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + source_oid oid; + source_owner oid; + bound_model_contract_id text; + bound_storage text; + descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF source_relation IS NULL THEN + RAISE EXCEPTION 'source_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + EXECUTE pg_catalog.format( + 'SELECT source_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_tilemaxsim_sources + WHERE source_oid = $1', + ext_schema + ) INTO + source_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING source_relation::oid; + IF source_oid IS NULL THEN + RAISE EXCEPTION 'TileMaxSim tensor source is not registered for relation %', source_relation; + END IF; + + SELECT c.relowner INTO source_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = source_oid AND c.relkind IN ('r', 'm'); + IF source_owner IS NULL THEN + RAISE EXCEPTION 'registered TileMaxSim source relation is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, source_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, source_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered TileMaxSim source'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, NULL::oid) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = source_oid + AND a.attnum = expected.attnum + AND (expected.atttypid IS NULL OR a.atttypid = expected.atttypid) + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.atttypid IS NOT NULL + OR a.atttypid IN ('integer'::regtype, 'bigint'::regtype); + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered TileMaxSim source columns are invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = source_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim source public ID is no longer unique'; + END IF; + + IF bound_storage = 'external_ref' THEN + IF descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered external_ref TileMaxSim binding is invalid'; + END IF; + tensor_relation_oid := source_oid; + ELSIF bound_storage = 'external_relation' THEN + IF descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered external_relation TileMaxSim binding is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_oid AND c.relkind IN ('r', 'm'); + IF NOT FOUND OR NOT pg_catalog.has_table_privilege(caller_oid, descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor relation is unavailable'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid IN ('integer'::regtype, 'bigint'::regtype) + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered TileMaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + RAISE EXCEPTION 'registered TileMaxSim storage is invalid'; + END IF; + + SELECT count(*) INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 5 THEN + RAISE EXCEPTION 'registered TileMaxSim tensor descriptor columns are invalid'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = source_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + source_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query vector[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_vector_wrapper'; + +CREATE FUNCTION vchordrq_tilemaxsim_rerank( + source_relation regclass, + query halfvec[], + candidate_ids bigint[], + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_tilemaxsim_rerank_halfvec_wrapper'; + +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; +COMMENT ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) +IS 'Exact TileMaxSim over a caller-scoped ID set without an artificial candidate-count cap. ACL and application filtering remain the caller responsibility.'; + +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, vector[], bigint[], integer) TO PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_tilemaxsim_rerank(regclass, halfvec[], bigint[], integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_tilemaxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL OR pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_tilemaxsim_sources', ext_schema) + ) IS NULL THEN + RETURN; + END IF; + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_tilemaxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE ( + d.objid = s.source_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.model_contract_attnum, + s.public_id_attnum, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_ref_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_rows_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dim_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_dtype_attnum END, + CASE WHEN s.storage = ''external_ref'' THEN s.tensor_checksum_attnum END + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid IN ( + s.descriptor_public_id_attnum, + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_tilemaxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_tilemaxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_tilemaxsim_source_sql_drop(); diff --git a/src/types.rs b/src/types.rs deleted file mode 100644 index 788504ab..00000000 --- a/src/types.rs +++ /dev/null @@ -1,43 +0,0 @@ -use serde::{Deserialize, Serialize}; -use validator::Validate; - -#[derive(Debug, Clone, Serialize, Deserialize, Validate)] -#[serde(deny_unknown_fields)] -pub struct RabbitholeIndexingOptions { - #[serde(default = "RabbitholeIndexingOptions::default_nlist")] - #[validate(range(min = 1, max = 1_000_000))] - pub nlist: u32, - #[serde(default = "RabbitholeIndexingOptions::default_spherical_centroids")] - pub spherical_centroids: bool, - #[serde(default = "RabbitholeIndexingOptions::default_residual_quantization")] - pub residual_quantization: bool, - #[serde(default = "RabbitholeIndexingOptions::default_build_threads")] - #[validate(range(min = 1, max = 255))] - pub build_threads: u16, -} - -impl RabbitholeIndexingOptions { - fn default_nlist() -> u32 { - 1000 - } - fn default_spherical_centroids() -> bool { - false - } - fn default_residual_quantization() -> bool { - false - } - fn default_build_threads() -> u16 { - 1 - } -} - -impl Default for RabbitholeIndexingOptions { - fn default() -> Self { - Self { - nlist: Self::default_nlist(), - spherical_centroids: Self::default_spherical_centroids(), - residual_quantization: Self::default_residual_quantization(), - build_threads: Self::default_build_threads(), - } - } -} diff --git a/src/upgrade.rs b/src/upgrade.rs new file mode 100644 index 00000000..119c6802 --- /dev/null +++ b/src/upgrade.rs @@ -0,0 +1,56 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +// Referenced symbols must exist in the dynamic library when dropping functions. +// So we should never remove symbols used by schema, otherwise there will be errors in upgrade. +// Reference: +// * https://www.postgresql.org/message-id/CACX+KaPOzzRHEt4w_=iqKbTpMKjyrUGVng1C749yP3r6dprtcg@mail.gmail.com +// * https://github.com/tensorchord/pgvecto.rs/issues/397 + +macro_rules! symbol { + ($t:ident) => { + paste::paste! { + #[unsafe(no_mangle)] + #[doc(hidden)] + #[pgrx::pg_guard] + extern "C-unwind" fn [<$t _wrapper>](_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::pg_sys::Datum { + pgrx::error!( + "the symbol {} is removed in the extension; please run extension update scripts", + stringify!($t), + ); + } + #[unsafe(no_mangle)] + #[doc(hidden)] + pub extern "C" fn []() -> &'static ::pgrx::pg_sys::Pg_finfo_record { + const V1_API: ::pgrx::pg_sys::Pg_finfo_record = ::pgrx::pg_sys::Pg_finfo_record { + api_version: 1, + }; + &V1_API + } + } + }; +} + +symbol!(_vchord_vector_quantize_to_scalar8_wrapper); +symbol!(_vchord_halfvec_quantize_to_scalar8_wrapper); +symbol!(_vchord_scalar8_operator_cosine_wrapper); +symbol!(_vchord_scalar8_operator_ip_wrapper); +symbol!(_vchord_scalar8_operator_l2_wrapper); +symbol!(_vchord_scalar8_sphere_cosine_in_wrapper); +symbol!(_vchord_scalar8_sphere_ip_in_wrapper); +symbol!(_vchord_scalar8_sphere_l2_in_wrapper); +symbol!(_vchord_typmod_in_65535_wrapper); +symbol!(_vchord_typmod_out_wrapper); +symbol!(_vchord_scalar8_recv_wrapper); +symbol!(_vchord_scalar8_send_wrapper); diff --git a/src/upgrade/mod.rs b/src/upgrade/mod.rs deleted file mode 100644 index 6eb441db..00000000 --- a/src/upgrade/mod.rs +++ /dev/null @@ -1 +0,0 @@ -mod symbols; diff --git a/src/upgrade/symbols.rs b/src/upgrade/symbols.rs deleted file mode 100644 index c275bb10..00000000 --- a/src/upgrade/symbols.rs +++ /dev/null @@ -1,30 +0,0 @@ -// Referenced symbols must exist in the dynamic library when dropping functions. -// So we should never remove symbols used by schema, otherwise there will be errors in upgrade. -// Reference: -// * https://www.postgresql.org/message-id/CACX+KaPOzzRHEt4w_=iqKbTpMKjyrUGVng1C749yP3r6dprtcg@mail.gmail.com -// * https://github.com/tensorchord/pgvecto.rs/issues/397 - -#[allow(unused)] -macro_rules! symbol { - ($t:ident) => { - paste::paste! { - #[no_mangle] - #[doc(hidden)] - #[pgrx::pg_guard] - extern "C" fn [<$t _wrapper>](_fcinfo: pgrx::pg_sys::FunctionCallInfo) -> pgrx::pg_sys::Datum { - pgrx::error!( - "the symbol {} is removed in the extension; please run extension update scripts", - stringify!($t), - ); - } - #[no_mangle] - #[doc(hidden)] - pub extern "C" fn []() -> &'static ::pgrx::pg_sys::Pg_finfo_record { - const V1_API: ::pgrx::pg_sys::Pg_finfo_record = ::pgrx::pg_sys::Pg_finfo_record { - api_version: 1, - }; - &V1_API - } - } - }; -} diff --git a/src/utils/cells.rs b/src/utils/cells.rs deleted file mode 100644 index ff4a31c4..00000000 --- a/src/utils/cells.rs +++ /dev/null @@ -1,21 +0,0 @@ -use std::cell::Cell; - -pub struct PgCell(Cell); - -unsafe impl Send for PgCell {} -unsafe impl Sync for PgCell {} - -impl PgCell { - pub const unsafe fn new(x: T) -> Self { - Self(Cell::new(x)) - } -} - -impl PgCell { - pub fn get(&self) -> T { - self.0.get() - } - pub fn set(&self, value: T) { - self.0.set(value); - } -} diff --git a/src/utils/mod.rs b/src/utils/mod.rs deleted file mode 100644 index 742cee90..00000000 --- a/src/utils/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod cells; diff --git a/taplo.toml b/taplo.toml new file mode 100644 index 00000000..b913433c --- /dev/null +++ b/taplo.toml @@ -0,0 +1,15 @@ +[formatting] +indent_string = " " + +[[rule]] +include = ["**/Cargo.toml"] +keys = [ + "dependencies", + "dev-dependencies", + "build-dependencies", + "target.*.dependencies", + "lints", + "workspace.dependencies", + "workspace.lints", +] +formatting = { reorder_keys = true, reorder_arrays = true, align_comments = true } diff --git a/tests/fail/null.fail b/tests/fail/null.fail new file mode 100644 index 00000000..3878d502 --- /dev/null +++ b/tests/fail/null.fail @@ -0,0 +1,36 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +INSERT INTO t (val) SELECT ('[NaN, Infinity, -Infinity]') FROM generate_series(1, 100); + +statement ok +INSERT INTO t (val) SELECT (NULL) FROM generate_series(1, 100); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +CREATE INDEX rabitq ON t USING vchordrq (val vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX rabitq; + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +---- +DROP TABLE t; diff --git a/tests/general/dequantize.slt b/tests/general/dequantize.slt new file mode 100644 index 00000000..adf30c39 --- /dev/null +++ b/tests/general/dequantize.slt @@ -0,0 +1,19 @@ +query I +SELECT dequantize_to_vector(quantize_to_rabitq8('[1,2,3,4,5,6,7,8]'::vector)) <-> '[1,2,3,4,5,6,7,8]'::vector < 0.07; +---- +t + +query I +SELECT dequantize_to_halfvec(quantize_to_rabitq8('[1,2,3,4,5,6,7,8]'::halfvec)) <-> '[1,2,3,4,5,6,7,8]'::halfvec < 0.07; +---- +t + +query I +SELECT dequantize_to_vector(quantize_to_rabitq4('[1,2,3,4,5,6,7,8]'::vector)) <-> '[1,2,3,4,5,6,7,8]'::vector < 1.00; +---- +t + +query I +SELECT dequantize_to_halfvec(quantize_to_rabitq4('[1,2,3,4,5,6,7,8]'::halfvec)) <-> '[1,2,3,4,5,6,7,8]'::halfvec < 1.00; +---- +t diff --git a/tests/general/distance.slt b/tests/general/distance.slt new file mode 100644 index 00000000..eb798772 --- /dev/null +++ b/tests/general/distance.slt @@ -0,0 +1,29 @@ +query I +SELECT round(('[1,2,3]'::vector <-> '[2,3,4]'::vector):: numeric, 3); +---- +1.732 + +query I +SELECT round(('[1,2,3]'::vector <#> '[2,3,4]'::vector):: numeric, 3); +---- +-20.000 + +query I +SELECT round(('[1,2,3]'::vector <=> '[2,3,4]'::vector):: numeric, 3); +---- +0.007 + +query I +SELECT round(('[1,2,3]'::halfvec <-> '[2,3,4]'::halfvec):: numeric, 3); +---- +1.732 + +query I +SELECT round(('[1,2,3]'::halfvec <#> '[2,3,4]'::halfvec):: numeric, 3); +---- +-20.000 + +query I +SELECT round(('[1,2,3]'::halfvec <=> '[2,3,4]'::halfvec):: numeric, 3); +---- +0.007 diff --git a/tests/general/issue_427.slt b/tests/general/issue_427.slt new file mode 100644 index 00000000..2a06a7d4 --- /dev/null +++ b/tests/general/issue_427.slt @@ -0,0 +1,19 @@ +# https://github.com/tensorchord/pgvecto.rs/issues/427 + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT NULL::vector FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> (SELECT val FROM t LIMIT 1) limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/general/vector.slt b/tests/general/vector.slt new file mode 100644 index 00000000..ce7edf96 --- /dev/null +++ b/tests/general/vector.slt @@ -0,0 +1,23 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/freespace.slt b/tests/vchordg/freespace.slt new file mode 100644 index 00000000..9f3e1ed9 --- /dev/null +++ b/tests/vchordg/freespace.slt @@ -0,0 +1,53 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +3 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 18); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +3 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +DELETE FROM t; + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +VACUUM (INDEX_CLEANUP ON) t; + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 19); + +query I +SELECT pg_relation_size('t_val_idx') / current_setting('block_size')::bigint; +---- +4 + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/index_vector.slt b/tests/vchordg/index_vector.slt new file mode 100644 index 00000000..8e3a4382 --- /dev/null +++ b/tests/vchordg/index_vector.slt @@ -0,0 +1,138 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg (val vector_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_l2_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_ip_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((val::halfvec(64)) halfvec_cosine_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/index_vector_rabitq.slt b/tests/vchordg/index_vector_rabitq.slt new file mode 100644 index 00000000..8837d9fd --- /dev/null +++ b/tests/vchordg/index_vector_rabitq.slt @@ -0,0 +1,256 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 9; +---- +1608 +155 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +1608 +155 +174 +818 +1603 +1629 +60 +218 +1080 +1639 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 9; +---- +1608 +155 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +174 +1603 +818 +1629 +60 +218 +1080 +79 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +174 +818 +1603 +1629 +60 +218 +1080 +79 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +174 +818 +1603 +1629 +60 +218 +1080 +79 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1839 +1568 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/partition.slt b/tests/vchordg/partition.slt new file mode 100644 index 00000000..d42ceeb8 --- /dev/null +++ b/tests/vchordg/partition.slt @@ -0,0 +1,51 @@ +# partition table +statement ok +CREATE TABLE t (val vector(3), category_id int) PARTITION BY LIST(category_id); + +statement ok +CREATE TABLE id_123 PARTITION OF t FOR VALUES IN (1, 2, 3); + +statement ok +CREATE TABLE id_456 PARTITION OF t FOR VALUES IN (4, 5, 6); + +statement ok +CREATE TABLE id_789 PARTITION OF t FOR VALUES IN (7, 8, 9); + +statement ok +INSERT INTO t (val, category_id) +SELECT + ARRAY[random(), random(), random()]::real[], + (random() * 6 + 1)::int +FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordg (val public.vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement error the relation "t_val_idx" is not an index +select vchordg_prewarm('t_val_idx'); + +statement ok +CREATE INDEX ON id_123 USING vchordg (val vector_cosine_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +# partial index +statement ok +CREATE INDEX ON t USING vchordg (val public.vector_ip_ops) WHERE (category_id = 1); + +query I +SELECT COUNT(1) FROM +(SELECT 1 FROM t WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE id_789, id_456, id_123, t; diff --git a/tests/vchordg/reindex.slt b/tests/vchordg/reindex.slt new file mode 100644 index 00000000..77da3a0b --- /dev/null +++ b/tests/vchordg/reindex.slt @@ -0,0 +1,33 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 700); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +REINDEX INDEX t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX CONCURRENTLY t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/sequential_build.slt b/tests/vchordg/sequential_build.slt new file mode 100644 index 00000000..826de42f --- /dev/null +++ b/tests/vchordg/sequential_build.slt @@ -0,0 +1,17 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET max_parallel_workers = 0; + +statement ok +SET max_parallel_maintenance_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordg (val vector_ip_ops); + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/vacuum.slt b/tests/vchordg/vacuum.slt new file mode 100644 index 00000000..4b1a548c --- /dev/null +++ b/tests/vchordg/vacuum.slt @@ -0,0 +1,77 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.667; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/vchordg/vacuum_parallel.slt b/tests/vchordg/vacuum_parallel.slt new file mode 100644 index 00000000..2bb62cad --- /dev/null +++ b/tests/vchordg/vacuum_parallel.slt @@ -0,0 +1,53 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET min_parallel_index_scan_size = 1; + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +CREATE INDEX ON t USING vchordg (val vector_l2_ops); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 100); + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.667; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/approximate.slt b/tests/vchordrq/approximate.slt new file mode 100644 index 00000000..c19f6bc6 --- /dev/null +++ b/tests/vchordrq/approximate.slt @@ -0,0 +1,57 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.internal] +lists = [33] +spherical_centroids = false +kmeans_dimension = 2 +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_cosine_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +kmeans_algorithm.hierarchical = {} +$$); + +statement ok +SET vchordrq.probes = '16'; + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt new file mode 100644 index 00000000..87cf85e7 --- /dev/null +++ b/tests/vchordrq/cost_estimator.slt @@ -0,0 +1,597 @@ +# Tests for amcostestimate in src/index/vchordrq/am/mod.rs. +# +# The cost estimator must: +# 1. Use the baserel's planner-computed row estimate to derive filter +# selectivity, then report index selectivity from the candidate count +# that the scan is expected to retrieve. +# 2. Cap the candidate-processing term (`next_count`) by the IVF candidate +# budget (`node_count`), so absurdly small selectivity cannot inflate +# the reported cost past the work the index actually does. +# 3. Degrade safely when planner stats are missing or extreme. +# +# --------------------------------------------------------------------------- +# Setup +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = off; + +statement ok +SET max_parallel_workers_per_gather = 0; + +statement ok +SET random_page_cost = 1.1; + +statement ok +CREATE TABLE cost_test ( + id int PRIMARY KEY, + category int NOT NULL, + v vector(3) NOT NULL +); + +# 10 000 rows, 100 evenly-sized categories. category=k matches 1% of rows. +# Vectors are deterministic per row so the test is reproducible. +statement ok +INSERT INTO cost_test +SELECT i, + (i % 100), + ARRAY[ + (i % 97) / 97.0, + (i % 89) / 89.0, + (i % 83) / 83.0 + ]::real[]::vector +FROM generate_series(1, 10000) i; + +statement ok +CREATE INDEX cost_test_v ON cost_test USING vchordrq (v vector_l2_ops); + +statement ok +CREATE INDEX cost_test_cat ON cost_test (category); + +# A predicate with an explicitly high procost that PG can't optimize away. +# Three things matter here: +# * `LANGUAGE plpgsql` — prevents SQL inlining (an inlined `IS NOT NULL` +# check on a PRIMARY KEY column gets folded to `true` and dropped from +# the plan, defeating the test). +# * `COST 10000` — sets `pg_proc.procost` so cost_index() charges +# 10000 * cpu_operator_cost ≈ 25 cost-units per tuple of filter eval. +# This mimics PostGIS ST_DWithin, JSONB containment, expensive regex, +# etc. — the real-world shapes where the bug bites. +# * Body is `mod($1,7) <> 99` so the result is always true but PG can't +# prove that across the function boundary, so the filter survives +# constant-folding. +statement ok +CREATE OR REPLACE FUNCTION slow_true(int) RETURNS boolean + LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE COST 10000 AS +$$ BEGIN RETURN mod($1, 7) <> 99; END $$; + +statement ok +ANALYZE cost_test; + +# Helper: returns the index name used at the top of the plan, or NULL if +# no Index Scan / Bitmap Index Scan appears. Robust across PG versions. +statement ok +CREATE OR REPLACE FUNCTION plan_top_index(query text) RETURNS text + LANGUAGE plpgsql AS $$ +DECLARE + line text; + m text[]; +BEGIN + FOR line IN EXECUTE 'EXPLAIN (FORMAT TEXT, COSTS OFF) ' || query LOOP + -- 'Index Scan using on ...' or 'Index Only Scan using on ...' + m := regexp_match(line, 'Index (?:Only )?Scan using (\S+) on '); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + -- 'Bitmap Index Scan on ' + m := regexp_match(line, 'Bitmap Index Scan on (\S+)'); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns the top plan node's total cost. +statement ok +CREATE OR REPLACE FUNCTION top_total_cost(query text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns total cost for the first plan node matching `pat`. +statement ok +CREATE OR REPLACE FUNCTION plan_node_total_cost(query text, pat text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + IF rec LIKE pat THEN + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END IF; + END LOOP; + RETURN NULL; +END $$; + +# --------------------------------------------------------------------------- +# Case 1: pure ORDER BY, no filter. +# Vchord must win. (Passes pre- and post-fix; sanity check.) +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# --------------------------------------------------------------------------- +# Case 2: tight filter (selectivity 1%) with an alternative index. +# Btree must still beat vchord — the IVF candidate budget is too large +# to be competitive at this selectivity. (Passes pre- and post-fix.) +# This guards against over-correction: the fix must not make vchord win +# on tight filters, which was the original motivation for PR #234. +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_cat + +statement ok +DROP INDEX cost_test_cat; + +# --------------------------------------------------------------------------- +# Case 3: selective heap filter combined with an expensive predicate and no +# supporting non-vector index. Pre-fix: index_selectivity=1.0 charges +# filter eval against every table row, so seq scan + sort can look cheaper. +# Post-fix: vchord estimates the candidates needed for the LIMIT and wins. +# The node-cost floor guards against under-pricing the scan as if it fetched +# only the final surviving rows. +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) > 10000; +---- +t + +statement ok +SET enable_seqscan = off; + +# --------------------------------------------------------------------------- +# Case 4: filter on a column with no supporting non-vector index, combined +# with the expensive predicate. The only alternatives are seq scan (disabled) +# or vchord. Vchord must be chosen. Pre-fix this still picks vchord because +# seq is disabled, but the *reported* cost must drop enough that LIMIT-aware +# planners above this node (joins, gather, partitions) don't overestimate. +# We assert plan shape + that the cost is bounded below 1e9 (disable_cost). +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# disable_cost in PG 14-18 is 1.0e10; a healthy estimate is well below that. +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) < 1e9; +---- +t + +# --------------------------------------------------------------------------- +# Case 5: no LIMIT clause. PlannerInfo.limit_tuples is -1, so the +# estimator should fall back to the IVF candidate budget instead of doing +# LIMIT/selectivity math. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 6: LIMIT larger than the number of expected survivors. +# The `next_count.min(node_count)` clamp must hold — the AM cannot claim +# to process more candidates than the IVF visits. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 7: near-zero filter selectivity. The fix clamps to 1e-9 instead of +# 0 so 1/selectivity does not produce +Inf. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 8: empty table. baserel->tuples is 0, fix falls back to +# selectivity = 1.0. Plan and result must not crash. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_empty (id int, v vector(3)); + +statement ok +CREATE INDEX ON cost_test_empty USING vchordrq (v vector_l2_ops); + +statement ok +ANALYZE cost_test_empty; + +query I +SELECT count(*) FROM ( + SELECT id FROM cost_test_empty ORDER BY v <-> '[0,0,0]'::vector LIMIT 10 +) t; +---- +0 + +# --------------------------------------------------------------------------- +# Case 9: never-ANALYZE'd table. baserel->tuples may be negative or zero +# in the planner's view. The fallback in the fix returns 1.0, matching the +# pre-fix behavior exactly so we don't regress on cold tables. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_cold (id int, v vector(3)); + +statement ok +INSERT INTO cost_test_cold +SELECT i, ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 100) i; + +statement ok +CREATE INDEX ON cost_test_cold USING vchordrq (v vector_l2_ops); + +# Intentionally no ANALYZE. + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test_cold ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 10: partial vchord index. The filter-selectivity estimate reflects +# baserestrictinfo, including clauses that overlap the partial-index +# predicate. We verify that this builds, plans, and runs. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_partial (id int, kind int, v vector(3)); + +statement ok +INSERT INTO cost_test_partial +SELECT i, (i % 3), ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 2000) i; + +statement ok +CREATE INDEX cost_test_partial_v ON cost_test_partial USING vchordrq (v vector_l2_ops) + WHERE kind = 0; + +statement ok +ANALYZE cost_test_partial; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_partial WHERE kind = 0 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 5' +); +---- +cost_test_partial_v + +# --------------------------------------------------------------------------- +# Case 11: vchordrq.enable_scan = off must still disable the path +# regardless of selectivity. Re-enable seqscan inside this case so there is +# a non-disabled alternative for the planner to fall back to — otherwise on +# PG18 every path looks "disabled" and the planner picks vchord by default. +# Use IS DISTINCT FROM so a seq-scan top (plan_top_index returns NULL) +# still compares as "not vchord." +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +statement ok +SET vchordrq.enable_scan = off; + +query B +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS DISTINCT FROM 'cost_test_v'; +---- +t + +statement ok +RESET vchordrq.enable_scan; + +statement ok +RESET enable_seqscan; + +# --------------------------------------------------------------------------- +# Case 12: MaxSim must have a nonzero, query-token-aware, backend-aware cost. +# The old special branch returned total_cost=0 and selectivity=1 for every +# MaxSim path, hiding all token expansion and exact rerank work. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_maxsim ( + id int PRIMARY KEY, + v vector(3)[] NOT NULL +); + +statement ok +INSERT INTO cost_test_maxsim +SELECT i, ARRAY[ + '[1,0,0]'::vector, + '[0,1,0]'::vector, + '[0,0,1]'::vector, + '[0.5,0.5,0]'::vector +] +FROM generate_series(1, 1000) i; + +statement ok +CREATE INDEX cost_test_maxsim_v +ON cost_test_maxsim +USING vchordrq (v vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [4] +$$); + +statement ok +ANALYZE cost_test_maxsim; + +statement ok +SET enable_seqscan = off; + +statement ok +SET vchordrq.probes = '4'; + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +CREATE TEMP TABLE maxsim_cost_observations ( + name text PRIMARY KEY, + value double precision NOT NULL +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 1; + +# A newly built index has a native indexed-vector count. The document-token +# GUC is only a compatibility fallback for pre-statistics indexes, so changing +# it must not change this index's cost. +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4096; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4096', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4') + = + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4096'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_1', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 64; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_64', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_64') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_1'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 32; + +statement ok +SET vchordrq.maxsim_candidate_limit = 16; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_16', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_candidate_limit = 256; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_256', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_256') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_16'); +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); +---- +cost_test_maxsim_v + +statement ok +RESET vchordrq.maxsim_backend; + +statement ok +RESET vchordrq.maxsim_candidate_limit; + +statement ok +RESET vchordrq.maxsim_planner_query_tokens; + +statement ok +RESET vchordrq.maxsim_planner_document_tokens; + +statement ok +RESET vchordrq.probes; + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + +statement ok +DROP TABLE cost_test; + +statement ok +DROP TABLE cost_test_empty; + +statement ok +DROP TABLE cost_test_cold; + +statement ok +DROP TABLE cost_test_partial; + +statement ok +DROP TABLE cost_test_maxsim; + +statement ok +DROP FUNCTION slow_true(int); + +statement ok +DROP FUNCTION plan_top_index(text); + +statement ok +DROP FUNCTION plan_node_total_cost(text, text); + +statement ok +DROP FUNCTION top_total_cost(text); + +statement ok +RESET enable_seqscan; + +statement ok +RESET max_parallel_workers_per_gather; + +statement ok +RESET random_page_cost; diff --git a/tests/vchordrq/external_build.slt b/tests/vchordrq/external_build.slt new file mode 100644 index 00000000..81dc35d8 --- /dev/null +++ b/tests/vchordrq/external_build.slt @@ -0,0 +1,117 @@ +statement ok +CREATE TABLE t (val0 vector(3), val1 halfvec(3)); + +statement ok +INSERT INTO t (val0, val1) +SELECT + ARRAY[random(), random(), random()]::real[]::vector, + ARRAY[random(), random(), random()]::real[]::halfvec +FROM generate_series(1, 100); + +statement ok +CREATE TABLE vector_centroid (id integer, parent integer, vector vector(3)); + +statement ok +INSERT INTO vector_centroid (id, vector) VALUES + (0, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +statement ok +CREATE TABLE halfvec_centroid (id integer, parent integer, vector halfvec(3)); + +statement ok +INSERT INTO halfvec_centroid (id, vector) VALUES + (0, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +statement ok +CREATE TABLE real_centroid (id integer, parent integer, vector real[]); + +statement ok +INSERT INTO real_centroid (id, vector) VALUES + (0, '{1.0, 0.0, 0.0}'), + (1, '{0.0, 1.0, 0.0}'), + (2, '{0.0, 0.0, 1.0}'); + +statement ok +CREATE TABLE bad_type_centroid (id integer, parent integer, vector integer); + +statement ok +INSERT INTO bad_type_centroid (id, vector) VALUES + (0, 0), + (1, 0), + (2, 0); + +statement ok +CREATE TABLE bad_duplicate_id (id integer, parent integer, vector vector(3)); + +statement ok +INSERT INTO bad_duplicate_id (id, vector) VALUES + (1, '[1.0, 0.0, 0.0]'), + (1, '[0.0, 1.0, 0.0]'), + (2, '[0.0, 0.0, 1.0]'); + +# external build for vector column + +statement ok +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.vector_centroid' +$$); + +# external build for halfvec column + +statement ok +CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.vector_centroid' +$$); + +# external build for halfvec column by a halfvec table + +statement ok +CREATE INDEX ON t USING vchordrq (val1 halfvec_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.halfvec_centroid' +$$); + +# external build for halfvec column by a real[] table + +statement ok +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.real_centroid' +$$); + +# failed: bad vector data type + +statement error cannot cast type integer to (.*)vector +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.bad_type_centroid' +$$); + +# failed: duplicate id + +statement error external build: there are at least two lines have same id, id = 1 +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.external] +table = 'public.bad_duplicate_id' +$$); + +statement ok +DROP TABLE t, vector_centroid, halfvec_centroid, real_centroid, bad_type_centroid, bad_duplicate_id; \ No newline at end of file diff --git a/tests/vchordrq/external_build_sql_inject.slt b/tests/vchordrq/external_build_sql_inject.slt new file mode 100644 index 00000000..6129ede9 --- /dev/null +++ b/tests/vchordrq/external_build_sql_inject.slt @@ -0,0 +1,11 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement error table name is not well-formed +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.external.table = "s; SELECT s" +$$); + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/filter_rerank_in_index.slt b/tests/vchordrq/filter_rerank_in_index.slt new file mode 100644 index 00000000..69c6e9e5 --- /dev/null +++ b/tests/vchordrq/filter_rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# non-heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + simple prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + simple prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/filter_rerank_in_table.slt b/tests/vchordrq/filter_rerank_in_table.slt new file mode 100644 index 00000000..c556207b --- /dev/null +++ b/tests/vchordrq/filter_rerank_in_table.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + simple prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + simple prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'prefetch_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/index.slt b/tests/vchordrq/index.slt new file mode 100644 index 00000000..0c95e8c8 --- /dev/null +++ b/tests/vchordrq/index.slt @@ -0,0 +1,60 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement error +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +unknown_options=true +$$); + +# multiple index on single column +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = true +[build.internal] +lists = [33] +spherical_centroids = false +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_cosine_ops) +WITH (options = $$ +residual_quantization = false +[build.internal] +lists = [33] +spherical_centroids = true +$$); + +statement ok +SET vchordrq.probes = '16'; + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_multivector.slt b/tests/vchordrq/index_multivector.slt new file mode 100644 index 00000000..044835dc --- /dev/null +++ b/tests/vchordrq/index_multivector.slt @@ -0,0 +1,75 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)[2]); + +statement ok +INSERT INTO t (val) +SELECT + ARRAY[ + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector), + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) + ] +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_maxsim_ops); + +query I +SELECT index FROM t ORDER BY +val @# ARRAY[ + array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector, + array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector +] +LIMIT 10; +---- +1207 +919 +1639 +1821 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)[2]) halfvec_maxsim_ops); + +query I +SELECT index FROM t ORDER BY +val::halfvec(64)[] @# ARRAY[ + array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec, + array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec +] +LIMIT 10; +---- +1207 +919 +1639 +1821 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_multivector_rabitq.slt b/tests/vchordrq/index_multivector_rabitq.slt new file mode 100644 index 00000000..0b7d824b --- /dev/null +++ b/tests/vchordrq/index_multivector_rabitq.slt @@ -0,0 +1,89 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE pg_temp.t (index serial primary key, val vector(64)[2]); + +statement ok +INSERT INTO pg_temp.t (val) +SELECT + ARRAY[ + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector), + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) + ] +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE FUNCTION pg_temp.quantize_to_rabitq8(val vector[]) +RETURNS rabitq8[] AS $$ + SELECT array_agg(public.quantize_to_rabitq8(x)) + FROM unnest(val) AS x; +$$ LANGUAGE sql IMMUTABLE; + +statement ok +CREATE FUNCTION pg_temp.quantize_to_rabitq4(val vector[]) +RETURNS rabitq4[] AS $$ + SELECT array_agg(public.quantize_to_rabitq4(x)) + FROM unnest(val) AS x; +$$ LANGUAGE sql IMMUTABLE; + +statement ok +CREATE INDEX ti ON pg_temp.t USING vchordrq ((pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[2]) rabitq8_maxsim_ops); + +query I +SELECT index FROM pg_temp.t ORDER BY +pg_temp.quantize_to_rabitq8(val)::rabitq8(64)[] @# ARRAY[ + quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), + quantize_to_rabitq8(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) +] +LIMIT 10; +---- +1207 +919 +1821 +1639 +174 +79 +1076 +1125 +239 +194 + +statement ok +DROP INDEX pg_temp.ti; + +statement ok +CREATE INDEX ti ON pg_temp.t USING vchordrq ((pg_temp.quantize_to_rabitq4(val)::rabitq4(64)[2]) rabitq4_maxsim_ops); + +query I +SELECT index FROM pg_temp.t ORDER BY +pg_temp.quantize_to_rabitq4(val)::rabitq4(64)[] @# ARRAY[ + quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector), + quantize_to_rabitq4(array_cat(ARRAY[0.3, 0.4], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) +] +LIMIT 10; +---- +1207 +919 +1821 +1639 +537 +174 +79 +194 +1920 +239 + +statement ok +DROP INDEX pg_temp.ti; + +statement ok +DROP TABLE pg_temp.t; diff --git a/tests/vchordrq/index_vector.slt b/tests/vchordrq/index_vector.slt new file mode 100644 index 00000000..af69591f --- /dev/null +++ b/tests/vchordrq/index_vector.slt @@ -0,0 +1,138 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_l2_ops); + +query I +SELECT index FROM t ORDER BY val <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_ip_ops); + +query I +SELECT index FROM t ORDER BY val <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq (val vector_cosine_ops); + +query I +SELECT index FROM t ORDER BY val <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_l2_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <-> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_ip_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <#> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((val::halfvec(64)) halfvec_cosine_ops); + +query I +SELECT index FROM t ORDER BY val::halfvec <=> array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/index_vector_rabitq.slt b/tests/vchordrq/index_vector_rabitq.slt new file mode 100644 index 00000000..33e7bc16 --- /dev/null +++ b/tests/vchordrq/index_vector_rabitq.slt @@ -0,0 +1,258 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (index serial primary key, val vector(64)); + +statement ok +INSERT INTO t (val) +SELECT + l2_normalize(ARRAY( + SELECT + ('x' || substring(md5((64 * i + j)::text), 1, 16))::bit(64)::bigint / 18446744073709551615.0 + FROM generate_series(1, 64) d(j) + )::vector) +FROM generate_series(1, 2048) s(i); + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +1643 +174 +1603 +818 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +1608 +155 +1643 +174 +818 +1603 +1629 +60 +218 +1080 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1839 +1568 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +CREATE INDEX ti ON t USING vchordrq ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops); + +query I +SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10; +---- +155 +1643 +174 +818 +1207 +1608 +1940 +331 +1568 +1839 + +statement ok +DROP INDEX ti; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/internal_build_kmeans.slt b/tests/vchordrq/internal_build_kmeans.slt new file mode 100644 index 00000000..5c5a7518 --- /dev/null +++ b/tests/vchordrq/internal_build_kmeans.slt @@ -0,0 +1,19 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [1000] +build.internal.kmeans_algorithm.lloyd = {} +$$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [1000] +build.internal.kmeans_algorithm.hierarchical = {} +$$); + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/issue_335.slt b/tests/vchordrq/issue_335.slt new file mode 100644 index 00000000..7a5dfd9e --- /dev/null +++ b/tests/vchordrq/issue_335.slt @@ -0,0 +1,37 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3), visible boolean default true); + +statement ok +INSERT INTO items (embedding) SELECT ARRAY[i * 0.001, i * 0.001, i * 0.001]::real[] FROM generate_series(1, 1000) s(i); + +statement ok +CREATE INDEX ON items USING vchordrq (embedding vector_l2_ops); + +statement ok +SET vchordrq.prefilter = on; + +query I +SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3; +---- +3 +4 +2 + +statement ok +UPDATE items SET visible = false WHERE id = 3; + +statement ok +UPDATE items SET visible = true WHERE id = 3; + +query I +SELECT id FROM items WHERE visible = true ORDER BY embedding <-> '[0.0031,0.0031,0.0031]' LIMIT 3; +---- +3 +4 +2 + +statement ok +DROP TABLE items; diff --git a/tests/vchordrq/maxsim_correctness.slt b/tests/vchordrq/maxsim_correctness.slt new file mode 100644 index 00000000..a5d9652f --- /dev/null +++ b/tests/vchordrq/maxsim_correctness.slt @@ -0,0 +1,221 @@ +# Deterministic MaxSim semantics and input-boundary coverage. + +statement ok +SET enable_seqscan TO off; + +# @# is a distance: negative late-interaction similarity, ordered ascending. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# MaxSim is asymmetric: the right-hand array contains query vectors. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,1]'::vector] +)::numeric, 3); +---- +-1.000 + +query I +SELECT round(( + ARRAY['[1,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# halfvec uses the same sign and orientation contract. +query I +SELECT round(( + ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] + @# ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] +)::numeric, 3); +---- +-2.000 + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY[]::vector[] @# ARRAY['[1,0]'::vector]; + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY['[1,0]'::vector] @# ARRAY[]::vector[]; + +statement error MaxSim arrays must not contain NULL vectors +SELECT ARRAY['[1,0]'::vector, NULL::vector] + @# ARRAY['[1,0]'::vector]; + +statement error dimension is not matched +SELECT ARRAY['[1,0]'::vector] @# ARRAY['[1,0,0]'::vector]; + +statement error MaxSim arrays cannot contain more than 65536 vectors +SELECT ARRAY['[1]'::vector] + @# ARRAY(SELECT '[1]'::vector FROM generate_series(1, 65537)); + +# Indexing an empty document array must fail instead of silently omitting it. +statement ok +CREATE TABLE maxsim_invalid_document (val vector(2)[]); + +statement ok +INSERT INTO maxsim_invalid_document VALUES (ARRAY[]::vector[]); + +statement error MaxSim arrays must contain at least one vector +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +TRUNCATE maxsim_invalid_document; + +statement ok +INSERT INTO maxsim_invalid_document +SELECT ARRAY( + SELECT '[1,0]'::vector + FROM generate_series(1, 65537) +); + +statement error MaxSim arrays cannot contain more than 65536 vectors +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +DROP TABLE maxsim_invalid_document; + +# Lock down current ranking with a non-flat index before the Phase 3 refactor. +statement ok +CREATE TABLE maxsim_deterministic ( + id integer primary key, + val vector(2)[] not null +); + +statement ok +INSERT INTO maxsim_deterministic VALUES + (1, ARRAY['[1,0]'::vector, '[0,1]'::vector]), + (2, ARRAY['[0.8,0]'::vector, '[0,0.8]'::vector]), + (3, ARRAY['[0.5,0.5]'::vector]), + (4, ARRAY['[-1,0]'::vector, '[0,-1]'::vector]); + +statement ok +CREATE INDEX maxsim_deterministic_idx +ON maxsim_deterministic +USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [2] +$$); + +statement ok +SET vchordrq.probes = '2'; + +statement ok +SET vchordrq.maxsim_refine = 100; + +statement ok +SET vchordrq.maxsim_threshold = 0; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_candidate_limit = 2; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement error exact MaxSim requires a positive vchordrq.maxsim_candidate_limit +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_candidate_limit = 4; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +query I +SELECT round((val @# ARRAY['[1,0]'::vector, '[0,1]'::vector])::numeric, 3) +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +-2.000 +-1.600 +-1.000 +0.000 + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error GPU MaxSim transport error: endpoint is empty +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_backend = 'auto'; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +# An empty query must also fail on the index scan path. +statement error MaxSim arrays must contain at least one vector +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY[]::vector[] +LIMIT 1; + +statement error dimension is not matched +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0,0]'::vector] +LIMIT 1; + +statement ok +DROP TABLE maxsim_deterministic; diff --git a/tests/vchordrq/maxsim_source_registry.slt b/tests/vchordrq/maxsim_source_registry.slt new file mode 100644 index 00000000..c408a469 --- /dev/null +++ b/tests/vchordrq/maxsim_source_registry.slt @@ -0,0 +1,341 @@ +# Phase 3B tensor-source registration must bind to a real MaxSim index by +# relation/attribute OID, reject incompatible descriptors, and fail closed when +# a bound index/table column is dropped. + +statement ok +CREATE TABLE maxsim_source_test ( + id bigint PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2)[] NOT NULL, + single_embedding vector(2) NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL, + application_note text +); + +statement ok +INSERT INTO maxsim_source_test VALUES ( + 1, + 'colqwen@test', + ARRAY['[1,0]'::vector, '[0,1]'::vector], + '[1,0]'::vector, + 'tensor://1', + 2, + 2, + 'float32', + 'sha256:test' +); + +statement ok +CREATE INDEX maxsim_source_test_idx +ON maxsim_source_test +USING vchordrq (embedding vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement ok +CREATE INDEX maxsim_source_wrong_idx +ON maxsim_source_test +USING vchordrq (single_embedding vector_l2_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement error not a valid single-key vchordrq MaxSim index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_wrong_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement error tensor rows column tensor_dtype must be a NOT NULL integer column +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_dtype', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_REF', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +query TT +SELECT s.storage, s.model_contract_id +FROM _vchordrq_maxsim_sources AS s +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +external_ref colqwen@test + +query TTT +SELECT source_storage, tensor_ref_column::text, public_id_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_ref tensor_ref id + +# Production external tensors may live in a compact descriptor relation keyed +# by the stable application public ID, avoiding a rewrite of the indexed heap. +statement ok +CREATE TABLE maxsim_descriptor_no_unique ( + public_id bigint NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement error must have a non-partial single-key unique index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_no_unique'::regclass, + descriptor_public_id_column => 'public_id' +); + +statement ok +CREATE TABLE maxsim_descriptor_test ( + public_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO maxsim_descriptor_test VALUES ( + 1, 'tensor://1', 2, 2, 'float32', 'sha256:test' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'public_id' +); + +query TTTT +SELECT source_storage, descriptor_relation::text, + descriptor_public_id_column::text, tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_relation maxsim_descriptor_test public_id tensor_ref + +statement ok +ALTER TABLE maxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +statement ok +ALTER TABLE maxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +# Unbound application columns are outside the registry contract. Dropping one +# must not invalidate an otherwise live source binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN application_note; + +query T +SELECT model_contract_id +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +colqwen@test + +# The explicit score surface validates bounded work and exact query/index type +# before any sidecar access. +statement error candidate_limit must be between 1 and 65536 +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 0, + 1 +); + +statement error top_k must be positive and no greater than candidate_limit +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 1, + 2 +); + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error MaxSim tensor kind or dimension is not matched +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::halfvec], + 1, + 1 +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_ref + +# Attribute numbers survive renames, so bindings do not depend on SQL text. +statement ok +ALTER TABLE maxsim_source_test RENAME tensor_ref TO tensor_location; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_location + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +# Type-altering DDL does not drop an attribute, so runtime resolution must +# revalidate the complete descriptor and fail closed. +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE varchar(16); + +statement error registered MaxSim tensor source has invalid descriptor columns +SELECT * FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); + +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE text; + +# Dropping a bound descriptor column invalidates the complete binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +query B +SELECT tensor_ref_attnum IS NULL +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +f + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement ok +DROP INDEX maxsim_source_test_idx; + +query I +SELECT count(*) FROM _vchordrq_maxsim_sources; +---- +0 + +statement ok +DROP TABLE maxsim_source_test; + +statement ok +DROP TABLE maxsim_descriptor_test, maxsim_descriptor_no_unique; diff --git a/tests/vchordrq/multivector.slt b/tests/vchordrq/multivector.slt new file mode 100644 index 00000000..5271e46c --- /dev/null +++ b/tests/vchordrq/multivector.slt @@ -0,0 +1,55 @@ +statement ok +CREATE TABLE t (id integer, val vector(2)[]); + +statement ok +INSERT INTO t (id, val) +SELECT id, + ARRAY[ + ARRAY[cos(((id + 0) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 0) % 10000) / 10000.0 * 6.283185307179586)]::vector, + ARRAY[cos(((id + 22) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 22) % 10000) / 10000.0 * 6.283185307179586)]::vector, + ARRAY[cos(((id + 777) % 10000) / 10000.0 * 6.283185307179586), sin(((id + 777) % 10000) / 10000.0 * 6.283185307179586)]::vector + ] +FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX t_val_idx ON t USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +build.internal.lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +statement ok +SET vchordrq.maxsim_refine = 3000; + +statement ok +SET enable_seqscan TO off; + +query I +SELECT id FROM t ORDER BY val @# ARRAY['[0.7197411498053302, 0.6942425205048314]'::vector, '[0.10645067063129976, 0.9943179847122079]'::vector] limit 18; +---- +1387 +1388 +1386 +1389 +1385 +1390 +1384 +1391 +1383 +1392 +1382 +1393 +1381 +1394 +1380 +1395 +1379 +1396 + +statement ok +DROP INDEX t_val_idx; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/options.slt b/tests/vchordrq/options.slt new file mode 100644 index 00000000..972917c1 --- /dev/null +++ b/tests/vchordrq/options.slt @@ -0,0 +1,46 @@ +statement ok +SET enable_seqscan TO off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) VALUES ('[1,1,1]'), ('[2,2,2]'); + +statement ok +CREATE INDEX i ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ +build.internal.lists = [2] +$$, probes = '2'); + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] +[2,2,2] + +statement ok +SET vchordrq.probes TO '1'; + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] + +statement ok +ALTER INDEX i SET (probes = '0'); + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- +[1,1,1] + +statement ok +SET vchordrq.probes TO DEFAULT; + +query T +SELECT val FROM t ORDER BY val <-> '[0,0,0]' LIMIT 10; +---- + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/partition.slt b/tests/vchordrq/partition.slt new file mode 100644 index 00000000..1e681c5e --- /dev/null +++ b/tests/vchordrq/partition.slt @@ -0,0 +1,51 @@ +# partition table +statement ok +CREATE TABLE t (val vector(3), category_id int) PARTITION BY LIST(category_id); + +statement ok +CREATE TABLE id_123 PARTITION OF t FOR VALUES IN (1, 2, 3); + +statement ok +CREATE TABLE id_456 PARTITION OF t FOR VALUES IN (4, 5, 6); + +statement ok +CREATE TABLE id_789 PARTITION OF t FOR VALUES IN (7, 8, 9); + +statement ok +INSERT INTO t (val, category_id) +SELECT + ARRAY[random(), random(), random()]::real[], + (random() * 6 + 1)::int +FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordrq (val public.vector_l2_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement error the relation "t_val_idx" is not an index +select vchordrq_prewarm('t_val_idx'); + +statement ok +CREATE INDEX ON id_123 USING vchordrq (val vector_cosine_ops); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +# partial index +statement ok +CREATE INDEX ON t USING vchordrq (val public.vector_ip_ops) WHERE (category_id = 1); + +query I +SELECT COUNT(1) FROM +(SELECT 1 FROM t WHERE (category_id = 1) ORDER BY val <#> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE id_789, id_456, id_123, t; diff --git a/tests/vchordrq/pg17/filter_rerank_in_index.slt b/tests/vchordrq/pg17/filter_rerank_in_index.slt new file mode 100644 index 00000000..3b0e9b10 --- /dev/null +++ b/tests/vchordrq/pg17/filter_rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# non-heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + stream prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# non-heap rerank + stream prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/pg17/filter_rerank_in_table.slt b/tests/vchordrq/pg17/filter_rerank_in_table.slt new file mode 100644 index 00000000..19f0bf60 --- /dev/null +++ b/tests/vchordrq/pg17/filter_rerank_in_table.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +# heap rerank + no prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + stream prefetch + postfilter +statement ok +SET vchordrq.prefilter to off; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + no prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_rerank to 'read_buffer'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +# heap rerank + stream prefetch + prefilter +statement ok +SET vchordrq.prefilter to on; + +statement ok +SET vchordrq.io_search to 'read_stream'; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/pin.slt b/tests/vchordrq/pin.slt new file mode 100644 index 00000000..fbe923c3 --- /dev/null +++ b/tests/vchordrq/pin.slt @@ -0,0 +1,18 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +build.pin = true +[build.internal] +lists = [32] +spherical_centroids = true +$$); + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/pushdown_plan.slt b/tests/vchordrq/pushdown_plan.slt new file mode 100644 index 00000000..62327229 --- /dev/null +++ b/tests/vchordrq/pushdown_plan.slt @@ -0,0 +1,206 @@ +# TODO: Some tests are disabled due to unimplemented types: sparse vector and f16 vector + +statement ok +CREATE TABLE t (val0 vector(3)); + +statement ok +INSERT INTO t (val0) +SELECT + ARRAY[random(), random(), random()]::real[]::vector +FROM generate_series(1, 10000); + +statement ok +CREATE INDEX ind0 ON t USING vchordrq (val0 vector_l2_ops); + +statement ok +SET enable_seqscan TO off; + +# statement ok +# CREATE INDEX ind1 ON t USING vchordrq (val1 halfvec_dot_ops); + +# 1 vector key + 1 corresponding order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + Order By: (val0 <-> '[0,0,0]'::vector) + +# 1 vector key + 0 order_by key + original style + +onlyif pg14 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg15 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg16 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg17 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +onlyif pg18 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1; +---- + Seq Scan on t + Disabled: true + Filter: ((val0 <-> '[0,0,0]'::vector) < '1'::double precision) + +# 1 vector key + 0 order_by key + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1); +---- + Index Scan using ind0 on t + Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + +# 0 vector key + 1 order_by key +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0,0,0]'::vector) + +# 2 vector key(1 of them is corresponding) + 1 order_by key + original style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <-> '[0, 0, 0]' < 1 +# AND val1 <#> '[0, 0, 0]' < 1 +# ORDER BY val0 <-> '[0, 0, 0]'; +# ---- +# Index Scan using ind0 on t +# Order By: (val0 <-> '[0,0,0]'::vector) +# Filter: (((val0 <-> '[0,0,0]'::vector) < '1'::double precision) AND ((val1 <#> '[0,0,0]'::halfvec) < '1'::double precision)) + +# 2 vector key(1 of them is corresponding) + 1 order_by key + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +# AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1) +# ORDER BY val0 <-> '[0, 0, 0]'; +# ---- +# Index Scan using ind0 on t +# Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) +# Order By: (val0 <-> '[0,0,0]'::vector) +# Filter: (val1 <<#>> '("[0,0,0]",1)'::sphere_vecf16) + +# 2 vector key(none of them is corresponding) + 1 order_by key + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val2 <<->> sphere('{}/3'::svector, 1) +# AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1) +# ORDER BY val0 <-> '[0, 0, 0]'; +# ---- +# Index Scan using ind0 on t +# Order By: (val0 <-> '[0,0,0]'::vector) +# Filter: ((val2 <<->> '({}/3,1)'::sphere_svector) AND (val1 <<#>> '("[0,0,0]",1)'::sphere_vecf16)) + +# 2 vector keys(both indexed) + 0 order_by key + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +# AND val1 <<#>> sphere('[0, 0, 0]'::halfvec, 1); +# ---- +# Index Scan using ind1 on t +# Index Cond: (val1 <<#>> '("[0,0,0]",1)'::sphere_vecf16) +# Filter: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + +# 2 vector keys(both not indexed) + 0 order_by key + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <<#>> sphere('[0, 0, 0]'::vector, 1) +# AND val1 <<->> sphere('[0, 0, 0]'::halfvec, 1); +# ---- +# Seq Scan on t +# Filter: ((val0 <<#>> '("[0,0,0]",1)'::sphere_vector) AND (val1 <<->> '("[0,0,0]",1)'::sphere_vecf16)) + +# 2 vector key(1 indexed) + 0 order_by key + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +# AND val1 <<->> sphere('[0, 0, 0]'::halfvec, 1); +# ---- +# Index Scan using ind0 on t +# Index Cond: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) +# Filter: (val1 <<->> '("[0,0,0]",1)'::sphere_vecf16) + +# 1 vector key + 1 not corresponding order_by key(operator) + sphere style +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) +SELECT val0 FROM t WHERE val0 <<#>> sphere('[0, 0, 0]'::vector, 1) +ORDER BY val0 <-> '[0, 0, 0]'; +---- + Index Scan using ind0 on t + Order By: (val0 <-> '[0,0,0]'::vector) + Filter: (val0 <<#>> '("[0,0,0]",1)'::sphere_vector) + +# 1 vector key + 1 not corresponding order_by key(variable) + sphere style +# query I +# EXPLAIN (COSTS FALSE, TIMING FALSE) +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0, 0, 0]'::vector, 1) +# ORDER BY val1 <#> '[1, 1, 1]'; +# ---- +# Index Scan using ind1 on t +# Order By: (val1 <#> '[1,1,1]'::halfvec) +# Filter: (val0 <<->> '("[0,0,0]",1)'::sphere_vector) + +# 0 vector key + 0 order_by key(variable) + +onlyif pg14 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg15 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg16 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg17 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + +onlyif pg18 +query I +EXPLAIN (COSTS FALSE, TIMING FALSE) SELECT val0 FROM t; +---- + Seq Scan on t + Disabled: true + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/pushdown_range.slt b/tests/vchordrq/pushdown_range.slt new file mode 100644 index 00000000..fb6431fd --- /dev/null +++ b/tests/vchordrq/pushdown_range.slt @@ -0,0 +1,47 @@ +# TODO: Some tests are disabled due to unimplemented types: sparse vector and f16 vector + +statement ok +CREATE TABLE t (val0 vector(3)); + +statement ok +INSERT INTO t (val0) VALUES + ('[0.1, 0.1, 0.1]'), + ('[0.2, 0.2, 0.2]'), + ('[0.3, 0.3, 0.3]'), + ('[0.4, 0.4, 0.4]'); + +statement ok +CREATE INDEX ON t USING vchordrq (val0 vector_l2_ops); + +# original style +query I +SELECT val0 FROM t WHERE val0 <-> '[0.24, 0.24, 0.24]' < 0.12 ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +---- +[0.2,0.2,0.2] +[0.3,0.3,0.3] + +# sphere style +query I +SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.12) ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +---- +[0.2,0.2,0.2] +[0.3,0.3,0.3] + +# sphere style: multiple vector keys and no order-by key +# query I +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) +# AND val1 <<#>> sphere('[0.24, -0.24, 0.24]'::halfvec, 0.05) +# ORDER BY val0 <-> '[0.24, 0.24, 0.24]'; +# ---- +# [0.2,0.2,0.2] + +# sphere style: vectors in key and order-by key are different +# query I +# SELECT val0 FROM t WHERE val0 <<->> sphere('[0.24, 0.24, 0.24]'::vector, 0.012) +# ORDER BY val1 <#> '[1, 1, -1]'; +# ---- +# [0.3,0.3,0.3] +# [0.2,0.2,0.2] + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/recall.slt b/tests/vchordrq/recall.slt new file mode 100644 index 00000000..8246f31a --- /dev/null +++ b/tests/vchordrq/recall.slt @@ -0,0 +1,188 @@ +statement ok +CREATE TABLE t (id SERIAL PRIMARY KEY, val vector(3)); + +statement ok +INSERT INTO t (val) +SELECT ARRAY[i * 0.0001, i * 0.00005, i * 0.0002]::vector(3) FROM generate_series(1, 10000) as s(i); + +statement ok +CREATE INDEX idx1 ON t USING vchordrq (val vector_l2_ops); + +statement ok +SET vchordrq.epsilon = 0.8; + +statement ok +SET vchordrq.probes = '1'; + +statement error MaxSim operator cannot be used for estimated recall +SELECT * from vchordrq_evaluate_query_recall(query=>'@#'); + +statement error Error executing ANN query (.+) need 0 probes, but 1 probes provided +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +statement ok +SET vchordrq.probes = ''; + +statement error Error executing ANN query (.+) could not convert type +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT val FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +statement error Error executing ANN query (.+) could not convert type +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT * FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); + +statement error Error executing Ground Truth query (.+) need 0 probes, but 1 probes provided +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$, accu_probes=>'1'); + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); +---- +1 + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$, exact_search=>true); +---- +1 + +query I +SELECT * from vchordrq_evaluate_query_recall(query=>$$SELECT ctid FROM t WHERE FALSE ORDER BY val <-> '[0.5, 0.25, 1.0]' LIMIT 10$$); +---- +NaN + +query I +SHOW vchordrq.epsilon; +---- +0.8 + +statement ok +CREATE TABLE t_dim4 (val vector(4), id SERIAL PRIMARY KEY); + +statement ok +INSERT INTO t_dim4 (val) +SELECT ARRAY[i * 0.0001, i * 0.00005, i * 0.0002, i * 0.001]::vector(4) FROM generate_series(1, 10000) as s(i); + +statement ok +CREATE INDEX idx2 ON t_dim4 USING vchordrq (val vector_l2_ops); + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_max_records = 1; + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_rate = 1; + +statement ok +ALTER SYSTEM SET vchordrq.query_sampling_enable = on; + +statement ok +SELECT pg_reload_conf(); + +query I retry 5 backoff 1s +SHOW vchordrq.query_sampling_enable; +---- +on + +statement ok +SET enable_seqscan = off; + +statement ok +SELECT * from t ORDER BY val <-> '[0.50, 0.25, 1.00]'; + +statement ok +SELECT * from t_dim4 ORDER BY val <-> '[1.00, 0.50, 0.25, 0]'; + +query I +SELECT value from vchordrq_sampled_queries('idx1'); +---- +[0.5,0.25,1] + +query I +SELECT value from vchordrq_sampled_queries('idx2'); +---- +[1,0.5,0.25,0] + +query I +SELECT COUNT(*) from vchordrq_sampled_queries; +---- +2 + +statement ok +SELECT * from t_dim4 ORDER BY val <-> '[2.1, 0.3, 0.7, 0.9]'; + +query I +SELECT * from vchordrq_sampled_queries('idx2'); +---- +public idx2 t_dim4 val <-> [2.1,0.3,0.7,0.9] + +query I +SELECT AVG(recall_value) +FROM ( + SELECT + vchordrq_evaluate_query_recall( + query => format( + 'SELECT ctid FROM %I.%I ORDER BY %I OPERATOR(%s) %L LIMIT 10', + lq.schema_name, + lq.table_name, + lq.column_name, + lq.operator, + lq.value + ) + ) AS recall_value + FROM + vchordrq_sampled_queries('idx2') AS lq +) AS eval_results; +---- +1 + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX idx3 ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops); + +statement ok +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.0000, 0.5000, 0.2500]' limit 1; + +query I +SELECT column_name from vchordrq_sampled_queries('idx3'); +---- +NULL + +query I +SELECT value from vchordrq_sampled_queries('idx3'); +---- +[1,0.5,0.25] + +statement ok +SET search_path='@'; + +query I +SELECT value from public.vchordrq_sampled_queries('public.idx3'); +---- +[1,0.5,0.25] + +query I +SELECT COUNT(*) from public.vchordrq_sampled_queries; +---- +3 + +statement ok +RESET enable_seqscan; + +statement ok +RESET search_path; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_enable; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_max_records; + +statement ok +ALTER SYSTEM RESET vchordrq.query_sampling_rate; + +statement ok +SELECT pg_reload_conf(); + +statement ok +DROP TABLE t, t_dim4, t_expr; diff --git a/tests/vchordrq/reindex.slt b/tests/vchordrq/reindex.slt new file mode 100644 index 00000000..9a58049b --- /dev/null +++ b/tests/vchordrq/reindex.slt @@ -0,0 +1,33 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 10000); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +REINDEX INDEX t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +REINDEX INDEX CONCURRENTLY t_val_idx; + +statement ok +INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]'); + +query I +SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10) t2; +---- +10 + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/rerank_in_index.slt b/tests/vchordrq/rerank_in_index.slt new file mode 100644 index 00000000..02d1e954 --- /dev/null +++ b/tests/vchordrq/rerank_in_index.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_column (id integer, val vector(3)); + +statement ok +INSERT INTO t_column (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_column USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_column ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +statement ok +DROP TABLE t_column; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = false +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' LIMIT 9; +---- +2 +1 +3 +4 +5 +7 +9 +11 +13 + +statement ok +SET vchordrq.prefilter to off; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +SET vchordrq.prefilter to on; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/rerank_in_table.slt b/tests/vchordrq/rerank_in_table.slt new file mode 100644 index 00000000..e9287976 --- /dev/null +++ b/tests/vchordrq/rerank_in_table.slt @@ -0,0 +1,101 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t_column (id integer, val vector(3)); + +statement ok +INSERT INTO t_column (id, val) SELECT id, ARRAY[id, id, id]::real[] FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_column USING vchordrq (val vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_column ORDER BY val <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +statement ok +DROP TABLE t_column; + +statement ok +CREATE TABLE t_expr (id integer); + +statement ok +INSERT INTO t_expr (id) SELECT id FROM generate_series(1, 10000) s(id); + +statement ok +CREATE INDEX ON t_expr USING vchordrq ((ARRAY[id::real, id::real, id::real]::vector(3)) vector_l2_ops) +WITH (options = $$ +residual_quantization = false +rerank_in_table = true +[build.internal] +lists = [] +$$); + +statement ok +SET vchordrq.probes = ''; + +query I +SELECT id FROM t_expr ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' limit 9; +---- +2 +1 +3 +4 +5 +6 +7 +8 +9 + +query I +SELECT id FROM t_expr WHERE id <= 5 OR id % 2 = 1 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> '[1.9, 1.9, 1.9]' LIMIT 9; +---- +2 +1 +3 +4 +5 +7 +9 +11 +13 + +statement ok +SET vchordrq.prefilter to off; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +SET vchordrq.prefilter to on; + +query I +SELECT ARRAY(SELECT id FROM t_expr WHERE (id <= 5 OR id % 2 = 1) OR e >= 2000 ORDER BY ARRAY[id::real, id::real, id::real]::vector(3) <-> q LIMIT 9) FROM (VALUES ('[1.9,1.99,1.999]'::vector, 1999), ('[2.1,2.11,2.111]', 2111)) AS t(q, e); +---- +{2,1,3,4,5,7,9,11,13} +{2,3,1,4,5,6,7,8,9} + +statement ok +DROP TABLE t_expr; diff --git a/tests/vchordrq/sequential_build.slt b/tests/vchordrq/sequential_build.slt new file mode 100644 index 00000000..b214b64f --- /dev/null +++ b/tests/vchordrq/sequential_build.slt @@ -0,0 +1,24 @@ +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SET max_parallel_workers = 0; + +statement ok +SET max_parallel_maintenance_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_ip_ops) +WITH (options = $$ +residual_quantization = false +build.pin = true +build.internal.lists = [32] +build.internal.build_threads = 2 +build.internal.spherical_centroids = true +$$); + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/tilemaxsim_source_registry.slt b/tests/vchordrq/tilemaxsim_source_registry.slt new file mode 100644 index 00000000..340bfd7d --- /dev/null +++ b/tests/vchordrq/tilemaxsim_source_registry.slt @@ -0,0 +1,91 @@ +# Caller-scoped TileMaxSim binds external tensor descriptors directly to an +# application relation. It must not require or inspect a VectorChord ANN index. + +statement ok +CREATE TABLE tilemaxsim_source_test ( + id integer PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2) +); + +statement ok +CREATE TABLE tilemaxsim_descriptor_test ( + source_id integer PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO tilemaxsim_source_test VALUES + (1, 'colqwen@test', '[1,0]'), + (2, 'colqwen@test', '[0,1]'); + +statement ok +INSERT INTO tilemaxsim_descriptor_test VALUES + (1, 'tensor://1', 2, 2, 'float32', 'sha256:test1'), + (2, 'tensor://2', 2, 2, 'float32', 'sha256:test2'); + +statement ok +SELECT vchordrq_register_tilemaxsim_source( + source_relation => 'tilemaxsim_source_test'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_RELATION', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'tilemaxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'source_id' +); + +query TTTT +SELECT registered_source::text, source_storage, + descriptor_relation::text, public_id_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tilemaxsim_source_test external_relation tilemaxsim_descriptor_test id + +# Candidate validation occurs before external tensor or sidecar access. +statement error candidate_ids must not contain duplicates +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::vector], + ARRAY[1,1]::bigint[], + 1 +); + +statement error top_k must be positive and no greater than candidate_ids length +SELECT * FROM vchordrq_tilemaxsim_rerank( + 'tilemaxsim_source_test'::regclass, + ARRAY['[1,0]'::halfvec], + ARRAY[1,2]::bigint[], + 3 +); + +# Attribute OIDs survive renames and the resolver returns the live name. +statement ok +ALTER TABLE tilemaxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_tilemaxsim_source_info('tilemaxsim_source_test'::regclass); +---- +tensor_location + +# Dropping a bound descriptor column removes the binding through sql_drop. +statement ok +ALTER TABLE tilemaxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) FROM _vchordrq_tilemaxsim_sources; +---- +0 + +statement ok +DROP TABLE tilemaxsim_source_test, tilemaxsim_descriptor_test; diff --git a/tests/vchordrq/vacuum.slt b/tests/vchordrq/vacuum.slt new file mode 100644 index 00000000..4cbdc08a --- /dev/null +++ b/tests/vchordrq/vacuum.slt @@ -0,0 +1,80 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +# statement ok +# SET vchordrq.max_parallel_vacuum_workers = 0; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.333; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.667; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +VACUUM t; + +statement ok +SELECT val FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 10; + +statement ok +DROP TABLE t; diff --git a/tests/vchordrq/vacuum_parallel.slt b/tests/vchordrq/vacuum_parallel.slt new file mode 100644 index 00000000..427ea608 --- /dev/null +++ b/tests/vchordrq/vacuum_parallel.slt @@ -0,0 +1,65 @@ +statement ok +SET enable_seqscan = off; + +statement ok +CREATE TABLE t (val vector(3)); + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +# statement ok +# SET vchordrq.max_parallel_vacuum_workers = 32; + +statement ok +SET min_parallel_index_scan_size = 1; + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [31] $$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [33] $$); + +statement ok +CREATE INDEX ON t USING vchordrq (val vector_l2_ops) +WITH (options = $$ build.internal.lists = [67] $$); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DELETE FROM t WHERE vector_norm(val) > 0.334; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000); + +statement ok +DELETE FROM t WHERE vector_norm(val) < 0.667; + +statement ok +VACUUM (PARALLEL 8) t; + +statement ok +DROP TABLE t; diff --git a/vchord.control b/vchord.control new file mode 100644 index 00000000..a08ebfe7 --- /dev/null +++ b/vchord.control @@ -0,0 +1,6 @@ +comment = 'vchord: Vector database plugin for Postgres, written in Rust, specifically designed for LLM' +default_version = '1.2.0' +module_pathname = 'vchord' +relocatable = true +superuser = true +requires = 'vector'