diff --git a/.gitfiles b/.gitfiles index a955184d73..e6e5f479a1 100644 --- a/.gitfiles +++ b/.gitfiles @@ -64,6 +64,7 @@ .github/helm/values/values-mirror-02.yaml .github/helm/values/values-profile.yaml .github/helm/values/values-readreplica.yaml +.github/helm/values/values-zvec.yaml .github/issue_label_bot.yaml .github/kubelinter.yaml .github/labeler.yaml @@ -460,6 +461,7 @@ charts/vald/templates/agent/sidecar/configmap.yaml charts/vald/templates/agent/sidecar/svc.yaml charts/vald/templates/agent/statefulset.yaml charts/vald/templates/agent/svc.yaml +charts/vald/templates/agent/zvec/configmap.yaml charts/vald/templates/discoverer/clusterrole.yaml charts/vald/templates/discoverer/clusterrolebinding.yaml charts/vald/templates/discoverer/configmap.yaml @@ -2295,6 +2297,8 @@ rust/libs/algorithms/qbg/build.rs rust/libs/algorithms/qbg/src/input.cpp rust/libs/algorithms/qbg/src/input.h rust/libs/algorithms/qbg/src/lib.rs +rust/libs/algorithms/zvec/Cargo.toml +rust/libs/algorithms/zvec/src/lib.rs rust/libs/build-utils/Cargo.toml rust/libs/build-utils/src/lib.rs rust/libs/kvs/Cargo.toml @@ -2380,6 +2384,7 @@ tests/v2/e2e/assets/readreplica.yaml tests/v2/e2e/assets/rollout.yaml tests/v2/e2e/assets/stream_crud.yaml tests/v2/e2e/assets/unary_crud.yaml +tests/v2/e2e/assets/zvec_crud.yaml tests/v2/e2e/config/config.go tests/v2/e2e/config/enums.go tests/v2/e2e/config/enums_test.go diff --git a/.github/actions/detect-docker-image-tags/action.yaml b/.github/actions/detect-docker-image-tags/action.yaml index 69ed823a10..c397f8de55 100644 --- a/.github/actions/detect-docker-image-tags/action.yaml +++ b/.github/actions/detect-docker-image-tags/action.yaml @@ -23,7 +23,7 @@ inputs: images: description: "Image names" required: false - default: "vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" + default: "vdaas/vald-agent vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" outputs: HELM_EXTRA_OPTIONS: description: "Helm extra options that specifies image tags" @@ -39,6 +39,7 @@ runs: id: specify_container_versions run: | declare -A m=( + ["vdaas/vald-agent"]="agent.image.tag" ["vdaas/vald-agent-ngt"]="agent.image.tag" ["vdaas/vald-agent-faiss"]="agent.image.tag" ["vdaas/vald-agent-sidecar"]="agent.sidecar.image.tag" diff --git a/.github/actions/setup-e2e/action.yaml b/.github/actions/setup-e2e/action.yaml index e801b0a0bd..414d07aeb2 100644 --- a/.github/actions/setup-e2e/action.yaml +++ b/.github/actions/setup-e2e/action.yaml @@ -51,7 +51,7 @@ inputs: target_images: description: "Image names" required: false - default: "vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" + default: "vdaas/vald-agent vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" outputs: HELM_EXTRA_OPTIONS: description: "Helm extra options that specifies E2E target image tags" diff --git a/.github/actions/wait-for-docker-image/action.yaml b/.github/actions/wait-for-docker-image/action.yaml index 294bebbc27..dd87eae458 100644 --- a/.github/actions/wait-for-docker-image/action.yaml +++ b/.github/actions/wait-for-docker-image/action.yaml @@ -19,7 +19,7 @@ inputs: images: description: "image names" required: false - default: "vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" + default: "vdaas/vald-agent vdaas/vald-agent-ngt vdaas/vald-agent-faiss vdaas/vald-agent-sidecar vdaas/vald-discoverer-k8s vdaas/vald-lb-gateway vdaas/vald-filter-gateway vdaas/vald-mirror-gateway vdaas/vald-manager-index vdaas/vald-index-operator vdaas/vald-readreplica-rotate" outputs: {} runs: using: "composite" diff --git a/.github/helm/values/values-zvec.yaml b/.github/helm/values/values-zvec.yaml new file mode 100644 index 0000000000..2c54169ac4 --- /dev/null +++ b/.github/helm/values/values-zvec.yaml @@ -0,0 +1,66 @@ +# +# Copyright (C) 2019-2026 vdaas.org vald team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +defaults: + logging: + level: info + networkPolicy: + enabled: true +gateway: + lb: + enabled: true + minReplicas: 1 + hpa: + enabled: false + resources: + requests: + cpu: 100m + memory: 50Mi + gateway_config: + index_replica: 3 +agent: + algorithm: zvec + minReplicas: 3 + maxReplicas: 3 + podManagementPolicy: Parallel + hpa: + enabled: false + image: + repository: vdaas/vald-agent + pullPolicy: IfNotPresent + resources: + requests: + cpu: 100m + memory: 50Mi + limits: + memory: 50Mi + zvec: + dimension: 784 + index_path: /var/lib/vald/zvec + fields: + - category + fts_fields: + - content +discoverer: + minReplicas: 1 + hpa: + enabled: false + resources: + requests: + cpu: 100m + memory: 50Mi +manager: + index: + enabled: false diff --git a/.github/workflows/dockers-agent-image.yaml b/.github/workflows/dockers-agent-image.yaml index 3a432b5219..01c4d0119c 100644 --- a/.github/workflows/dockers-agent-image.yaml +++ b/.github/workflows/dockers-agent-image.yaml @@ -56,6 +56,7 @@ on: - rust/Cargo.lock - rust/Cargo.toml - rust/bin/agent + - rust/libs/algorithms/zvec/** - rust/libs/ngt-rs/** - rust/libs/ngt/** - rust/libs/proto/** @@ -81,6 +82,7 @@ on: - rust/Cargo.lock - rust/Cargo.toml - rust/bin/agent + - rust/libs/algorithms/zvec/** - rust/libs/ngt-rs/** - rust/libs/ngt/** - rust/libs/proto/** diff --git a/.github/workflows/e2e.v2.yaml b/.github/workflows/e2e.v2.yaml index 3032205750..89e65dfa54 100644 --- a/.github/workflows/e2e.v2.yaml +++ b/.github/workflows/e2e.v2.yaml @@ -73,6 +73,7 @@ jobs: { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "management" }, { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "mirror" }, { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "profile" }, + { scenario: "zvec_crud", deployment: "helm-chart", cluster: "k3d", environment: "zvec", values: "zvec", insert_count: "1000", expected_index: "3000" }, { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "failure" }, { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "kill" }, { scenario: "stream_crud", deployment: "helm-chart", cluster: "k3d", environment: "partition" }, @@ -110,6 +111,15 @@ jobs: core.info('Matched ' + filtered.length + ' job(s).'); } + filtered = filtered.map(function(o) { + const values = o.values || ((o.environment === 'profile' || o.environment === 'faiss') ? o.environment : 'lb'); + return Object.assign({ + values: values, + insert_count: '10000', + expected_index: '30000', + }, o); + }); + // Set output "result" as a JSON string const matrix = { include: filtered }; core.info('Generated matrix:\n' + JSON.stringify(matrix, null, 2)); @@ -223,6 +233,12 @@ jobs: if: ${{ matrix.environment == 'management' }} run: | echo "HELM_EXTRA_OPTIONS=\"--values .github/helm/values/values-index-management-jobs.yaml\"" >> $GITHUB_ENV + - name: Build local Rust agent image for Zvec + if: ${{ matrix.environment == 'zvec' && matrix.cluster == 'k3d' }} + shell: bash + run: | + make TAG=nightly docker/build/agent + k3d image import vdaas/vald-agent:nightly -c vald - name: Deploy Vald by Helm Chart if: ${{ matrix.deployment == 'helm-chart' && matrix.environment != 'mirror' && matrix.scenario != 'readreplica' }} uses: ./.github/actions/e2e-deploy-vald @@ -232,8 +248,8 @@ jobs: # that override last so it wins over k8s/vald/manifests' own # `--set agent.image.repository=...-agent-ngt` (see Makefile.d/e2e.mk # e2e/v2/actions/run/faiss for the equivalent local dev target). - helm_extra_options: "${{ steps.setup_e2e.outputs.HELM_EXTRA_OPTIONS }}${{ matrix.environment == 'faiss' && ' --set agent.image.repository=vdaas/vald-agent-faiss' || '' }}" - values: .github/helm/values/values-${{ (matrix.environment == 'profile' || matrix.environment == 'faiss') && matrix.environment || 'lb' }}.yaml + helm_extra_options: "${{ steps.setup_e2e.outputs.HELM_EXTRA_OPTIONS }}${{ matrix.environment == 'faiss' && ' --set agent.image.repository=vdaas/vald-agent-faiss' || matrix.environment == 'zvec' && ' --set agent.image.repository=vdaas/vald-agent' || '' }}" + values: .github/helm/values/values-${{ matrix.values }}.yaml wait_for_selector: "app=vald-lb-gateway" - name: Deploy Vald Read Replica if: ${{ 'readreplica' == matrix.scenario }} @@ -327,8 +343,8 @@ jobs: kubectl get pods make e2e/v2 E2E_CONFIG=$(pwd)/tests/v2/e2e/assets/${{ matrix.scenario }}.yaml \ E2E_PARALLELISM="4" \ - E2E_INSERT_COUNT="10000" \ - E2E_EXPECTED_INDEX="30000" \ + E2E_INSERT_COUNT="${{ matrix.insert_count }}" \ + E2E_EXPECTED_INDEX="${{ matrix.expected_index }}" \ E2E_QPS="30" \ E2E_SEARCH_COUNT="10" \ E2E_UPDATE_COUNT="100" \ diff --git a/Makefile.d/build.mk b/Makefile.d/build.mk index d3116d94ef..0a97281c36 100644 --- a/Makefile.d/build.mk +++ b/Makefile.d/build.mk @@ -132,12 +132,14 @@ example/client/client: $(eval CGO_ENABLED = 1) $(call go-example-build,example/client,-linkmode 'external',$(HDF5_LDFLAGS), cgo,$(HDF5_VERSION),$@) +RUST_AGENT_FEATURES ?= qbg,zvec + rust/target/release/agent: pushd rust && \ $(CC_ENV_VARS) \ OPENSSL_STATIC=1 \ PKG_CONFIG_ALL_STATIC=1 \ - cargo build -p agent --release && \ + cargo build -p agent --no-default-features --features "$(RUST_AGENT_FEATURES)" --release && \ popd rust/target/debug/agent: @@ -145,7 +147,7 @@ rust/target/debug/agent: $(CC_ENV_VARS) \ OPENSSL_STATIC=1 \ PKG_CONFIG_ALL_STATIC=1 \ - cargo build -p agent && \ + cargo build -p agent --no-default-features --features "$(RUST_AGENT_FEATURES)" && \ popd tests/v2/e2e/e2e: diff --git a/Makefile.d/test.mk b/Makefile.d/test.mk index 674bf68c33..54e5b6293f 100644 --- a/Makefile.d/test.mk +++ b/Makefile.d/test.mk @@ -203,6 +203,7 @@ test/cmd: \ ## run tests for rust test/rust: \ test/rust/qbg \ + test/rust/zvec \ test/rust/agent .PHONY: test/rust/qbg @@ -219,11 +220,17 @@ test/rust/qbg: cargo test --manifest-path rust/Cargo.toml --package qbg --lib -- tests::test_index --exact --show-output rm -rf rust/libs/algorithms/qbg/index/ +.PHONY: test/rust/zvec +## run tests for zvec +test/rust/zvec: + $(CC_ENV_VARS) \ + cargo test --manifest-path rust/Cargo.toml --package zvec --lib -- --show-output + .PHONY: test/rust/agent ## run tests for agent test/rust/agent: $(CC_ENV_VARS) \ - cargo test --manifest-path rust/Cargo.toml --package agent -- handler::common::tests --show-output + cargo test --manifest-path rust/Cargo.toml --package agent --no-default-features --features "$(RUST_AGENT_FEATURES)" -- handler::common::tests --show-output .PHONY: test/hack ## run tests for hack diff --git a/charts/vald/templates/agent/zvec/configmap.yaml b/charts/vald/templates/agent/zvec/configmap.yaml new file mode 100644 index 0000000000..01d6161646 --- /dev/null +++ b/charts/vald/templates/agent/zvec/configmap.yaml @@ -0,0 +1,46 @@ +# +# Copyright (C) 2019-2026 vdaas.org vald team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +{{- $agent := .Values.agent -}} +{{- if and ($agent.enabled) (eq (lower $agent.algorithm) "zvec")}} +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ $agent.name }}-config + labels: + app.kubernetes.io/name: {{ include "vald.name" . }} + helm.sh/chart: {{ include "vald.chart" . }} + app.kubernetes.io/managed-by: {{ .Release.Service }} + app.kubernetes.io/instance: {{ .Release.Name }} + app.kubernetes.io/version: {{ .Chart.Version }} + app.kubernetes.io/component: agent +data: + config.yaml: | + --- + algorithm: zvec + version: {{ $agent.version }} + time_zone: {{ default .Values.defaults.time_zone $agent.time_zone }} + logging: + {{- $logging := dict "Values" $agent.logging "default" .Values.defaults.logging }} + {{- include "vald.logging" $logging | nindent 6 }} + server_config: + {{- $servers := dict "Values" $agent.server_config "default" .Values.defaults.server_config }} + {{- include "vald.servers" $servers | nindent 6 }} + observability: + {{- $observability := dict "Values" $agent.observability "default" .Values.defaults.observability }} + {{- include "vald.observability" $observability | nindent 6 }} + zvec: + {{- toYaml $agent.zvec | nindent 6 }} +{{- end }} diff --git a/charts/vald/values.schema.json b/charts/vald/values.schema.json index 4684bcc129..62e2e7f6ca 100644 --- a/charts/vald/values.schema.json +++ b/charts/vald/values.schema.json @@ -80,8 +80,8 @@ }, "algorithm": { "type": "string", - "description": "agent algorithm type. it should be `ngt` or `faiss`.", - "enum": ["ngt", "faiss"] + "description": "agent algorithm type. it should be `ngt`, `faiss` or `zvec`.", + "enum": ["ngt", "faiss", "zvec"] }, "annotations": { "type": "object", @@ -640,6 +640,34 @@ "name": "config" } }, + "zvec": { + "type": "object", + "properties": { + "dimension": { + "type": "integer", + "description": "vector dimension", + "minimum": 1 + }, + "index_path": { + "type": "string", + "description": "Zvec collection storage path" + }, + "fields": { + "type": "array", + "description": "string fields stored in the Zvec collection", + "items": { + "type": "string" + } + }, + "fts_fields": { + "type": "array", + "description": "full-text indexed string fields stored in the Zvec collection", + "items": { + "type": "string" + } + } + } + }, "persistentVolume": { "type": "object", "properties": { diff --git a/charts/vald/values.yaml b/charts/vald/values.yaml index 122791f0ed..7874a27172 100644 --- a/charts/vald/values.yaml +++ b/charts/vald/values.yaml @@ -1998,9 +1998,9 @@ agent: # @schema {"name": "agent.version", "alias": "version"} # agent.version -- version of agent config version: v0.0.0 - # @schema {"name": "agent.algorithm", "type": "string", "enum": ["ngt", "faiss"]} + # @schema {"name": "agent.algorithm", "type": "string", "enum": ["ngt", "faiss", "zvec"]} # agent.algorithm -- agent algorithm type. - # it should be `ngt` or `faiss`. + # it should be `ngt`, `faiss` or `zvec`. algorithm: ngt # @schema {"name": "agent.time_zone", "type": "string"} # agent.time_zone -- Time zone @@ -2468,6 +2468,22 @@ agent: # @schema {"name": "agent.faiss.kvsdb.concurrency", "type": "integer"} # agent.faiss.kvsdb.concurrency -- kvsdb processing concurrency concurrency: 6 + # @schema {"name": "agent.zvec", "type": "object"} + zvec: + # @schema {"name": "agent.zvec.dimension", "type": "integer", "minimum": 1} + # agent.zvec.dimension -- vector dimension + dimension: 4096 + # @schema {"name": "agent.zvec.index_path", "type": "string"} + # agent.zvec.index_path -- Zvec collection storage path + index_path: /var/lib/vald/zvec + # @schema {"name": "agent.zvec.fields", "type": "array", "items": {"type": "string"}} + # agent.zvec.fields -- string fields stored in the Zvec collection + fields: + - category + # @schema {"name": "agent.zvec.fts_fields", "type": "array", "items": {"type": "string"}} + # agent.zvec.fts_fields -- full-text indexed string fields stored in the Zvec collection + fts_fields: + - content # @schema {"name": "agent.sidecar", "type": "object"} sidecar: # @schema {"name": "agent.sidecar.enabled", "type": "boolean"} diff --git a/dockers/agent/core/agent/Dockerfile b/dockers/agent/core/agent/Dockerfile index eb947210b8..c1966e0b16 100644 --- a/dockers/agent/core/agent/Dockerfile +++ b/dockers/agent/core/agent/Dockerfile @@ -94,12 +94,14 @@ RUN --mount=type=bind,target=.,rw \ && CC=clang CXX=clang++ make ngt/install \ && make faiss/install \ && make rust/target/release/${APP_NAME} \ + && find rust/target/release/build -name libzvec_c_api.so -exec cp {} /usr/lib/libzvec_c_api.so \; \ && mv "rust/target/release/${APP_NAME}" "/usr/bin/${APP_NAME}" \ && rm -rf rust/target # skipcq: DOK-DL3026,DOK-DL3007 FROM gcr.io/distroless/cc-debian12:nonroot LABEL maintainer="vdaas.org vald team " COPY --from=builder /usr/bin/agent /usr/bin/agent +COPY --from=builder /usr/lib/libzvec_c_api.so /usr/lib/libzvec_c_api.so # skipcq: DOK-DL3002 USER nonroot:nonroot ENTRYPOINT ["/usr/bin/agent"] diff --git a/hack/docker/gen/main.go b/hack/docker/gen/main.go index 7f0c364c2c..a98ed62523 100644 --- a/hack/docker/gen/main.go +++ b/hack/docker/gen/main.go @@ -42,6 +42,7 @@ const ( agent = "agent" agentFaiss = agent + "-faiss" agentNGT = agent + "-ngt" + agentZVec = agent + "-zvec" agentSidecar = agent + "-sidecar" bench = "benchmark" benchJob = bench + "-job" @@ -130,6 +131,7 @@ const ( rustNgtRsPath = "rust/libs/ngt-rs/**" rustNgtPath = "rust/libs/ngt/**" rustProtoPath = "rust/libs/proto/**" + rustZVecPath = "rust/libs/algorithms/zvec/**" excludeTestFilesPath = "!**/*_test.go" excludeMockFilesPath = "!**/*_mock.go" @@ -456,6 +458,7 @@ var ( } rustBuildCommands = []string{ "make rust/target/release/${APP_NAME}", + "find rust/target/release/build -name libzvec_c_api.so -exec cp {} /usr/lib/libzvec_c_api.so \\;", "mv \"rust/target/release/${APP_NAME}\" \"{{$.BinDir}}/${APP_NAME}\"", "rm -rf rust/target", } @@ -718,6 +721,7 @@ func main() { ngtClangLTOPreprocess, faissPreprocess, }, + StageFiles: []string{"/usr/lib/libzvec_c_api.so"}, }, vald + "-" + agentSidecar: { AppName: "sidecar", @@ -970,6 +974,7 @@ func main() { rustNgtRsPath, rustNgtPath, rustProtoPath, + rustZVecPath, rustVersionPath, ) } diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1123032d67..af03bd5adc 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -42,6 +42,7 @@ dependencies = [ "tonic", "tonic-types", "tower", + "zvec", ] [[package]] @@ -58,10 +59,7 @@ name = "algorithm" version = "0.1.0" dependencies = [ "anyhow", - "faiss", - "ngt", "proto", - "qbg", "tonic", ] @@ -3113,3 +3111,28 @@ name = "zmij" version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zvec" +version = "0.1.0" +dependencies = [ + "algorithm", + "prost", + "proto", + "zvec-rust", +] + +[[package]] +name = "zvec-rust" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68005391d2debdee9dcdd295830f7ca4f6bfba2f4f2e3cadfb8f1cc6c7929987" +dependencies = [ + "zvec-rust-sys", +] + +[[package]] +name = "zvec-rust-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f89deb26358085908592b89a97af9155e8f91b88d78006289589a394931b7664" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 63b1936661..f48dbc110e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -23,6 +23,7 @@ members = [ "libs/algorithms/ngt", "libs/algorithms/qbg", "libs/algorithms/faiss", + "libs/algorithms/zvec", "libs/kvs", "libs/observability", "libs/proto", diff --git a/rust/bin/agent/Cargo.toml b/rust/bin/agent/Cargo.toml index f397af1c48..3ff5dee7fa 100644 --- a/rust/bin/agent/Cargo.toml +++ b/rust/bin/agent/Cargo.toml @@ -22,7 +22,8 @@ edition = "2024" [dependencies] algorithm = { version = "0.1.0", path = "../../libs/algorithm" } -qbg = { version = "0.1.0", path = "../../libs/algorithms/qbg" } +qbg = { version = "0.1.0", path = "../../libs/algorithms/qbg", optional = true } +zvec = { version = "0.1.0", path = "../../libs/algorithms/zvec", optional = true } anyhow = "1.0.104" chrono = "0.4.45" config = "0.15.25" @@ -46,3 +47,8 @@ cc = "1" [dev-dependencies] bytes = "1.12.1" http-body = "1.1.0" + +[features] +default = ["qbg", "zvec"] +qbg = ["dep:qbg"] +zvec = ["dep:zvec"] diff --git a/rust/bin/agent/build.rs b/rust/bin/agent/build.rs index 0685e2ab32..0e84bc9a7a 100644 --- a/rust/bin/agent/build.rs +++ b/rust/bin/agent/build.rs @@ -19,7 +19,9 @@ fn main() { .file("src/glibc_compat.c") .compile("glibc_compat"); - println!("cargo:rustc-link-lib=static=stdc++"); + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux") { + println!("cargo:rustc-link-lib=static=stdc++"); + } println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=src/glibc_compat.c"); } diff --git a/rust/bin/agent/src/handler.rs b/rust/bin/agent/src/handler.rs index f25d19d945..80eaaec86c 100644 --- a/rust/bin/agent/src/handler.rs +++ b/rust/bin/agent/src/handler.rs @@ -23,7 +23,7 @@ use std::sync::Arc; use tokio::sync::RwLock; pub struct Agent { - s: Arc>, + s: Arc>>, name: String, ip: String, resource_type: String, @@ -33,7 +33,7 @@ pub struct Agent { impl Agent { pub fn new( - s: impl algorithm::ANN + 'static, + s: Box, name: &str, ip: &str, resource_type: &str, diff --git a/rust/bin/agent/src/handler/insert.rs b/rust/bin/agent/src/handler/insert.rs index 98586d3578..888ccfb5d9 100644 --- a/rust/bin/agent/src/handler/insert.rs +++ b/rust/bin/agent/src/handler/insert.rs @@ -18,7 +18,7 @@ use proto::{ payload::v1::{insert, object}, vald::v1::insert_server, }; -use std::{collections::HashMap, string::String, sync::Arc}; +use std::{string::String, sync::Arc}; use tokio::sync::RwLock; use tonic::{Code, Status}; use tonic_types::StatusExt; @@ -26,7 +26,7 @@ use tonic_types::StatusExt; use super::common::{bidirectional_stream, build_error_details}; pub(super) async fn insert( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, @@ -69,7 +69,12 @@ pub(super) async fn insert( warn!("{:?}", status); return Err(status); } - let result = s.insert(vec.id.clone(), vec.vector.clone(), config.timestamp); + let result = s.insert_with_options( + vec.id.clone(), + vec.vector.clone(), + config.timestamp, + &config.options, + ); match result { Err(err) => { let resource_type = format!("{}/qbg.Insert", resource_type); @@ -227,7 +232,7 @@ impl insert_server::Insert for super::Agent { let hostname = super::common::get_hostname(); let domain = hostname.as_str(); let mut uuids: Vec = Vec::new(); - let mut vmap = HashMap::new(); + let mut vectors = Vec::with_capacity(mreq.requests.len()); { let mut s = self.s.write().await; for req in mreq.requests.clone() { @@ -259,10 +264,28 @@ impl insert_server::Insert for super::Agent { warn!("{:?}", status); return Err(status); } + let config = req + .config + .ok_or_else(|| Status::invalid_argument("Missing configuration in request"))?; uuids.push(vec.id.clone()); - vmap.insert(vec.id, vec.vector); + vectors.push((vec.id, vec.vector, config.timestamp, config.options)); + } + let mut duplicated = Vec::new(); + let mut result = Ok(()); + for (uuid, vector, timestamp, options) in vectors { + if let Err(err) = s.insert_with_options(uuid.clone(), vector, timestamp, &options) { + match err { + Error::UUIDAlreadyExists { uuid } => duplicated.push(uuid), + _ => { + result = Err(err); + break; + } + } + } + } + if result.is_ok() && !duplicated.is_empty() { + result = Err(Error::new_uuid_already_exists(duplicated)); } - let result = s.insert_multiple(vmap); match result { Err(err) => { let resource_type = format!("{}/qbg.MultiInsert", self.resource_type); diff --git a/rust/bin/agent/src/handler/object.rs b/rust/bin/agent/src/handler/object.rs index 5b1dc4b758..6acf64fcf1 100644 --- a/rust/bin/agent/src/handler/object.rs +++ b/rust/bin/agent/src/handler/object.rs @@ -23,7 +23,7 @@ use tonic_types::StatusExt; use super::common::{bidirectional_stream, build_error_details}; async fn get_object( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, diff --git a/rust/bin/agent/src/handler/remove.rs b/rust/bin/agent/src/handler/remove.rs index 15a6907338..138017b5cc 100644 --- a/rust/bin/agent/src/handler/remove.rs +++ b/rust/bin/agent/src/handler/remove.rs @@ -26,7 +26,7 @@ use tonic_types::StatusExt; use super::common::{bidirectional_stream, build_error_details}; async fn remove( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, diff --git a/rust/bin/agent/src/handler/search.rs b/rust/bin/agent/src/handler/search.rs index 23453717b9..b28a0bc350 100644 --- a/rust/bin/agent/src/handler/search.rs +++ b/rust/bin/agent/src/handler/search.rs @@ -23,7 +23,7 @@ use tonic_types::StatusExt; use super::common::{bidirectional_stream, build_error_details}; async fn search( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, @@ -62,11 +62,12 @@ async fn search( warn!("{:?}", status); return Err(status); } - let result = s.search( + let result = s.search_with_options( request.vector.clone(), config.num, config.epsilon, config.radius, + &config.options, ); match result { Err(err) => { diff --git a/rust/bin/agent/src/handler/update.rs b/rust/bin/agent/src/handler/update.rs index b39e8fdb6b..bd9a15788f 100644 --- a/rust/bin/agent/src/handler/update.rs +++ b/rust/bin/agent/src/handler/update.rs @@ -26,7 +26,7 @@ use tonic_types::StatusExt; use super::common::{bidirectional_stream, build_error_details}; pub(crate) async fn update( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, diff --git a/rust/bin/agent/src/handler/upsert.rs b/rust/bin/agent/src/handler/upsert.rs index cfb7495b1b..7cfebc2028 100644 --- a/rust/bin/agent/src/handler/upsert.rs +++ b/rust/bin/agent/src/handler/upsert.rs @@ -28,7 +28,7 @@ use super::insert::insert as insert_fn; use super::update::update as update_fn; async fn upsert( - s: Arc>, + s: Arc>>, resource_type: &str, api_name: &str, name: &str, diff --git a/rust/bin/agent/src/main.rs b/rust/bin/agent/src/main.rs index d60830ad85..408b6e9824 100644 --- a/rust/bin/agent/src/main.rs +++ b/rust/bin/agent/src/main.rs @@ -12,16 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -use algorithm::{Error, MultiError}; +use algorithm::Error; +#[cfg(feature = "qbg")] +use algorithm::MultiError; use anyhow::Result; +#[cfg(feature = "qbg")] use chrono::{Local, Timelike}; use config::Config; +#[cfg(feature = "qbg")] use proto::payload::v1::object::Distance; use proto::payload::v1::search; -use qbg::index::Index; -use qbg::property::Property; +#[cfg(feature = "qbg")] +use qbg::{index::Index, property::Property}; use std::collections::HashMap; use std::time::Duration; +#[cfg(feature = "zvec")] +use zvec::Service as ZVecService; mod handler; mod middleware; @@ -120,12 +126,14 @@ impl algorithm::ANN for _MockService { } } +#[cfg(feature = "qbg")] struct QBGService { path: String, index: Index, property: Property, } +#[cfg(feature = "qbg")] impl QBGService { fn new(settings: Config) -> Self { let path = settings @@ -192,6 +200,7 @@ impl QBGService { } } +#[cfg(feature = "qbg")] impl algorithm::ANN for QBGService { fn exists(&self, _uuid: String) -> bool { // convert uuid to id @@ -401,10 +410,47 @@ async fn main() -> Result<(), Box> { .unwrap(); let _logger = flexi_logger::Logger::try_with_str(settings.get::("logging.level")?)?.start()?; - let service = QBGService::new(settings.clone()); + let algorithm = settings + .get::("algorithm") + .unwrap_or_else(|_| "qbg".to_string()) + .to_lowercase(); + let (service, agent_name): (Box, &str) = match algorithm.as_str() { + #[cfg(feature = "zvec")] + "zvec" => { + let path = settings + .get::("zvec.index_path") + .unwrap_or_else(|_| "/var/lib/vald/zvec".to_string()); + let dimension = settings.get::("zvec.dimension").unwrap_or(784); + let fields = settings + .get_array("zvec.fields") + .unwrap_or_default() + .into_iter() + .filter_map(|value| value.into_string().ok()) + .collect(); + let fts_fields = settings + .get_array("zvec.fts_fields") + .unwrap_or_default() + .into_iter() + .filter_map(|value| value.into_string().ok()) + .collect(); + ( + Box::new(ZVecService::new(&path, dimension, fields, fts_fields)?), + "agent-zvec", + ) + } + #[cfg(feature = "qbg")] + "qbg" => (Box::new(QBGService::new(settings.clone())), "agent-qbg"), + #[cfg(not(feature = "zvec"))] + "zvec" => return Err("Rust agent was built without the zvec feature".into()), + #[cfg(not(feature = "qbg"))] + "qbg" => return Err("Rust agent was built without the qbg feature".into()), + name => { + return Err(format!("unsupported Rust agent algorithm: {name}").into()); + } + }; let agent = handler::Agent::new( service, - "agent-qbg", + agent_name, "127.0.0.1", "vald/internal/core/algorithm", "vald-agent", diff --git a/rust/libs/algorithm/Cargo.toml b/rust/libs/algorithm/Cargo.toml index ebad3ceeaf..c5bf5acc82 100644 --- a/rust/libs/algorithm/Cargo.toml +++ b/rust/libs/algorithm/Cargo.toml @@ -20,8 +20,5 @@ edition = "2024" [dependencies] anyhow = "1.0.104" -faiss = { version = "0.1.0", path = "../algorithms/faiss" } -ngt = { version = "0.1.0", path = "../algorithms/ngt" } -qbg = { version = "0.1.0", path = "../algorithms/qbg" } proto = { version = "0.1.0", path = "../proto" } tonic = "0.14.6" diff --git a/rust/libs/algorithm/src/lib.rs b/rust/libs/algorithm/src/lib.rs index 3cc9d84e8c..45bcb1cd75 100644 --- a/rust/libs/algorithm/src/lib.rs +++ b/rust/libs/algorithm/src/lib.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. use anyhow::Result; +use proto::google::protobuf::Any; use proto::payload::v1::search; use std::{collections::HashMap, error, fmt, i64}; @@ -55,6 +56,9 @@ pub enum Error { uuid: String, }, Unknown {}, + Backend { + message: String, + }, } impl MultiError for Error { @@ -137,6 +141,7 @@ impl fmt::Display for Error { } Error::ObjectIDNotFound { uuid } => write!(f, "uuid {}'s object id not found", uuid), Error::Unknown {} => write!(f, "unknown error"), + Error::Backend { message } => write!(f, "algorithm backend error: {}", message), } } } @@ -146,6 +151,15 @@ pub trait ANN: Send + Sync { fn create_index(&mut self) -> Result<(), Error>; fn save_index(&mut self) -> Result<(), Error>; fn insert(&mut self, uuid: String, vector: Vec, ts: i64) -> Result<(), Error>; + fn insert_with_options( + &mut self, + uuid: String, + vector: Vec, + ts: i64, + _options: &[Any], + ) -> Result<(), Error> { + self.insert(uuid, vector, ts) + } fn insert_multiple(&mut self, vectors: HashMap>) -> Result<(), Error>; fn update(&mut self, uuid: String, vector: Vec, ts: i64) -> Result<(), Error>; fn update_multiple(&mut self, vectors: HashMap>) -> Result<(), Error>; @@ -159,6 +173,16 @@ pub trait ANN: Send + Sync { epsilon: f32, radius: f32, ) -> Result; + fn search_with_options( + &self, + vector: Vec, + k: u32, + epsilon: f32, + radius: f32, + _options: &[Any], + ) -> Result { + self.search(vector, k, epsilon, radius) + } fn get_object(&self, uuid: String) -> Result<(Vec, i64), Error>; fn get_dimension_size(&self) -> usize; fn len(&self) -> u32; diff --git a/rust/libs/algorithms/zvec/Cargo.toml b/rust/libs/algorithms/zvec/Cargo.toml new file mode 100644 index 0000000000..d52316e066 --- /dev/null +++ b/rust/libs/algorithms/zvec/Cargo.toml @@ -0,0 +1,25 @@ +# +# Copyright (C) 2019-2026 vdaas.org vald team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +[package] +name = "zvec" +version = "0.1.0" +edition = "2024" + +[dependencies] +algorithm = { version = "0.1.0", path = "../../algorithm" } +prost = "0.14.4" +proto = { version = "0.1.0", path = "../../proto" } +zvec-rust = "0.6.0" diff --git a/rust/libs/algorithms/zvec/src/lib.rs b/rust/libs/algorithms/zvec/src/lib.rs new file mode 100644 index 0000000000..b7fc6d30fe --- /dev/null +++ b/rust/libs/algorithms/zvec/src/lib.rs @@ -0,0 +1,656 @@ +// Copyright (C) 2019-2026 vdaas.org vald team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// You may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use algorithm::{ANN, Error, MultiError}; +use prost::Message; +use proto::google::protobuf::Any; +use proto::payload::v1::{object::Distance, search, z_vec}; +use std::collections::HashMap; +use std::path::Path; +use std::sync::OnceLock; +use zvec_rust::{ + Collection, CollectionSchema, DataType, Doc, ErrorCode, Fts, IndexParams, MetricType, + MultiQuery, SearchQuery, SubQuery, +}; + +const VECTOR_FIELD: &str = "embedding"; +const TIMESTAMP_FIELD: &str = "timestamp"; +static INITIALIZED: OnceLock> = OnceLock::new(); + +pub struct Service { + collection: Collection, + dimension: usize, + fields: Vec, +} + +impl Service { + pub fn new( + path: &str, + dimension: usize, + fields: Vec, + fts_fields: Vec, + ) -> Result { + if dimension < 2 { + return Err(Error::InvalidDimensionSize { + uuid: String::new(), + current: dimension.to_string(), + limit: "0".to_string(), + }); + } + match INITIALIZED.get_or_init(|| zvec_rust::initialize(None).map_err(|err| err.to_string())) + { + Ok(()) => {} + Err(message) => { + return Err(Error::Backend { + message: message.clone(), + }); + } + } + + let collection = if Path::new(path).exists() { + Collection::open(path, None).map_err(backend_error)? + } else { + let mut schema = CollectionSchema::builder("vald_zvec") + .add_field( + zvec_rust::FieldSchema::new(TIMESTAMP_FIELD, DataType::Int64, false, 0) + .map_err(backend_error)?, + ) + .add_vector_field( + VECTOR_FIELD, + DataType::VectorFp32, + dimension as u32, + IndexParams::hnsw(MetricType::L2, 16, 200).map_err(backend_error)?, + ); + for field in &fields { + let mut field_schema = + zvec_rust::FieldSchema::new(field, DataType::String, true, 0) + .map_err(backend_error)?; + field_schema + .set_index_params(&IndexParams::invert(false, false).map_err(backend_error)?) + .map_err(backend_error)?; + schema = schema.add_field(field_schema); + } + for field in &fts_fields { + let mut field_schema = + zvec_rust::FieldSchema::new(field, DataType::String, true, 0) + .map_err(backend_error)?; + field_schema + .set_index_params(&IndexParams::fts(None, None, None).map_err(backend_error)?) + .map_err(backend_error)?; + schema = schema.add_field(field_schema); + } + let schema = schema.build().map_err(backend_error)?; + Collection::create_and_open(path, &schema, None).map_err(backend_error)? + }; + Ok(Self { + collection, + dimension, + fields: fields.into_iter().chain(fts_fields).collect(), + }) + } + + fn document( + &self, + uuid: &str, + vector: &[f32], + timestamp: i64, + options: &[Any], + ) -> Result { + if uuid.is_empty() { + return Err(Error::InvalidUUID { + uuid: uuid.to_string(), + }); + } + self.validate_dimension(uuid, vector)?; + let options = decode_any::(options)?; + let mut doc = Doc::new().map_err(backend_error)?; + doc.set_pk(uuid); + doc.add_i64(TIMESTAMP_FIELD, timestamp) + .map_err(backend_error)?; + doc.add_vector_f32(VECTOR_FIELD, vector) + .map_err(backend_error)?; + if let Some(options) = options { + for (name, value) in options.fields { + if self.fields.iter().any(|field| field == &name) { + doc.add_string(&name, &value).map_err(backend_error)?; + } + } + } + Ok(doc) + } + + fn validate_dimension(&self, uuid: &str, vector: &[f32]) -> Result<(), Error> { + if vector.len() != self.dimension { + return Err(Error::InvalidDimensionSize { + uuid: uuid.to_string(), + current: vector.len().to_string(), + limit: self.dimension.to_string(), + }); + } + Ok(()) + } + + fn write_result(result: zvec_rust::WriteResult, uuids: &[&str]) -> Result<(), Error> { + if result.error_count == 0 { + return Ok(()); + } + let (index, failed) = result + .results + .into_iter() + .enumerate() + .find(|(_, result)| !result.is_success()) + .ok_or_else(|| Error::Backend { + message: "zvec write failed".to_string(), + })?; + let uuid = uuids.get(index).copied().unwrap_or_default().to_string(); + match failed.code { + ErrorCode::AlreadyExists => Err(Error::UUIDAlreadyExists { uuid }), + ErrorCode::NotFound => Err(Error::ObjectIDNotFound { uuid }), + _ => Err(Error::Backend { + message: failed.message, + }), + } + } + + fn search_documents( + &self, + vector: &[f32], + k: u32, + options: Option, + ) -> Result, Error> { + let Some(options) = options else { + let query = SearchQuery::new(VECTOR_FIELD, vector, k as i32).map_err(backend_error)?; + return self.collection.query(&query).map_err(backend_error); + }; + if options.hybrid_queries.is_empty() { + let mut builder = SearchQuery::builder() + .field_name(VECTOR_FIELD) + .vector(vector) + .topk(k as i32); + if !options.pre_filter.is_empty() { + builder = builder.filter(&options.pre_filter); + } + let query = builder.build().map_err(backend_error)?; + return self.collection.query(&query).map_err(backend_error); + } + if options.hybrid_queries.len() == 1 { + let item = &options.hybrid_queries[0]; + let mut query = match &item.query { + Some(z_vec::query::Query::Vector(vector)) => { + SearchQuery::new(&item.field_name, &vector.values, k as i32) + .map_err(backend_error)? + } + Some(z_vec::query::Query::Fts(fts)) => { + let mut value = Fts::new().map_err(backend_error)?; + if !fts.match_string.is_empty() { + value + .set_match_string(&fts.match_string) + .map_err(backend_error)?; + } + if !fts.query_string.is_empty() { + value + .set_query_string(&fts.query_string) + .map_err(backend_error)?; + } + SearchQuery::fts(&item.field_name, &value, k as i32).map_err(backend_error)? + } + None => { + return Err(Error::Backend { + message: "Zvec hybrid query has no payload".to_string(), + }); + } + }; + if !options.pre_filter.is_empty() { + query + .set_filter(&options.pre_filter) + .map_err(backend_error)?; + } + return self.collection.query(&query).map_err(backend_error); + } + + let mut query = MultiQuery::new().map_err(backend_error)?; + query.set_topk(k as i32).map_err(backend_error)?; + if !options.pre_filter.is_empty() { + query + .set_filter(&options.pre_filter) + .map_err(backend_error)?; + } + for item in options.hybrid_queries { + let mut sub = SubQuery::new().map_err(backend_error)?; + sub.set_field_name(&item.field_name) + .map_err(backend_error)?; + sub.set_num_candidates((k.max(10) * 5) as i32) + .map_err(backend_error)?; + match item.query { + Some(z_vec::query::Query::Vector(vector)) => sub + .set_query_vector(&vector.values) + .map_err(backend_error)?, + Some(z_vec::query::Query::Fts(fts)) => { + let mut value = Fts::new().map_err(backend_error)?; + if !fts.match_string.is_empty() { + value + .set_match_string(&fts.match_string) + .map_err(backend_error)?; + } + if !fts.query_string.is_empty() { + value + .set_query_string(&fts.query_string) + .map_err(backend_error)?; + } + sub.set_fts(&value).map_err(backend_error)?; + } + None => continue, + } + query.add_sub_query(&sub).map_err(backend_error)?; + } + if !options.hybrid_weights.is_empty() { + let weights = options + .hybrid_weights + .iter() + .map(|weight| f64::from(*weight)) + .collect::>(); + query.set_rerank_weighted(&weights).map_err(backend_error)?; + } else { + query.set_rerank_rrf(60).map_err(backend_error)?; + } + self.collection.multi_query(&query).map_err(backend_error) + } +} + +impl ANN for Service { + fn exists(&self, uuid: String) -> bool { + self.collection + .fetch_with_options(&[uuid.as_str()], None, false) + .map(|docs| !docs.is_empty()) + .unwrap_or(false) + } + + fn create_index(&mut self) -> Result<(), Error> { + self.collection.optimize().map_err(backend_error) + } + + fn save_index(&mut self) -> Result<(), Error> { + self.collection.flush().map_err(backend_error) + } + + fn insert(&mut self, uuid: String, vector: Vec, ts: i64) -> Result<(), Error> { + self.insert_with_options(uuid, vector, ts, &[]) + } + + fn insert_with_options( + &mut self, + uuid: String, + vector: Vec, + ts: i64, + options: &[Any], + ) -> Result<(), Error> { + let doc = self.document(&uuid, &vector, ts, options)?; + Self::write_result( + self.collection.insert(&[&doc]).map_err(backend_error)?, + &[uuid.as_str()], + ) + } + + fn insert_multiple(&mut self, vectors: HashMap>) -> Result<(), Error> { + let mut duplicated = Vec::new(); + for (uuid, vector) in vectors { + if let Err(err) = self.insert(uuid, vector, 0) { + match err { + Error::UUIDAlreadyExists { uuid } => duplicated.push(uuid), + _ => return Err(err), + } + } + } + if duplicated.is_empty() { + Ok(()) + } else { + Err(Error::new_uuid_already_exists(duplicated)) + } + } + + fn update(&mut self, uuid: String, vector: Vec, ts: i64) -> Result<(), Error> { + let doc = self.document(&uuid, &vector, ts, &[])?; + Self::write_result( + self.collection.update(&[&doc]).map_err(backend_error)?, + &[uuid.as_str()], + ) + } + + fn update_multiple(&mut self, vectors: HashMap>) -> Result<(), Error> { + for (uuid, vector) in vectors { + self.update(uuid, vector, 0)?; + } + Ok(()) + } + + fn ready_for_update(&mut self, uuid: String, vector: Vec, ts: i64) -> Result<(), Error> { + self.validate_dimension(&uuid, &vector)?; + if !self.exists(uuid.clone()) { + return Err(Error::UUIDNotFound { uuid }); + } + self.update(uuid, vector, ts) + } + + fn remove(&mut self, uuid: String, _ts: i64) -> Result<(), Error> { + Self::write_result( + self.collection + .delete(&[uuid.as_str()]) + .map_err(backend_error)?, + &[uuid.as_str()], + ) + } + + fn remove_multiple(&mut self, uuids: Vec) -> Result<(), Error> { + let values = uuids.iter().map(String::as_str).collect::>(); + Self::write_result( + self.collection.delete(&values).map_err(backend_error)?, + &values, + ) + } + + fn search( + &self, + vector: Vec, + k: u32, + epsilon: f32, + radius: f32, + ) -> Result { + self.search_with_options(vector, k, epsilon, radius, &[]) + } + + fn search_with_options( + &self, + vector: Vec, + k: u32, + _epsilon: f32, + _radius: f32, + options: &[Any], + ) -> Result { + self.validate_dimension(String::new().as_str(), &vector)?; + let options = decode_any::(options)?; + let docs = self.search_documents(&vector, k, options)?; + Ok(search::Response { + request_id: String::new(), + results: docs + .into_iter() + .filter_map(|doc| { + doc.get_pk().map(|id| Distance { + id: id.to_string(), + distance: doc.get_score(), + }) + }) + .collect(), + }) + } + + fn get_object(&self, uuid: String) -> Result<(Vec, i64), Error> { + let mut docs = self + .collection + .fetch(&[uuid.as_str()]) + .map_err(backend_error)?; + let doc = docs.pop().ok_or(Error::ObjectIDNotFound { uuid })?; + let vector = doc + .get_vector_f32(VECTOR_FIELD) + .map_err(backend_error)? + .ok_or_else(|| Error::Backend { + message: "zvec document has no embedding".to_string(), + })?; + let timestamp = doc + .get_i64(TIMESTAMP_FIELD) + .map_err(backend_error)? + .unwrap_or_default(); + Ok((vector, timestamp)) + } + + fn get_dimension_size(&self) -> usize { + self.dimension + } + + fn len(&self) -> u32 { + self.collection + .stats() + .map(|stats| stats.doc_count.min(u64::from(u32::MAX)) as u32) + .unwrap_or_default() + } + + fn insert_vqueue_buffer_len(&self) -> u32 { + 0 + } + + fn delete_vqueue_buffer_len(&self) -> u32 { + 0 + } + + fn is_indexing(&self) -> bool { + false + } + + fn is_saving(&self) -> bool { + false + } +} + +fn decode_any(options: &[Any]) -> Result, Error> +where + T: Message + Default + prost::Name, +{ + let expected = format!("/{}", T::full_name()); + for option in options { + if option.type_url.ends_with(&expected) { + return T::decode(option.value.as_slice()) + .map(Some) + .map_err(|err| Error::Backend { + message: err.to_string(), + }); + } + } + Ok(None) +} + +fn backend_error(error: zvec_rust::Error) -> Error { + Error::Backend { + message: error.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost::Name; + use std::{ + fs, + path::PathBuf, + sync::atomic::{AtomicU64, Ordering}, + }; + + static NEXT_TEST_ID: AtomicU64 = AtomicU64::new(0); + + struct TestDir(PathBuf); + + impl TestDir { + fn new() -> Self { + let id = NEXT_TEST_ID.fetch_add(1, Ordering::Relaxed); + Self(std::env::temp_dir().join(format!("vald-zvec-{}-{id}", std::process::id()))) + } + } + + impl Drop for TestDir { + fn drop(&mut self) { + if let Err(err) = fs::remove_dir_all(&self.0) { + if err.kind() != std::io::ErrorKind::NotFound { + eprintln!("failed to remove Zvec test directory: {err}"); + } + } + } + } + + fn document_options(fields: &[(&str, &str)]) -> Any { + let value = z_vec::DocumentOptions { + fields: fields + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(), + }; + Any { + type_url: format!( + "type.googleapis.com/{}", + z_vec::DocumentOptions::full_name() + ), + value: value.encode_to_vec(), + } + } + + fn search_options(value: z_vec::SearchOptions) -> Any { + Any { + type_url: format!("type.googleapis.com/{}", z_vec::SearchOptions::full_name()), + value: value.encode_to_vec(), + } + } + + #[test] + fn decode_any_options() { + let tests = [ + ( + "matching type", + vec![document_options(&[("category", "news")])], + Some("news"), + false, + ), + ( + "unrelated type", + vec![Any { + type_url: "type.googleapis.com/example.Unrelated".to_string(), + value: Vec::new(), + }], + None, + false, + ), + ( + "malformed payload", + vec![Any { + type_url: format!( + "type.googleapis.com/{}", + z_vec::DocumentOptions::full_name() + ), + value: vec![0xff], + }], + None, + true, + ), + ]; + + for (name, options, expected, want_error) in tests { + let result = decode_any::(&options); + assert_eq!(result.is_err(), want_error, "{name}"); + if let Ok(value) = result { + assert_eq!( + value.and_then(|value| value.fields.get("category").cloned()), + expected.map(str::to_string), + "{name}" + ); + } + } + } + + #[test] + fn collection_crud_and_search() { + let directory = TestDir::new(); + { + let mut service = Service::new( + directory.0.to_str().expect("test path must be UTF-8"), + 3, + vec!["category".to_string()], + vec!["content".to_string()], + ) + .expect("service creation must succeed"); + let options = [document_options(&[ + ("category", "news"), + ("content", "vector search news"), + ])]; + + service + .insert_with_options("first".to_string(), vec![0.0, 0.0, 0.0], 42, &options) + .expect("insert must succeed"); + service + .insert("second".to_string(), vec![1.0, 1.0, 1.0], 43) + .expect("second insert must succeed"); + + assert!(service.exists("first".to_string())); + assert_eq!(service.len(), 2); + assert_eq!( + service + .get_object("first".to_string()) + .expect("fetch must succeed"), + (vec![0.0, 0.0, 0.0], 42) + ); + + let response = service + .search(vec![0.0, 0.0, 0.0], 2, 0.0, -1.0) + .expect("search must succeed"); + assert_eq!(response.results.len(), 2); + assert_eq!(response.results[0].id, "first"); + assert!(response.results[0].distance <= response.results[1].distance); + + let filtered = service + .search_with_options( + vec![0.0, 0.0, 0.0], + 2, + 0.0, + -1.0, + &[search_options(z_vec::SearchOptions { + pre_filter: "category = 'news'".to_string(), + ..Default::default() + })], + ) + .expect("filtered search must succeed"); + assert_eq!(filtered.results.len(), 1); + assert_eq!(filtered.results[0].id, "first"); + + let full_text = service + .search_with_options( + vec![0.0, 0.0, 0.0], + 2, + 0.0, + -1.0, + &[search_options(z_vec::SearchOptions { + hybrid_queries: vec![z_vec::Query { + field_name: "content".to_string(), + query: Some(z_vec::query::Query::Fts(z_vec::Fts { + match_string: "news".to_string(), + query_string: String::new(), + })), + }], + ..Default::default() + })], + ) + .expect("full-text search must succeed"); + assert_eq!(full_text.results.len(), 1); + assert_eq!(full_text.results[0].id, "first"); + + assert!(matches!( + service.insert("first".to_string(), vec![0.0, 0.0, 0.0], 44), + Err(Error::UUIDAlreadyExists { uuid }) if uuid == "first" + )); + assert!(matches!( + service.insert("bad-dimension".to_string(), vec![0.0, 0.0], 0), + Err(Error::InvalidDimensionSize { .. }) + )); + + service + .remove("first".to_string(), 0) + .expect("remove must succeed"); + assert!(!service.exists("first".to_string())); + } + } +} diff --git a/tests/v2/e2e/assets/zvec_crud.yaml b/tests/v2/e2e/assets/zvec_crud.yaml new file mode 100644 index 0000000000..0d70b70093 --- /dev/null +++ b/tests/v2/e2e/assets/zvec_crud.yaml @@ -0,0 +1,95 @@ +# +# Copyright (C) 2019-2026 vdaas.org vald team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +time_zone: UTC +logging: + format: raw + level: debug + logger: glg +dataset: + name: _E2E_DATASET_PATH_ +kubernetes: + kube_config: _KUBECONFIG_ + port_forward: + enabled: true + local_port: 8082 + namespace: _E2E_TARGET_NAMESPACE_ + service_name: _E2E_TARGET_NAME_ + target_port: 8081 +target: + addrs: + - 127.0.0.1:8082 + health_check_duration: 1s + connection_pool: + enable_dns_resolver: true + enable_rebalance: true + old_conn_close_duration: 2m + rebalance_duration: 30m + size: 3 + backoff: + backoff_factor: 1.1 + backoff_time_limit: 5s + enable_error_log: false + initial_duration: 5ms + jitter_limit: 100ms + maximum_duration: 5s + retry_count: 100 + call_option: + wait_for_ready: true + dial_option: + insecure: true + min_connection_timeout: 20s + net: + network: tcp + socket_option: + reuse_addr: true + reuse_port: true +metadata: {} +metrics: + enabled: false +strategies: + - concurrency: 1 + name: Insert and search through zvec agent + operations: + - name: Insert -> IndexInfo -> Search + executions: + - mode: unary + name: Initial IndexInfo + type: index_info + expect: + - value: {} + - mode: unary + name: Insert + type: insert + parallelism: _E2E_PARALLELISM_ + num: _E2E_INSERT_COUNT_ + qps: _E2E_QPS_ + wait: 10s + - mode: unary + name: IndexInfo + type: index_info + retry_until_success_timeout: 3m + expect: + - status_code: ok + path: $.stored + value: _E2E_EXPECTED_INDEX_ + - mode: unary + name: Search + type: search + parallelism: _E2E_PARALLELISM_ + num: _E2E_SEARCH_COUNT_ + search: + timeout: 3s + k: 10