wehub/alibaba--zvec
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.clang-format
+31
View File
@@ -0,0 +1,31 @@
# http://clang.llvm.org/docs/ClangFormatStyleOptions.html
# Defines the Google C++ style for automatic reformatting.
BasedOnStyle: Google
MaxEmptyLinesToKeep: 2
DerivePointerAlignment: false
PointerAlignment: Right
AllowShortFunctionsOnASingleLine: Empty
IncludeBlocks: Merge
IncludeCategories:
- Regex: '^<linux/[0-9A-Za-z]+'
Priority: 100
- Regex: '^<mach/[0-9A-Za-z]+'
Priority: 101
- Regex: '^<mach-o/[0-9A-Za-z]+'
Priority: 102
- Regex: '^<sys/[0-9A-Za-z]+'
Priority: 103
- Regex: '^<[0-9A-Za-z]+\.h>$'
Priority: 200
- Regex: '^<[0-9A-Za-z_]+>$'
Priority: 201
- Regex: '^<[0-9A-Za-z_]+\.[0-9A-Za-z]+>$'
Priority: 202
- Regex: '^<[0-9A-Za-z_]+/[0-9A-Za-z]+'
Priority: 203
- Regex: '^\"[0-9A-Za-z_]+/[0-9A-Za-z]+'
Priority: 300
- Regex: '^\"[0-9A-Za-z_]+\.[0-9A-Za-z]+\"$'
Priority: 301
- Regex: '.*'
Priority: 1000
.clang-tidy
+10
View File
@@ -0,0 +1,10 @@
Checks: >
-*,
bugprone-multiple-statement-macro,
bugprone-sizeof-expression,
modernize-use-nullptr,
modernize-use-override,
WarningsAsErrors: "*"
HeaderFilterRegex: "^(src|tests|tools)/(?!db/sqlengine/antlr/gen/|db/index/column/fts_column/gen/|include/zvec/ailego/encoding/json/mod_json\\.h).*"
FormatStyle: none
SystemHeaders: false
.gitattributes
+3
View File
@@ -0,0 +1,3 @@
# Auto-generated files — collapsed in GitHub PR diffs
src/db/index/column/fts_column/gen/** linguist-generated=true
src/db/sqlengine/antlr/gen/** linguist-generated=true
.github/CODEOWNERS
+38
View File
@@ -0,0 +1,38 @@
# CODEOWNERS - Auto-assign reviewers based on file paths
# Documentation: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
# Default: feihongxu0824 reviews everything not covered below
* @feihongxu0824
# .github directory
/.github/ @Cuiyus
/python/ @Cuiyus
# cmake, examples, scripts
/cmake/ @egolearner
/examples/ @egolearner
/scripts/ @egolearner
# src/db, tests
/src/db/ @zhourrr
/tests/db @zhourrr
# src/core, src/turbo
/src/core/ @richyreachy
/src/turbo/ @richyreachy
# src/ailego
/src/ailego/ @iaojnh
/tests/ailego @iaojnh
/tests/core @iaojnh
# tools
/tools/ @JalinWang
# src/binding
/src/binding/ @Cuiyus
# src/include, thirdparty, .gitmodules
/src/include/ @chinaux
/thirdparty/ @chinaux
.gitmodules @chinaux
.github/ISSUE_TEMPLATE/benchmark.yml
+49
View File
@@ -0,0 +1,49 @@
name: Benchmarking
description: Add, update, or fix benchmark cases for zvec
title: "[Benchmark]: "
labels: ["benchmark"]
body:
- type: markdown
attributes:
value: |
Use this for benchmark-related work: new test cases, CI integration, or performance regression tracking.
- type: input
id: benchmark_type
attributes:
label: Benchmark Type
description: e.g., filtered search, batch insert, recall@k, ARM64 vs x86
validations:
required: true
- type: textarea
id: goal
attributes:
label: Goal
description: What performance aspect are you measuring or improving?
validations:
required: true
- type: textarea
id: methodology
attributes:
label: Methodology
description: Dataset, query size, hardware, metrics (latency, throughput, memory)
validations:
required: true
- type: textarea
id: baseline
attributes:
label: Baseline (if applicable)
description: Current performance numbers or competing systems for comparison.
validations:
required: false
- type: textarea
id: ci_integration
attributes:
label: CI Integration Plan
description: Should this run in CI? How often?
validations:
required: false
.github/ISSUE_TEMPLATE/bug_report.yml
+73
View File
@@ -0,0 +1,73 @@
name: Bug Report
description: Report a bug or unexpected behavior (e.g., crash, incorrect vector query, memory leak)
title: "[Bug]: "
labels: ["bug", "triage"]
body:
- type: markdown
attributes:
value: |
Thank you for reporting! Please provide detailed info so we can reproduce and fix it quickly.
- type: textarea
id: description
attributes:
label: Description
description: What happened? What did you expect?
placeholder: |
e.g. "Query with vector field crashes when using Zvec Python API"
validations:
required: true
- type: textarea
id: steps_to_reproduce
attributes:
label: Steps to Reproduce
description: Exact steps to trigger the issue (code snippets welcome)
placeholder: |
1. Build Zvec with CMake (Debug/Release)
2. Run Python script: `python test.py`
3. Call `collection.query(VectorQuery())`
4. Process segfaults / hangs / returns wrong results
render: python
validations:
required: true
- type: textarea
id: logs_or_trace
attributes:
label: Logs / Stack Trace
description: Paste relevant logs, LLDB/GDB backtrace, or CI failures
placeholder: |
Thread 1 "python" received signal SIGSEGV, Segmentation fault.
0x0000000104a2c3f0 in std::__1::shared_ptr<...>::...
render: shell
validations:
required: false
- type: input
id: os
attributes:
label: Operating System
placeholder: macOS 14 (M1), Ubuntu 22.04, Windows 11 (WSL2)
validations:
required: true
- type: input
id: build_env
attributes:
label: Build & Runtime Environment
description: Compiler, CMake, Python, key dependencies
placeholder: |
clang 15.0.0, CMake 4.1.2, Python 3.11.9, magic_enum v0.9.7 (via git submodule)
validations:
required: true
- type: checkboxes
id: additional_context
attributes:
label: Additional Context
options:
- label: I've checked `git status` — no uncommitted submodule changes
- label: I built with `CMAKE_BUILD_TYPE=Debug`
- label: This occurs with or without `COVERAGE=ON`
- label: The issue involves Python ↔ C++ integration (pybind11)
.github/ISSUE_TEMPLATE/config.yml
+9
View File
@@ -0,0 +1,9 @@
blank_issues_enabled: false
contact_links:
- name: Documentation
url: https://zvec.org/en/
about: Check the quickstart, build guide, and API docs first.
- name: Python API Examples
url: https://zvec.org/en/docs/quickstart/
about: See working usage examples.
.github/ISSUE_TEMPLATE/enhancement.yml
+41
View File
@@ -0,0 +1,41 @@
name: Enhancement
description: Improve an existing feature or component
title: "[Enhance]: "
labels: ["enhancement"]
body:
- type: markdown
attributes:
value: |
This template is for improving existing functionality (e.g., performance, usability, robustness).
- type: input
id: component
attributes:
label: Affected Component
description: e.g., HNSW index, buffer manager, Python API
validations:
required: true
- type: textarea
id: current
attributes:
label: Current Behavior
description: What is the current state and its limitations?
validations:
required: true
- type: textarea
id: desired
attributes:
label: Desired Improvement
description: What should be improved and how?
validations:
required: true
- type: textarea
id: impact
attributes:
label: Impact
description: How will this benefit users? (e.g., faster queries, lower memory, easier integration)
validations:
required: true
.github/ISSUE_TEMPLATE/feature_request.yml
+55
View File
@@ -0,0 +1,55 @@
name: Feature Request
description: Suggest a new feature or improvement (e.g., better memory control, new query option)
title: "[Feature]: "
labels: ["feature"]
body:
- type: markdown
attributes:
value: |
Thanks for your idea! Help us understand the motivation and scope.
- type: textarea
id: problem_or_motivation
attributes:
label: Problem / Motivation
description: What problem does this solve? Why is it needed?
placeholder: |
e.g. "Current vector queries don't allow filtering by metadata + distance threshold at once"
validations:
required: true
- type: textarea
id: proposed_solution
attributes:
label: Proposed Solution
description: How should it work? API sketch or pseudocode welcome.
placeholder: |
Add `filter=` and `max_distance=` args to `Zvec.query()`:
```python
results = db.query(vector, filter="category == 'A'", max_distance=0.5)
```
render: python
validations:
required: false
- type: textarea
id: alternatives
attributes:
label: Alternatives Considered
description: Are there workarounds? Why not use them?
validations:
required: false
- type: dropdown
id: impact_area
attributes:
label: Affected Area
multiple: true
options:
- label: C++ Core (storage, indexing)
- label: Python API / Bindings
- label: Build System (CMake, Homebrew pkg)
- label: Testing / CI / Coverage
- label: Documentation
validations:
required: false
.github/ISSUE_TEMPLATE/integration.yml
+36
View File
@@ -0,0 +1,36 @@
name: Ecosystem Integration
description: Integrate zvec with external frameworks (e.g., LangChain, LlamaIndex)
title: "[Integration]: "
labels: ["integration"]
body:
- type: input
id: framework
attributes:
label: Target Framework
description: e.g., LangChain, LlamaIndex, Haystack
validations:
required: true
- type: textarea
id: motivation
attributes:
label: Motivation
description: Why integrate with this framework? Who benefits?
validations:
required: true
- type: textarea
id: interface
attributes:
label: Required Interface
description: What adapter or interface must be implemented? (e.g., VectorStore base class)
validations:
required: true
- type: textarea
id: reference
attributes:
label: Reference Implementations
description: Links to similar integrations in other vector DBs.
validations:
required: false
.github/ISSUE_TEMPLATE/profiling.yml
+41
View File
@@ -0,0 +1,41 @@
name: Profiling / Investigation
description: Profile performance, compatibility, or behavior in a specific scenario
title: "[Profile]: "
labels: ["profile"]
body:
- type: markdown
attributes:
value: |
Use this for tasks like performance profiling, architecture compatibility checks, or feasibility studies.
- type: input
id: scenario
attributes:
label: Target Scenario
description: e.g., ARM64 deployment, high-concurrency load, large dataset ingestion
validations:
required: true
- type: textarea
id: objective
attributes:
label: Objective
description: What do you want to learn or validate?
validations:
required: true
- type: textarea
id: methodology
attributes:
label: Proposed Methodology
description: How will you conduct the investigation? (tools, metrics, test data)
validations:
required: true
- type: textarea
id: expected_outcome
attributes:
label: Expected Outcome
description: What deliverables are expected? (e.g., report, optimization PR, benchmark results)
validations:
required: true
.github/auto-assign-config.yml
+24
View File
@@ -0,0 +1,24 @@
# Auto-assign configuration
# Documentation: https://github.com/kentaro-m/auto-assign-action
#
# NOTE: Reviewers are assigned via CODEOWNERS file based on file paths.
# This config only handles assignee assignment.
# Automatically add PR author as assignee
addAssignees: author
# Number of assignees to add (0 = add all from the list)
numberOfAssignees: 0
# Reviewers are handled by CODEOWNERS, not here
addReviewers: false
# Skip draft PRs
skipDraft: true
# Skip keywords in PR title (won't assign if title contains these)
skipKeywords:
- wip
- WIP
- draft
- DRAFT
.github/cmake/subproject-integration/CMakeLists.txt
+26
View File
@@ -0,0 +1,26 @@
cmake_minimum_required(VERSION 3.13)
project(zvec_subproject_integration LANGUAGES C CXX)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(ZVEC_SOURCE_DIR "" CACHE PATH "Path to the zvec source tree")
if(NOT ZVEC_SOURCE_DIR)
message(FATAL_ERROR "ZVEC_SOURCE_DIR is required")
endif()
add_subdirectory("${ZVEC_SOURCE_DIR}" zvec)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin")
add_executable(zvec_embed_db_example
"${ZVEC_SOURCE_DIR}/examples/c++/db/main.cc")
target_link_libraries(zvec_embed_db_example PRIVATE zvec_shared)
add_executable(zvec_embed_core_example
"${ZVEC_SOURCE_DIR}/examples/c++/core/main.cc")
target_link_libraries(zvec_embed_core_example PRIVATE zvec_core_shared)
add_executable(zvec_embed_ailego_example
"${ZVEC_SOURCE_DIR}/examples/c++/ailego/main.cc")
target_link_libraries(zvec_embed_ailego_example PRIVATE zvec_ailego_shared)
.github/codecov.yml
+32
View File
@@ -0,0 +1,32 @@
codecov:
require_ci_to_pass: true
coverage:
precision: 2
round: down
range: "60...75"
status:
project:
default: false
patch:
default: false
parsers:
gcov:
branch_detection:
conditional: true
loop: true
method: false
macro: false
comment:
require_changes: false
layout: "reach,diff,flags,tree"
behavior: default
ignore:
- "thirdparty/"
- "tests/"
.github/dependabot.yml
+17
View File
@@ -0,0 +1,17 @@
version: 2
updates:
# GitHub Actions dependencies
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
day: "monday"
time: "02:00"
timezone: "Asia/Shanghai"
labels:
- "dependencies"
- "github-actions"
commit-message:
prefix: "ci"
include: "scope"
open-pull-requests-limit: 5
.github/workflows/01-ci-pipeline.yml
+116
View File
@@ -0,0 +1,116 @@
name: Main
on:
push:
branches: [ "main" ]
paths-ignore: &paths-ignore
- '**.md'
- '.github/CODEOWNERS'
- '.github/ISSUE_TEMPLATE/**'
- '.github/auto-assign-config.yml'
- '.github/codecov.yml'
- '.github/dependabot.yml'
- '.github/workflows/auto-assign.yml'
- '.github/workflows/community-pr-handler.yml'
- '.github/workflows/issue-auto-assign.yml'
- '.github/workflows/build_wheel.yml'
- '.github/workflows/build_test_wheel.yml'
- '.github/workflows/_build_wheel_job.yml'
- '.github/workflows/continuous_bench.yml'
- '.github/workflows/nightly_coverage.yml'
- '.github/workflows/07-linux-riscv-build.yml'
- '.github/workflows/docker/**'
- '.github/workflows/scripts/**'
merge_group:
pull_request:
branches: [ "main" ]
paths-ignore: *paths-ignore
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || '' }}-${{ github.base_ref || '' }}-${{ github.ref != 'refs/heads/main' || github.sha }}
cancel-in-progress: true
permissions:
contents: read
jobs:
# Code quality checks (fast, run first)
lint:
if: >-
github.event_name != 'pull_request' ||
(!github.event.pull_request.draft &&
!contains(github.event.pull_request.title, 'wip'))
uses: ./.github/workflows/02-lint-check.yml
# Static analysis: clang-tidy on changed C/C++ files
clang-tidy:
if: >-
github.event_name != 'pull_request' ||
(!github.event.pull_request.draft &&
!contains(github.event.pull_request.title, 'wip'))
uses: ./.github/workflows/clang_tidy.yml
# Main build and test matrix
build-and-test-macos-arm64:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (macos-arm64)
needs: [lint, clang-tidy]
uses: ./.github/workflows/03-macos-linux-build.yml
with:
platform: macos-arm64
os: macos-15
build-and-test-macos-26-arm64:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (macos-26-arm64)
needs: [lint, clang-tidy]
uses: ./.github/workflows/03-macos-linux-build.yml
with:
platform: macos-arm64
os: macos-26
build-and-test-linux-arm64:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (linux-arm64)
needs: [lint, clang-tidy]
uses: ./.github/workflows/03-macos-linux-build.yml
with:
platform: linux-arm64
os: ubuntu-24.04-arm
build-and-test-linux-x64:
name: Build & Test (linux-x64)
needs: [lint, clang-tidy]
uses: ./.github/workflows/03-macos-linux-build.yml
with:
platform: linux-x64
os: ubuntu-24.04
build-and-test-linux-x64-clang:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (linux-x64-clang)
needs: [lint, clang-tidy]
uses: ./.github/workflows/03-macos-linux-build.yml
with:
platform: linux-x64-clang
os: ubuntu-24.04
compiler: clang
build-android:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (android)
needs: [lint, clang-tidy]
uses: ./.github/workflows/04-android-build.yml
build-and-test-on-windows:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (Windows)
needs: [lint, clang-tidy]
uses: ./.github/workflows/05-windows-build.yml
build-ios:
if: github.event_name != 'push' || github.ref != 'refs/heads/main'
name: Build & Test (iOS)
needs: [lint, clang-tidy]
uses: ./.github/workflows/06-ios-build.yml
.github/workflows/02-lint-check.yml
+52
View File
@@ -0,0 +1,52 @@
name: Lint
on:
workflow_call:
jobs:
lint:
name: Code Quality Checks
runs-on: ubuntu-24.04
steps:
- name: Checkout code
uses: actions/checkout@v7
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
cache: 'pip'
cache-dependency-path: 'pyproject.toml'
- name: Install linting tools
run: |
python -m pip install --upgrade pip \
ruff==v0.14.4 \
clang-format==18.1.8
shell: bash
- name: Run Ruff Linter
run: python -m ruff check .
shell: bash
- name: Run Ruff Formatter Check
run: python -m ruff format --check .
shell: bash
- name: Run clang-format Check
run: |
CPP_FILES=$(find . -type f \( -name "*.cpp" -o -name "*.h" -o -name "*.hpp" -o -name "*.cc" -o -name "*.cxx" \) \
! -path "./build/*" \
! -path "./tests/*" \
! -path "./scripts/*" \
! -path "./python/*" \
! -path "./thirdparty/*" \
! -path "./.git/*")
if [ -z "$CPP_FILES" ]; then
echo "No C++ files found to check."
exit 0
fi
clang-format --dry-run --Werror $CPP_FILES
shell: bash
.github/workflows/03-macos-linux-build.yml
+185
View File
@@ -0,0 +1,185 @@
name: MacOS & Linux Build
on:
workflow_call:
inputs:
platform:
description: 'Platform identifier'
required: true
type: string
os:
description: 'GitHub Actions runner OS'
required: true
type: string
compiler:
description: 'C++ compiler to use (gcc or clang)'
required: false
type: string
default: ''
permissions:
contents: read
jobs:
# Build and test matrix (parallel execution)
build-and-test:
name: Build & Test (${{ inputs.platform }})
runs-on: ${{ inputs.os }}
strategy:
fail-fast: false
matrix:
include:
- os: ${{ inputs.os }}
platform: ${{ inputs.platform }}
arch_flag: "" # Use appropriate architecture
steps:
- name: Checkout code
uses: actions/checkout@v7
with:
submodules: recursive
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: ${{ inputs.platform }}-${{ inputs.os }}-${{ inputs.compiler || 'gcc' }}
max-size: 150M
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
cache: 'pip'
cache-dependency-path: 'pyproject.toml'
- name: Install Clang
if: inputs.compiler == 'clang' && runner.os == 'Linux'
run: |
sudo apt-get update -y
sudo apt-get install -y clang libomp-dev
shell: bash
- name: Install AIO
if: runner.os == 'Linux' && runner.arch == 'X64'
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libaio-dev
shell: bash
- name: Print CPU info
if: runner.os == 'Linux'
run: lscpu
shell: bash
- name: Set up environment variables
run: |
# Set number of processors for parallel builds
if [[ "${{ matrix.platform }}" == "macos-arm64" ]]; then
NPROC=$(sysctl -n hw.ncpu 2>/dev/null || echo 2)
else
NPROC=$(nproc 2>/dev/null || echo 2)
fi
echo "NPROC=$NPROC" >> $GITHUB_ENV
echo "Using $NPROC parallel jobs for builds"
# Set compiler when clang is requested
if [[ "${{ inputs.compiler }}" == "clang" ]]; then
echo "CC=clang" >> $GITHUB_ENV
echo "CXX=clang++" >> $GITHUB_ENV
fi
# Add Python user base bin to PATH for pip-installed CLI tools
echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH
shell: bash
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install \
pybind11==3.0 \
cmake==3.30.0 \
ninja==1.11.1 \
pytest \
pytest-xdist \
scikit-build-core \
setuptools_scm
shell: bash
- name: Check CMake subproject integration
if: matrix.platform == 'linux-x64'
run: |
SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-smoke"
SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-build"
mkdir -p "$SMOKE_SOURCE" "$SMOKE_BUILD"
cp "$GITHUB_WORKSPACE/.github/cmake/subproject-integration/CMakeLists.txt" \
"$SMOKE_SOURCE/CMakeLists.txt"
cmake -S "$SMOKE_SOURCE" -B "$SMOKE_BUILD" -G Ninja \
-DZVEC_SOURCE_DIR="$GITHUB_WORKSPACE" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TOOLS=OFF \
-DBUILD_C_BINDINGS=OFF \
-DBUILD_PYTHON_BINDINGS=OFF \
-DENABLE_WERROR=ON \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
shell: bash
- name: Build from source
run: |
cd "$GITHUB_WORKSPACE"
CMAKE_GENERATOR="Ninja" \
CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \
python -m pip install -v . \
--no-build-isolation \
--config-settings='cmake.define.BUILD_TOOLS=ON' \
--config-settings='cmake.define.ENABLE_WERROR=ON' \
--config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \
--config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache' \
${{ matrix.arch_flag }}
shell: bash
- name: Run C++ Tests
run: |
cd "$GITHUB_WORKSPACE/build"
cmake --build . --target unittest --parallel $NPROC
shell: bash
- name: Run Python Tests
run: |
cd "$GITHUB_WORKSPACE"
python -m pytest python/tests/
shell: bash
- name: Run C++ Examples
run: |
cd "$GITHUB_WORKSPACE/examples/c++"
mkdir build && cd build
cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build . --parallel $NPROC
./db-example
./core-example
./ailego-example
shell: bash
- name: Run C Examples
run: |
cd "$GITHUB_WORKSPACE/examples/c"
mkdir build && cd build
cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build . --parallel $NPROC
./c_api_basic_example
./c_api_collection_schema_example
./c_api_doc_example
./c_api_field_schema_example
./c_api_index_example
./c_api_optimized_example
shell: bash
.github/workflows/04-android-build.yml
+437
View File
@@ -0,0 +1,437 @@
name: Android Cross Build & Test
on:
workflow_call:
workflow_dispatch:
inputs:
api:
description: 'Android API level'
required: false
default: '34'
type: string
permissions:
contents: read
env:
NDK_VERSION: '26.1.10909125'
jobs:
build-and-test:
runs-on: ubuntu-24.04
timeout-minutes: 120
strategy:
fail-fast: false
matrix:
abi: [x86_64]
api: ${{ github.event.inputs.api && fromJSON(format('["{0}"]', github.event.inputs.api)) || fromJSON('["34"]') }}
steps:
# ── Environment setup ──────────────────────────────────────────────
- name: Checkout
uses: actions/checkout@v7
with:
submodules: recursive
- name: Install dependencies
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
cmake ninja-build git ca-certificates python3 \
build-essential make unzip curl
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: android-${{ matrix.abi }}
max-size: 300M
- name: Setup Java 17
uses: actions/setup-java@v5
with:
distribution: temurin
java-version: '17'
- name: Setup Android SDK
uses: android-actions/setup-android@v4
- name: Enable KVM
run: sudo chmod 666 /dev/kvm || true
- name: Install NDK, emulator and system image
shell: bash
run: |
sdkmanager --install \
"ndk;$NDK_VERSION" \
"platform-tools" \
"platforms;android-${{ matrix.api }}" \
"emulator"
# Install x86_64 system image (try variants in order of availability)
sdkmanager --install "system-images;android-${{ matrix.api }};google_apis;x86_64" 2>/dev/null || \
sdkmanager --install "system-images;android-${{ matrix.api }};google_apis_playstore;x86_64" 2>/dev/null || \
sdkmanager --install "system-images;android-${{ matrix.api }};default;x86_64"
# ── Step 1: build host protoc (using HOST compiler, NOT NDK) ───────
- name: Cache host protoc
uses: actions/cache@v6
with:
path: build_host
key: ${{ runner.os }}-host-protoc-${{ hashFiles('thirdparty/protobuf/**', 'CMakeLists.txt') }}
restore-keys: |
${{ runner.os }}-host-protoc-
- name: 'Step 1: Build host protoc'
shell: bash
run: |
if [ ! -f "build_host/bin/protoc" ]; then
git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true
cmake -S . -B build_host \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_TOOLCHAIN_FILE="" \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-G Ninja
cmake --build build_host --target protoc --parallel
else
echo "Using cached host protoc"
fi
# ── Step 2: cross-compile zvec + tests for Android ─────────────────
- name: 'Step 2: Cross-compile zvec and tests'
shell: bash
env:
BUILD_DIR: build_android_${{ matrix.abi }}
run: |
ANDROID_NDK_HOME="$ANDROID_HOME/ndk/$NDK_VERSION"
# Reset thirdparty so the cross toolchain can patch cleanly
git submodule foreach --recursive 'git stash --include-untracked' 2>/dev/null || true
# Force reconfigure to pick up any cmake changes
rm -f "$BUILD_DIR/CMakeCache.txt"
cmake -S . -B "$BUILD_DIR" -G Ninja \
-DANDROID_NDK="$ANDROID_NDK_HOME" \
-DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \
-DANDROID_ABI=${{ matrix.abi }} \
-DANDROID_NATIVE_API_LEVEL=${{ matrix.api }} \
-DANDROID_STL=c++_static \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_PYTHON_BINDINGS=OFF \
-DBUILD_TOOLS=OFF \
-DENABLE_NATIVE=OFF \
-DAUTO_DETECT_ARCH=OFF \
-DENABLE_WERROR=ON \
-DCMAKE_INSTALL_PREFIX="$BUILD_DIR/install" \
-DGLOBAL_CC_PROTOBUF_PROTOC="$GITHUB_WORKSPACE/build_host/bin/protoc" \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
echo "Building all targets..."
cmake --build "$BUILD_DIR" --parallel
# Discover test targets from ctest metadata
echo "Discovering test targets..."
TEST_NAMES=()
while IFS= read -r line; do
name=$(echo "$line" | sed -n 's/.*Test[[:space:]]*#[0-9]*:[[:space:]]*//p')
[ -n "$name" ] && TEST_NAMES+=("$name")
done < <(cd "$BUILD_DIR" && ctest --show-only 2>/dev/null || true)
# Fallback: scan ninja targets for *_test
if [ ${#TEST_NAMES[@]} -eq 0 ]; then
echo "ctest unavailable, scanning ninja targets..."
while IFS= read -r line; do
name=$(echo "$line" | sed -n 's/^\([a-zA-Z0-9_]*_test\): .*/\1/p')
[ -n "$name" ] && TEST_NAMES+=("$name")
done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true)
fi
echo "Building ${#TEST_NAMES[@]} test executables..."
ninja -C "$BUILD_DIR" -j$(nproc) "${TEST_NAMES[@]}"
# ── Step 3: start emulator ─────────────────────────────────────────
- name: 'Step 3: Start Android emulator'
shell: bash
run: |
AVD_NAME="zvec_test_avd"
AVD_DIR="$HOME/.android/avd/${AVD_NAME}.avd"
AVD_INI="$HOME/.android/avd/${AVD_NAME}.ini"
# Find installed system image (same priority as build_android.sh)
SYS_IMG=""
for variant in google_apis google_apis_playstore default; do
candidate="$ANDROID_HOME/system-images/android-${{ matrix.api }}/$variant/x86_64"
if [ -d "$candidate" ]; then
SYS_IMG="$candidate"
break
fi
done
if [ -z "$SYS_IMG" ]; then
SYS_IMG=$(find "$ANDROID_HOME/system-images" -type d -name "x86_64" 2>/dev/null | head -1)
fi
if [ -z "$SYS_IMG" ]; then
echo "ERROR: No x86_64 system image found"
exit 1
fi
echo "System image: $SYS_IMG"
# Extract tag from path (e.g. .../google_apis/x86_64 -> google_apis)
SYS_TAG=$(basename "$(dirname "$SYS_IMG")")
# Create AVD via INI files (faster and more reliable than avdmanager)
mkdir -p "$AVD_DIR"
cat > "$AVD_INI" << EOAVD
avd.ini.encoding=UTF-8
path=${AVD_DIR}
path.rel=avd/${AVD_NAME}.avd
target=android-${{ matrix.api }}
EOAVD
cat > "$AVD_DIR/config.ini" << EOCFG
AvdId=${AVD_NAME}
PlayStore.enabled=false
abi.type=x86_64
avd.ini.displayname=${AVD_NAME}
avd.ini.encoding=UTF-8
disk.dataPartition.size=8G
hw.accelerator.isConfigured=true
hw.cpu.arch=x86_64
hw.cpu.ncore=4
hw.lcd.density=420
hw.lcd.height=1920
hw.lcd.width=1080
hw.ramSize=4096
image.sysdir.1=${SYS_IMG}/
tag.display=${SYS_TAG}
tag.id=${SYS_TAG}
EOCFG
echo "Created AVD: $AVD_NAME"
# Launch emulator in background
$ANDROID_HOME/emulator/emulator -avd "$AVD_NAME" \
-no-window -no-audio -no-boot-anim \
-gpu swiftshader_indirect \
-netdelay none -netspeed full \
-memory 4096 \
-no-snapshot \
-wipe-data &
echo "EMULATOR_PID=$!" >> "$GITHUB_ENV"
# Wait for device to be reachable
adb wait-for-device
# Poll for boot completion with timeout
echo "Waiting for boot to complete..."
TIMEOUT=300
ELAPSED=0
while true; do
BOOTED=$(adb shell getprop sys.boot_completed 2>/dev/null | tr -d '\r\n ' || true)
if [ "$BOOTED" = "1" ]; then
echo "Emulator ready (took ${ELAPSED}s)"
break
fi
if [ $ELAPSED -ge $TIMEOUT ]; then
echo "ERROR: Emulator failed to boot within ${TIMEOUT}s"
exit 1
fi
sleep 3
ELAPSED=$((ELAPSED + 3))
done
echo "Device ABI: $(adb shell getprop ro.product.cpu.abi | tr -d '\r')"
echo "ABI list : $(adb shell getprop ro.product.cpu.abilist | tr -d '\r')"
# ── Step 4: run unit tests on emulator ─────────────────────────────
- name: 'Step 4: Run unit tests on emulator'
shell: bash
env:
BUILD_DIR: build_android_${{ matrix.abi }}
run: |
DEVICE_TEST_DIR="/data/local/tmp/zvec_tests"
DEVICE_LIB_DIR="$DEVICE_TEST_DIR/lib"
adb shell "mkdir -p $DEVICE_TEST_DIR $DEVICE_LIB_DIR"
# Push shared libraries
echo "Pushing shared libraries..."
SO_COUNT=0
while IFS= read -r so_file; do
adb push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1
SO_COUNT=$((SO_COUNT + 1))
done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null)
echo "Pushed $SO_COUNT shared libraries"
# Push helper binaries (needed by crash_recovery tests which fork+exec them)
echo "Pushing helper binaries..."
for helper_name in data_generator collection_optimizer; do
helper_path=$(find "$BUILD_DIR" -name "$helper_name" -type f -executable ! -name "*_test" 2>/dev/null | head -1)
if [ -n "$helper_path" ]; then
adb push "$helper_path" "$DEVICE_TEST_DIR/$helper_name" > /dev/null 2>&1
adb shell "chmod 755 $DEVICE_TEST_DIR/$helper_name"
echo " Pushed $helper_name"
fi
done
# Discover test targets from ctest metadata
TEST_NAMES=()
while IFS= read -r line; do
name=$(echo "$line" | sed -n 's/.*Test[[:space:]]*#[0-9]*:[[:space:]]*//p')
[ -n "$name" ] && TEST_NAMES+=("$name")
done < <(cd "$BUILD_DIR" && ctest --show-only 2>/dev/null || true)
if [ ${#TEST_NAMES[@]} -eq 0 ]; then
while IFS= read -r line; do
name=$(echo "$line" | sed -n 's/^\([a-zA-Z0-9_]*_test\): .*/\1/p')
[ -n "$name" ] && TEST_NAMES+=("$name")
done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true)
fi
# Collect test binaries
TEST_BINS=()
for name in "${TEST_NAMES[@]}"; do
bin_path=$(find "$BUILD_DIR" -name "$name" -type f -executable 2>/dev/null | head -1)
if [ -n "$bin_path" ]; then
TEST_BINS+=("$bin_path")
else
echo "WARNING: binary not found for '$name'"
fi
done
TOTAL=${#TEST_BINS[@]}
PASSED=0
FAILED=0
FAILED_NAMES=()
IDX=0
echo "Running $TOTAL unit tests on emulator..."
for test_bin in "${TEST_BINS[@]}"; do
IDX=$((IDX + 1))
test_name=$(basename "$test_bin")
device_path="$DEVICE_TEST_DIR/$test_name"
# Give each test its own working directory to avoid name collisions
WORK_DIR="$DEVICE_TEST_DIR/workdir_${test_name}"
echo ""
echo "────────────────────────────────────────"
echo " [$IDX/$TOTAL] $test_name"
echo "────────────────────────────────────────"
set +e
# Create isolated working directory
adb shell "mkdir -p $WORK_DIR" 2>/dev/null
# Copy helper binaries into working directory so crash_recovery tests
# (which fork+exec data_generator / collection_optimizer) can find them
adb shell "for h in $DEVICE_TEST_DIR/data_generator $DEVICE_TEST_DIR/collection_optimizer; do [ -f \$h ] && cp \$h $WORK_DIR/; done" 2>/dev/null
# Push test binary
adb push "$test_bin" "$device_path" > /dev/null 2>&1
adb shell "chmod 755 $device_path" 2>/dev/null
# Run test from its own working directory with LD_LIBRARY_PATH
OUTPUT=$(adb shell "cd $WORK_DIR && LD_LIBRARY_PATH=$DEVICE_LIB_DIR $device_path 2>&1; echo EXIT_CODE=\$?" 2>&1)
# Extract exit code from the output
EXIT_CODE=$(echo "$OUTPUT" | grep -o 'EXIT_CODE=[0-9]*' | tail -1 | cut -d= -f2)
set -e
# Print test output (without the EXIT_CODE marker)
echo "$OUTPUT" | grep -v 'EXIT_CODE=' | sed 's/^/ /' || true
if [ "$EXIT_CODE" = "0" ]; then
echo " >>> PASSED"
PASSED=$((PASSED + 1))
else
# Detect "crash-on-exit" pattern: all gtest assertions passed but
# process crashed during static destructor teardown (common with c++_static STL)
GTEST_PASSED_LINE=$(echo "$OUTPUT" | grep '\[ PASSED \]' | tail -1 || true)
GTEST_FAILED_LINE=$(echo "$OUTPUT" | grep '\[ FAILED \]' | head -1 || true)
if [ -n "$GTEST_PASSED_LINE" ] && [ -z "$GTEST_FAILED_LINE" ] && \
{ [ "$EXIT_CODE" = "139" ] || [ "$EXIT_CODE" = "134" ] || [ "$EXIT_CODE" = "135" ]; }; then
echo " >>> PASSED (crash-on-exit ignored, exit=$EXIT_CODE)"
PASSED=$((PASSED + 1))
else
echo " >>> FAILED (exit=$EXIT_CODE)"
FAILED=$((FAILED + 1))
FAILED_NAMES+=("$test_name")
fi
fi
# Clean up binary and working directory to reclaim disk space
adb shell "rm -rf $device_path $WORK_DIR" 2>/dev/null || true
done
echo ""
echo "============================================================"
echo " Test Summary"
echo "============================================================"
echo " Total : $TOTAL"
echo " Passed : $PASSED"
echo " Failed : $FAILED"
if [ $FAILED -gt 0 ]; then
echo ""
echo " Failed tests:"
for name in "${FAILED_NAMES[@]}"; do
echo " - $name"
done
fi
echo "============================================================"
if [ $FAILED -gt 0 ]; then
exit 1
fi
echo "All tests passed!"
# ── Step 5: build and run examples ─────────────────────────────────
- name: 'Step 5: Build and run examples'
shell: bash
env:
BUILD_DIR: build_android_${{ matrix.abi }}
run: |
ANDROID_NDK_HOME="$ANDROID_HOME/ndk/$NDK_VERSION"
EXAMPLES_BUILD="examples/c++/build-android-examples-${{ matrix.abi }}"
cmake -S examples/c++ -B "$EXAMPLES_BUILD" -G Ninja \
-DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \
-DANDROID_ABI=${{ matrix.abi }} \
-DANDROID_PLATFORM=android-${{ matrix.api }} \
-DANDROID_STL=c++_static \
-DCMAKE_BUILD_TYPE=Release \
-DHOST_BUILD_DIR="$BUILD_DIR" \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build "$EXAMPLES_BUILD" --parallel
# Reuse the shared-library directory from Step 4; push again in
# case Step 4 was skipped or the directory was cleaned.
DEVICE_LIB_DIR="/data/local/tmp/zvec_tests/lib"
adb shell "mkdir -p $DEVICE_LIB_DIR" 2>/dev/null || true
SO_COUNT=0
while IFS= read -r so_file; do
adb push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1
SO_COUNT=$((SO_COUNT + 1))
done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null)
echo "Pushed $SO_COUNT shared libraries to $DEVICE_LIB_DIR"
for example in ailego-example core-example db-example; do
if [ -f "$EXAMPLES_BUILD/$example" ]; then
echo "=== Running $example ==="
adb push "$EXAMPLES_BUILD/$example" "/data/local/tmp/$example" > /dev/null 2>&1
adb shell "chmod 755 /data/local/tmp/$example && cd /data/local/tmp && LD_LIBRARY_PATH=$DEVICE_LIB_DIR ./$example"
adb shell "rm -f /data/local/tmp/$example"
fi
done
# ── Cleanup ────────────────────────────────────────────────────────
- name: Stop emulator
if: always()
shell: bash
run: |
adb emu kill 2>/dev/null || true
sleep 2
if [ -n "$EMULATOR_PID" ]; then
kill "$EMULATOR_PID" 2>/dev/null || true
fi
.github/workflows/05-windows-build.yml
+157
View File
@@ -0,0 +1,157 @@
name: Windows Build
on:
workflow_call:
workflow_dispatch:
permissions:
contents: read
jobs:
# Windows build and test matrix
build-and-test-windows:
name: Build & Test (${{ matrix.platform }})
runs-on: ${{ matrix.platform }}
strategy:
fail-fast: false
matrix:
include:
- platform: windows-2022
- platform: windows-2025
env:
SCCACHE_GHA_ENABLED: "true"
steps:
- name: Show env info
run: |
Get-CimInstance -ClassName Win32_Processor
where cl
& "C:\Program Files (x86)\Microsoft Visual Studio\Installer\vswhere.exe" -all -products * -prerelease -format json
shell: powershell
- name: Checkout code
uses: actions/checkout@v7
with:
submodules: recursive
- name: Cleanup left marker files
run: |
git submodule foreach --recursive 'git reset --hard && git clean -ffdx'
shell: powershell
- name: Setup sccache
uses: mozilla-actions/sccache-action@v0.0.10
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
cache: 'pip'
cache-dependency-path: 'pyproject.toml'
- name: Set up MSVC environment
uses: ilammy/msvc-dev-cmd@v1.13.0
with:
arch: x64
- name: Set up environment variables
run: |
$nproc = (Get-CimInstance Win32_ComputerSystem).NumberOfLogicalProcessors
echo "NPROC=$nproc" >> $env:GITHUB_ENV
echo "Using $nproc parallel jobs for builds"
shell: powershell
- name: Install dependencies
run: |
python -m pip install --upgrade pip `
pybind11==3.0 `
cmake==3.30.0 `
ninja==1.11.1 `
pytest `
pytest-xdist `
scikit-build-core `
setuptools_scm
shell: powershell
- name: Build from source
run: |
cd "$env:GITHUB_WORKSPACE"
$env:CMAKE_GENERATOR = "Ninja"
$env:CMAKE_BUILD_PARALLEL_LEVEL = "$env:NPROC"
python -m pip install -v . `
--no-build-isolation `
--config-settings='cmake.define.BUILD_TOOLS=ON' `
--config-settings='cmake.define.ENABLE_WERROR=ON' `
--config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=sccache' `
--config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=sccache'
shell: powershell
- name: Show sccache statistics after pip install
if: always()
run: sccache --show-stats
shell: powershell
- name: Run C++ Tests
run: |
cd "$env:GITHUB_WORKSPACE\build"
cmake --build . --target unittest --config Release --parallel $env:NPROC
shell: powershell
- name: Show sccache statistics after tests
if: always()
run: sccache --show-stats
shell: powershell
- name: Run Python Tests
run: |
cd "$env:GITHUB_WORKSPACE"
python -m pytest python/tests/ --basetemp=./.pytest_tmp
shell: powershell
- name: Run C++ Examples
run: |
cd "$env:GITHUB_WORKSPACE\examples\c++"
mkdir build
cd build
cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Release `
-DCMAKE_C_COMPILER_LAUNCHER=sccache `
-DCMAKE_CXX_COMPILER_LAUNCHER=sccache
cmake --build . --config Release --parallel $env:NPROC
# Copy zvec DLLs next to the example executables so Windows can find them.
# CMake places DLLs in bin/ (RUNTIME output) and import libs in lib/.
$buildDir = "$env:GITHUB_WORKSPACE\build"
foreach ($dllName in @("zvec.dll", "zvec_core.dll", "zvec_ailego.dll")) {
$found = $false
foreach ($sub in @("$buildDir\bin", "$buildDir\bin\Release", "$buildDir\lib", "$buildDir\lib\Release")) {
$dllPath = Join-Path $sub $dllName
if (Test-Path $dllPath) {
Copy-Item $dllPath -Destination . -Force
Write-Host "Copied $dllName from $sub"
$found = $true
break
}
}
if (-not $found) {
Write-Host "WARNING: $dllName not found, searching recursively..."
$dll = Get-ChildItem -Path $buildDir -Filter $dllName -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
if ($dll) {
Copy-Item $dll.FullName -Destination . -Force
Write-Host "Copied $dllName from $($dll.DirectoryName)"
} else {
Write-Error "$dllName not found anywhere under $buildDir"
}
}
}
.\db-example.exe
.\core-example.exe
.\ailego-example.exe
shell: powershell
- name: Show sccache statistics
if: always()
run: sccache --show-stats
shell: powershell
.github/workflows/06-ios-build.yml
+163
View File
@@ -0,0 +1,163 @@
name: iOS Cross Build
on:
workflow_call:
workflow_dispatch:
permissions:
contents: read
jobs:
build-ios:
runs-on: macos-15
strategy:
fail-fast: false
matrix:
include:
- platform: SIMULATORARM64
arch: arm64
sdk: iphonesimulator
test_on_simulator: true
- platform: OS
arch: arm64
sdk: iphoneos
test_on_simulator: false
name: iOS (${{ matrix.platform }})
steps:
- name: Checkout
uses: actions/checkout@v7
with:
submodules: recursive
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: ios-${{ matrix.platform }}
max-size: 150M
- name: Cache host protoc build
uses: actions/cache@v6
with:
path: build_host
key: macos-host-protoc-${{ hashFiles('thirdparty/protobuf/**', 'CMakeLists.txt') }}
restore-keys: |
macos-host-protoc-
- name: Build host protoc
run: |
if [ ! -f "build_host/bin/protoc" ]; then
cmake -S . -B build_host -DCMAKE_BUILD_TYPE=Release -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build build_host --target protoc --parallel $(sysctl -n hw.ncpu)
else
echo "Using cached host protoc"
fi
- name: Configure and Build
run: |
git submodule foreach --recursive 'git stash --include-untracked' || true
SDK_PATH=$(xcrun --sdk ${{ matrix.sdk }} --show-sdk-path)
NPROC=$(sysctl -n hw.ncpu)
cmake -S . -B build_ios_${{ matrix.platform }} \
-DCMAKE_SYSTEM_NAME=iOS \
-DCMAKE_OSX_DEPLOYMENT_TARGET="13.0" \
-DCMAKE_OSX_ARCHITECTURES="${{ matrix.arch }}" \
-DCMAKE_OSX_SYSROOT="$SDK_PATH" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_PYTHON_BINDINGS=OFF \
-DBUILD_TOOLS=OFF \
-DENABLE_WERROR=ON \
-DCMAKE_INSTALL_PREFIX="./install" \
-DGLOBAL_CC_PROTOBUF_PROTOC="$GITHUB_WORKSPACE/build_host/bin/protoc" \
-DIOS=ON \
-DCMAKE_POLICY_VERSION_MINIMUM=3.5 \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
cmake --build build_ios_${{ matrix.platform }} --parallel $NPROC
- name: Build test targets
if: matrix.test_on_simulator
run: |
NPROC=$(sysctl -n hw.ncpu)
cmake --build build_ios_${{ matrix.platform }} --target unittest --parallel $NPROC
- name: Boot iOS Simulator
if: matrix.test_on_simulator
run: |
DEVICE_ID=$(xcrun simctl list devices available -j \
| python3 -c "
import json, sys
data = json.load(sys.stdin)
for runtime, devices in data['devices'].items():
if 'iOS' in runtime:
for d in devices:
if 'iPhone' in d['name'] and d['isAvailable']:
print(d['udid'])
sys.exit(0)
sys.exit(1)
")
echo "DEVICE_ID=$DEVICE_ID" >> $GITHUB_ENV
xcrun simctl boot "$DEVICE_ID"
echo "Booted simulator: $DEVICE_ID"
- name: Run all tests on simulator
if: matrix.test_on_simulator
run: |
FAILED_TESTS=""
PASSED=0
TOTAL=0
for APP in build_ios_${{ matrix.platform }}/bin/*_test.app; do
[ -d "$APP" ] || continue
TEST_NAME=$(basename "$APP" .app)
BUNDLE_ID="com.zvec.${TEST_NAME}"
TOTAL=$((TOTAL + 1))
echo "::group::Running ${TEST_NAME}"
xcrun simctl install "$DEVICE_ID" "$APP"
set +eo pipefail
for attempt in 1 2 3; do
xcrun simctl launch --console "$DEVICE_ID" "$BUNDLE_ID" 2>&1 | tee /tmp/${TEST_NAME}.log
LAUNCH_EXIT=${PIPESTATUS[0]}
if ! grep -q "unknown to FrontBoard" /tmp/${TEST_NAME}.log; then
break
fi
echo "::warning::Attempt ${attempt}/3: FrontBoard has not registered ${TEST_NAME} yet, retrying in 3s..."
sleep 3
done
set -eo pipefail
if grep -q '\[ FAILED \]' /tmp/${TEST_NAME}.log; then
echo "::error::${TEST_NAME} has failing tests"
FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}"
elif grep -q '\[ PASSED \]' /tmp/${TEST_NAME}.log; then
PASSED=$((PASSED + 1))
elif grep -qE 'Failed: 0$' /tmp/${TEST_NAME}.log; then
# c_api_test uses a custom test framework (not GTest)
PASSED=$((PASSED + 1))
elif [ "$LAUNCH_EXIT" -eq 0 ]; then
echo "::warning::${TEST_NAME} exited 0 but produced no recognisable test summary"
PASSED=$((PASSED + 1))
else
echo "::error::${TEST_NAME} exited ${LAUNCH_EXIT} with no test summary"
FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}"
fi
echo "::endgroup::"
done
echo "Test summary: ${PASSED}/${TOTAL} passed"
if [ -n "$FAILED_TESTS" ]; then
echo "::error::Failed tests:${FAILED_TESTS}"
exit 1
fi
- name: Shutdown Simulator
if: matrix.test_on_simulator && always()
run: |
xcrun simctl shutdown "$DEVICE_ID" || true
.github/workflows/07-linux-riscv-build.yml
+342
View File
@@ -0,0 +1,342 @@
name: Weekly Linux RISC-V Build
on:
schedule:
# Saturday 00:00 UTC+8 (Friday 16:00 UTC)
- cron: '0 16 * * 5'
workflow_dispatch:
permissions:
contents: read
env:
RISE_PYPI: https://gitlab.com/api/v4/projects/56254198/packages/pypi/simple
PIP_BREAK_SYSTEM_PACKAGES: 1
jobs:
build:
name: Build (linux-riscv64)
runs-on: ubuntu-24.04-riscv
steps:
- name: Checkout code
uses: actions/checkout@v7
with:
submodules: recursive
- name: Install build dependencies
run: |
sudo mkdir -p /var/lib/dpkg/updates
sudo mkdir -p /var/lib/apt/lists/
sudo mkdir -p /var/cache/apt/archives/
sudo touch /var/lib/dpkg/status
sudo apt-get purge -y byobu || true
sudo apt-get update -o Dpkg::Lock::Timeout=300
sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \
-o Dpkg::Options::="--force-confdef" \
-o Dpkg::Options::="--force-confold" \
-o Dpkg::Lock::Timeout=300 \
ccache python3-pybind11 pybind11-dev
shell: bash
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: linux-riscv64-${{ runner.os }}-gcc
max-size: 150M
- name: Build from source
run: |
cd "$GITHUB_WORKSPACE"
NPROC=$(nproc 2>/dev/null || echo 2)
echo "Using $NPROC parallel jobs for builds"
cmake -S . -B build \
-G "Unix Makefiles" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TOOLS=ON \
-DBUILD_PYTHON_BINDINGS=ON \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
make -C build -j"$NPROC"
shell: bash
- name: Archive entire workspace
run: |
cd "$GITHUB_WORKSPACE"
tar -cf linux-riscv64-workspace.tar .
shell: bash
- name: Upload workspace artifacts
uses: actions/upload-artifact@v7
with:
name: linux-riscv64-workspace
path: ${{ github.workspace }}/linux-riscv64-workspace.tar
if-no-files-found: error
cpp-tests:
name: C++ Tests
runs-on: ubuntu-24.04-riscv
needs: build
steps:
- name: Download workspace artifacts
uses: actions/download-artifact@v8
with:
name: linux-riscv64-workspace
path: ${{ github.workspace }}
- name: Extract workspace
run: |
cd "$GITHUB_WORKSPACE"
tar -xf linux-riscv64-workspace.tar
shell: bash
- name: Install test dependencies
run: |
sudo mkdir -p /var/lib/dpkg/updates
sudo mkdir -p /var/lib/apt/lists/
sudo mkdir -p /var/cache/apt/archives/
sudo touch /var/lib/dpkg/status
sudo apt-get purge -y byobu || true
sudo apt-get update -o Dpkg::Lock::Timeout=300
sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \
-o Dpkg::Options::="--force-confdef" \
-o Dpkg::Options::="--force-confold" \
-o Dpkg::Lock::Timeout=300 \
ccache python3-pybind11 pybind11-dev libgtest-dev liburing-dev
shell: bash
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: linux-riscv64-${{ runner.os }}-gcc
max-size: 150M
- name: Reconfigure build directory
run: |
cd "$GITHUB_WORKSPACE"
cmake -S . -B build \
-G "Unix Makefiles" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TOOLS=ON \
-DBUILD_PYTHON_BINDINGS=ON \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
shell: bash
- name: Run C++ Tests
run: |
cd "$GITHUB_WORKSPACE/build"
NPROC=$(nproc 2>/dev/null || echo 2)
make unittest -j"$NPROC"
shell: bash
python-tests:
name: Python Tests
runs-on: ubuntu-24.04-riscv
needs: build
steps:
- name: Select Python
run: |
if command -v python3 >/dev/null 2>&1; then
PYTHON_BIN=python3
elif command -v python >/dev/null 2>&1; then
PYTHON_BIN=python
else
echo "No local Python interpreter found on PATH"
exit 1
fi
"$PYTHON_BIN" --version
echo "PYTHON=$PYTHON_BIN" >> "$GITHUB_ENV"
shell: bash
- name: Download workspace artifacts
uses: actions/download-artifact@v8
with:
name: linux-riscv64-workspace
path: ${{ github.workspace }}
- name: Extract workspace
run: |
cd "$GITHUB_WORKSPACE"
tar -xf linux-riscv64-workspace.tar
echo "$($PYTHON -c 'import site; print(site.USER_BASE)')/bin" >> "$GITHUB_PATH"
shell: bash
- name: Cache pip
uses: actions/cache@v6
with:
path: ~/.cache/pip
key: linux-riscv64-${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}
restore-keys: |
linux-riscv64-${{ runner.os }}-pip-
- name: Install dependencies
run: |
sudo mkdir -p /var/lib/dpkg/updates
sudo mkdir -p /var/lib/apt/lists/
sudo mkdir -p /var/cache/apt/archives/
sudo touch /var/lib/dpkg/status
sudo apt-get purge -y byobu || true
sudo apt-get update -o Dpkg::Lock::Timeout=300
sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \
-o Dpkg::Options::="--force-confdef" \
-o Dpkg::Options::="--force-confold" \
-o Dpkg::Lock::Timeout=300 \
ccache libgtest-dev liburing-dev
$PYTHON -m pip install --upgrade pip
$PYTHON -m pip install numpy==2.2.2 cmake==3.30.0 ninja==1.11.1.1 --index-url "$RISE_PYPI"
$PYTHON -m pip install pybind11==3.0 pytest scikit-build-core setuptools_scm pytest-xdist
shell: bash
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: linux-riscv64-${{ runner.os }}-gcc
max-size: 150M
- name: Reconfigure build directory
run: |
cd "$GITHUB_WORKSPACE"
cmake -S . -B build \
-G "Unix Makefiles" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TOOLS=ON \
-DBUILD_PYTHON_BINDINGS=ON \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-Dpybind11_DIR="$($PYTHON -c 'import pybind11; print(pybind11.get_cmake_dir())')"
shell: bash
- name: Install from existing build directory
run: |
cd "$GITHUB_WORKSPACE"
NPROC=$(nproc 2>/dev/null || echo 2)
export SKBUILD_BUILD_DIR="$GITHUB_WORKSPACE/build"
CMAKE_GENERATOR="Unix Makefiles" \
CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \
$PYTHON -m pip install -v . \
--no-build-isolation \
--config-settings='cmake.define.BUILD_TOOLS="ON"' \
--config-settings='cmake.define.ENABLE_WERROR=ON' \
--config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \
--config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache'
shell: bash
- name: Run Python Tests
run: |
cd "$GITHUB_WORKSPACE"
$PYTHON -m pytest python/tests/
shell: bash
cpp-examples:
name: C++ Examples
runs-on: ubuntu-24.04-riscv
needs: build
steps:
- name: Download workspace artifacts
uses: actions/download-artifact@v8
with:
name: linux-riscv64-workspace
path: ${{ github.workspace }}
- name: Extract workspace
run: |
cd "$GITHUB_WORKSPACE"
tar -xf linux-riscv64-workspace.tar
shell: bash
- name: Install ccache
run: |
sudo mkdir -p /var/lib/dpkg/updates
sudo mkdir -p /var/lib/apt/lists/
sudo mkdir -p /var/cache/apt/archives/
sudo touch /var/lib/dpkg/status
sudo apt-get purge -y byobu || true
sudo apt-get update -o Dpkg::Lock::Timeout=300
sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \
-o Dpkg::Options::="--force-confdef" \
-o Dpkg::Options::="--force-confold" \
-o Dpkg::Lock::Timeout=300 \
ccache
shell: bash
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: linux-riscv64-${{ runner.os }}-gcc
max-size: 150M
- name: Run C++ Examples
run: |
cd "$GITHUB_WORKSPACE/examples/c++"
NPROC=$(nproc 2>/dev/null || echo 2)
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
make -j "$NPROC"
./db-example
./core-example
./ailego-example
shell: bash
c-examples:
name: C Examples
runs-on: ubuntu-24.04-riscv
needs: build
steps:
- name: Download workspace artifacts
uses: actions/download-artifact@v8
with:
name: linux-riscv64-workspace
path: ${{ github.workspace }}
- name: Extract workspace
run: |
cd "$GITHUB_WORKSPACE"
tar -xf linux-riscv64-workspace.tar
shell: bash
- name: Install ccache
run: |
sudo mkdir -p /var/lib/dpkg/updates
sudo mkdir -p /var/lib/apt/lists/
sudo mkdir -p /var/cache/apt/archives/
sudo touch /var/lib/dpkg/status
sudo apt-get purge -y byobu || true
sudo apt-get update -o Dpkg::Lock::Timeout=300
sudo DEBIAN_FRONTEND=noninteractive apt-get install -yq \
-o Dpkg::Options::="--force-confdef" \
-o Dpkg::Options::="--force-confold" \
-o Dpkg::Lock::Timeout=300 \
ccache
shell: bash
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: linux-riscv64-${{ runner.os }}-gcc
max-size: 150M
- name: Run C Examples
run: |
cd "$GITHUB_WORKSPACE/examples/c"
NPROC=$(nproc 2>/dev/null || echo 2)
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
make -j "$NPROC"
./c_api_basic_example
./c_api_collection_schema_example
./c_api_doc_example
./c_api_field_schema_example
./c_api_index_example
./c_api_optimized_example
shell: bash
.github/workflows/08-cmake-subproject-integration.yml
+102
View File
@@ -0,0 +1,102 @@
name: Nightly CMake Subproject Integration
on:
schedule:
# Run after the other nightly jobs to avoid starting every heavy job at once.
- cron: '30 16 * * *'
workflow_dispatch:
permissions:
contents: read
jobs:
subproject-integration:
name: Subproject Build & Example
runs-on: ubuntu-24.04
timeout-minutes: 90
steps:
- name: Checkout code
uses: actions/checkout@v7
with:
submodules: recursive
- name: Setup ccache
uses: hendrikmuhs/ccache-action@v1.2
with:
key: cmake-subproject-linux-x64
max-size: 150M
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.10'
cache: 'pip'
cache-dependency-path: 'pyproject.toml'
- name: Install system dependencies
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libaio-dev
shell: bash
- name: Set up environment variables
run: |
NPROC=$(nproc 2>/dev/null || echo 2)
echo "NPROC=$NPROC" >> "$GITHUB_ENV"
echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> "$GITHUB_PATH"
shell: bash
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install \
cmake==3.30.0 \
ninja==1.11.1
shell: bash
- name: Configure subproject integration
run: |
SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-nightly"
SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-nightly-build"
mkdir -p "$SMOKE_SOURCE" "$SMOKE_BUILD"
cp "$GITHUB_WORKSPACE/.github/cmake/subproject-integration/CMakeLists.txt" \
"$SMOKE_SOURCE/CMakeLists.txt"
CMAKE_ARGS=(
-DZVEC_SOURCE_DIR="$GITHUB_WORKSPACE"
-DCMAKE_BUILD_TYPE=Release
-DENABLE_WERROR=ON
-DUSE_OSS_MIRROR=ON
-DCMAKE_C_COMPILER_LAUNCHER=ccache
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
-DBUILD_TOOLS=OFF
-DBUILD_C_BINDINGS=OFF
-DBUILD_PYTHON_BINDINGS=OFF
)
cmake -S "$SMOKE_SOURCE" -B "$SMOKE_BUILD" -G Ninja \
"${CMAKE_ARGS[@]}"
echo "SMOKE_BUILD=$SMOKE_BUILD" >> "$GITHUB_ENV"
shell: bash
- name: Build subproject targets
run: |
cmake --build "$SMOKE_BUILD" --target \
core_knn_diskann \
zvec_embed_db_example \
zvec_embed_core_example \
zvec_embed_ailego_example \
--parallel "$NPROC"
shell: bash
- name: Run subproject examples
run: |
cd "$SMOKE_BUILD"
./bin/zvec_embed_db_example
./bin/zvec_embed_core_example
./bin/zvec_embed_ailego_example
shell: bash
.github/workflows/_build_wheel_job.yml
+104
View File
@@ -0,0 +1,104 @@
name: "(Reusable) Build, Publish and Smoke-test a Wheel"
on:
workflow_call:
inputs:
runner:
description: "GitHub Actions runner label"
required: true
type: string
pypi_repository_url:
description: "PyPI repository URL (empty string means official PyPI)"
required: false
type: string
default: ""
secrets:
PYPI_API_TOKEN:
required: true
jobs:
build_publish_test:
name: Build / publish / smoke-test on ${{ inputs.runner }}
runs-on: ${{ inputs.runner }}
permissions:
contents: read
steps:
- name: Checkout code
uses: actions/checkout@v7
with:
submodules: recursive
fetch-tags: true
fetch-depth: 0
- name: Set up Python (for cibuildwheel controller)
uses: actions/setup-python@v6
with:
python-version: '3.11'
- name: Install cibuildwheel
run: |
pip install --upgrade pip
pip install cibuildwheel==3.4.0
- name: Build wheels using cibuildwheel
run: |
python -m cibuildwheel --output-dir wheelhouse
# Save list of built wheels for publishing
ls wheelhouse/*.whl | tee $GITHUB_STEP_SUMMARY
echo "wheels=$(ls wheelhouse/*.whl | tr '\n' ' ')" >> $GITHUB_ENV
- name: Publish to PyPI
if: success() && github.event_name == 'workflow_dispatch'
env:
TWINE_USERNAME: __token__
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
TWINE_REPOSITORY_URL: ${{ inputs.pypi_repository_url }}
run: |
pip install twine
twine upload --skip-existing --verbose wheelhouse/*.whl
- name: Smoke test from PyPI
if: success() && github.event_name == 'workflow_dispatch'
shell: bash
env:
PYPI_REPOSITORY_URL: ${{ inputs.pypi_repository_url }}
run: |
# Extract version from wheel filename (e.g. zvec-0.2.1.dev24-cp311-...whl -> 0.2.1.dev24)
WHEEL_FILE=$(ls wheelhouse/zvec-*.whl | head -1)
ZVEC_VERSION=$(basename "$WHEEL_FILE" | sed 's/zvec-\([^-]*\)-.*/\1/')
# Build index-url flags: use TestPyPI when repository URL is set, otherwise official PyPI
if [ -n "$PYPI_REPOSITORY_URL" ]; then
INDEX_FLAGS="--index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/"
echo "Waiting for zvec==$ZVEC_VERSION to become available on TestPyPI..."
else
INDEX_FLAGS=""
echo "Waiting for zvec==$ZVEC_VERSION to become available on PyPI..."
fi
# Poll until the version is available (max 5 minutes)
FOUND=0
for i in $(seq 1 30); do
if pip install $INDEX_FLAGS --dry-run "zvec==$ZVEC_VERSION" > /dev/null 2>&1; then
echo "Version $ZVEC_VERSION is available."
FOUND=1
break
fi
echo "Attempt $i/30: not yet available, retrying in 10s..."
sleep 10
done
if [ "$FOUND" -eq 0 ]; then
echo "ERROR: Timed out (5 min) waiting for zvec==$ZVEC_VERSION on PyPI. Aborting smoke test."
exit 1
fi
# Create a clean venv and install
python -m venv test_env
source test_env/bin/activate
pip install --upgrade pip
pip install $INDEX_FLAGS "zvec==$ZVEC_VERSION"
pip install --upgrade pip
pip install $INDEX_FLAGS "zvec==$ZVEC_VERSION"
# Run a simple smoke test
python -c "import zvec; print('Import OK:', zvec.__version__)"
.github/workflows/auto-assign.yml
+25
View File
@@ -0,0 +1,25 @@
name: Auto Assign PR
on:
pull_request_target:
types: [opened, reopened, ready_for_review]
permissions:
contents: read
pull-requests: write
jobs:
auto-assign:
runs-on: ubuntu-latest
steps:
- name: Checkout for config
uses: actions/checkout@v7
with:
ref: ${{ github.event.pull_request.base.ref }}
sparse-checkout: .github/auto-assign-config.yml
sparse-checkout-cone-mode: false
- name: Auto Assign PR Author
uses: kentaro-m/auto-assign-action@v2.0.2
with:
configuration-path: '.github/auto-assign-config.yml'
.github/workflows/clang_tidy.yml
+197
View File
@@ -0,0 +1,197 @@
name: Clang-Tidy
on:
workflow_call:
workflow_dispatch:
permissions:
contents: read
jobs:
clang_tidy:
name: Clang-Tidy Checks
runs-on: ubuntu-24.04
steps:
- name: Shallow checkout (no submodules)
uses: actions/checkout@v7
with:
submodules: false
fetch-depth: 0
- name: Collect changed C/C++ files
id: changed_files
uses: tj-actions/changed-files@v47
with:
files: |
**/*.c
**/*.cc
**/*.cpp
**/*.cxx
**/*.h
**/*.hpp
files_ignore: |
thirdparty/**
build/**
src/db/sqlengine/antlr/gen/**
src/db/index/column/fts_column/gen/**
- name: No C/C++ files changed - skip
if: steps.changed_files.outputs.any_changed != 'true'
run: echo "No C/C++ files changed. Skipping clang-tidy."
- name: Checkout submodules
if: steps.changed_files.outputs.any_changed == 'true'
run: git submodule update --init --recursive --depth 1
- name: Install dependencies
if: steps.changed_files.outputs.any_changed == 'true'
run: |
sudo apt-get update
sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev libaio-dev
- name: Setup ccache
if: steps.changed_files.outputs.any_changed == 'true'
uses: hendrikmuhs/ccache-action@v1.2
with:
key: clang-tidy
max-size: 500M
- name: Configure CMake and export compile commands
if: steps.changed_files.outputs.any_changed == 'true'
run: |
cmake -S . -B build -G Ninja \
-DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
-DBUILD_TOOLS=ON \
-DCMAKE_C_COMPILER_LAUNCHER=ccache \
-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
- name: Filter changed files against compile_commands.json
id: tidy_files
if: steps.changed_files.outputs.any_changed == 'true'
run: |
python3 - <<'PY'
import json
import os
from pathlib import Path
raw = os.environ.get("ALL_CHANGED_FILES", "")
changed = [f for f in raw.split() if f]
compile_db_path = Path("build/compile_commands.json")
with compile_db_path.open("r", encoding="utf-8") as fh:
compile_db = json.load(fh)
compile_entries = set()
for entry in compile_db:
file_field = entry.get("file", "")
if file_field:
compile_entries.add(os.path.normpath(file_field))
cwd = Path.cwd().resolve()
selected = []
skipped = []
for rel_path in changed:
abs_path = os.path.normpath(str((cwd / rel_path).resolve()))
if abs_path in compile_entries:
selected.append(rel_path)
elif not Path(rel_path).is_file():
skipped.append(f"{rel_path} (file not found)")
else:
skipped.append(f"{rel_path} (not in compile_commands.json)")
github_output = os.environ["GITHUB_OUTPUT"]
with open(github_output, "a", encoding="utf-8") as out:
out.write(f"any_tidy_files={'true' if selected else 'false'}\n")
out.write("all_tidy_files<<TIDY_EOF\n")
out.write("\n".join(selected) + "\n")
out.write("TIDY_EOF\n")
out.write("skipped_files<<SKIP_EOF\n")
out.write("\n".join(skipped) + "\n")
out.write("SKIP_EOF\n")
PY
env:
ALL_CHANGED_FILES: ${{ steps.changed_files.outputs.all_changed_files }}
- name: Show skipped files
if: steps.changed_files.outputs.any_changed == 'true' && steps.tidy_files.outputs.skipped_files != ''
run: |
echo "=== Files skipped by clang-tidy ==="
printf '%s\n' "${{ steps.tidy_files.outputs.skipped_files }}"
- name: Compute submodule commits hash
id: submodule_hash
if: steps.changed_files.outputs.any_changed == 'true'
run: |
hash=$(git submodule status --recursive | awk '{print $1}' | sort | sha256sum | awk '{print $1}')
echo "value=${hash}" >> "$GITHUB_OUTPUT"
- name: Restore generated headers cache
id: cache_headers
if: steps.changed_files.outputs.any_changed == 'true'
uses: actions/cache/restore@v6
with:
path: |
build/external
build/thirdparty
build/src/db/proto
key: clang-tidy-headers-${{ runner.os }}-${{ hashFiles('thirdparty/**/*.cmake', 'thirdparty/**/CMakeLists.txt', 'thirdparty/**/*.patch', 'src/db/proto/*.proto') }}-${{ steps.submodule_hash.outputs.value }}
- name: Build generated headers only
if: steps.tidy_files.outputs.any_tidy_files == 'true' && steps.cache_headers.outputs.cache-hit != 'true'
run: |
ninja -C build clang_tidy_deps
- name: Save generated headers cache
if: steps.tidy_files.outputs.any_tidy_files == 'true' && steps.cache_headers.outputs.cache-hit != 'true'
uses: actions/cache/save@v6
with:
path: |
build/external
build/thirdparty
build/src/db/proto
key: clang-tidy-headers-${{ runner.os }}-${{ hashFiles('thirdparty/**/*.cmake', 'thirdparty/**/CMakeLists.txt', 'thirdparty/**/*.patch', 'src/db/proto/*.proto') }}-${{ steps.submodule_hash.outputs.value }}
- name: Run clang-tidy on changed files (parallel)
if: steps.tidy_files.outputs.any_tidy_files == 'true'
run: |
mapfile -t files_to_check <<'TIDY_EOF'
${{ steps.tidy_files.outputs.all_tidy_files }}
TIDY_EOF
log_dir=$(mktemp -d)
printf '%s\n' "${files_to_check[@]}" \
| grep -v '^\s*$' \
| xargs -I{} -P "$(nproc)" sh -c '
file="{}"
if [ -f "$file" ]; then
log="'"$log_dir"'/$$.log"
echo "$file" > "$log"
if clang-tidy -p build --quiet --warnings-as-errors="*" "$file" >> "$log" 2>&1; then
echo "PASS: $file"
rm -f "$log"
else
echo "FAIL: $file"
fi
fi'
failed=0
for f in "$log_dir"/*.log; do
[ -e "$f" ] || break
failed=1
src=$(head -1 "$f")
echo ""
echo "::group::clang-tidy errors: $src"
tail -n +2 "$f"
echo "::endgroup::"
done
rm -rf "$log_dir"
if [ "$failed" -eq 1 ]; then
echo "::error::clang-tidy found issues in one or more files"
exit 1
fi
- name: No files to analyse
if: steps.changed_files.outputs.any_changed == 'true' && steps.tidy_files.outputs.any_tidy_files != 'true'
run: echo "Changed C/C++ files not in compile_commands.json. Nothing to analyse."
.github/workflows/community-pr-handler.yml
+148
View File
@@ -0,0 +1,148 @@
name: Community PR Handler
on:
pull_request_target:
types: [opened, reopened, ready_for_review]
permissions:
contents: read
pull-requests: write
jobs:
handle-community-pr:
runs-on: ubuntu-latest
if: github.event.pull_request.draft == false
steps:
- name: Check if community contributor
id: check
uses: actions/github-script@v9
with:
script: |
const pr = context.payload.pull_request;
const prAuthor = pr.user.login;
const association = pr.author_association;
console.log('========== PR Info ==========');
console.log(`PR Number: #${pr.number}`);
console.log(`PR Title: ${pr.title}`);
console.log(`PR Author: ${prAuthor}`);
console.log(`Author Association: ${association}`);
console.log(`Base Branch: ${pr.base.ref}`);
console.log(`Head Branch: ${pr.head.ref}`);
console.log(`Head Repo: ${pr.head.repo?.full_name || 'same repo'}`);
// OWNER, MEMBER, COLLABORATOR are maintainers; others are community contributors
const maintainerRoles = ['OWNER', 'MEMBER', 'COLLABORATOR'];
const isCommunity = !maintainerRoles.includes(association);
console.log('========== Decision ==========');
console.log(`Maintainer Roles: ${maintainerRoles.join(', ')}`);
console.log(`Is Community Contributor: ${isCommunity}`);
if (!isCommunity) {
console.log('Skipping: PR author is a maintainer, no action needed.');
}
core.setOutput('is_community', isCommunity);
- name: Checkout base branch for CODEOWNERS
if: steps.check.outputs.is_community == 'true'
uses: actions/checkout@v7
with:
ref: ${{ github.event.pull_request.base.ref }}
sparse-checkout: .github/CODEOWNERS
sparse-checkout-cone-mode: false
- name: Assign reviewer as assignee
if: steps.check.outputs.is_community == 'true'
uses: actions/github-script@v9
with:
script: |
const fs = require('fs');
const prNumber = context.payload.pull_request.number;
console.log('========== Fetching Changed Files ==========');
// Get changed files
const { data: files } = await github.rest.pulls.listFiles({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: prNumber
});
const changedFiles = files.map(f => f.filename);
console.log(`Total changed files: ${changedFiles.length}`);
changedFiles.forEach((f, i) => console.log(` [${i + 1}] ${f}`));
console.log('========== Parsing CODEOWNERS ==========');
// Parse CODEOWNERS
let codeowners = [];
try {
const content = fs.readFileSync('.github/CODEOWNERS', 'utf8');
const lines = content.split('\n').filter(l => l.trim() && !l.trim().startsWith('#'));
console.log(`Found ${lines.length} rules in CODEOWNERS`);
for (const line of lines) {
const parts = line.trim().split(/\s+/);
if (parts.length >= 2) {
const pattern = parts[0];
const owners = parts.slice(1).map(o => o.replace('@', ''));
codeowners.push({ pattern, owners });
console.log(` Rule: "${pattern}" -> [${owners.join(', ')}]`);
}
}
} catch (e) {
console.log('ERROR: Could not read CODEOWNERS:', e.message);
return;
}
console.log('========== Matching Files to Owners ==========');
// Find matching owners for changed files
const matchedOwners = new Set();
for (const file of changedFiles) {
let matchedOwner = null;
let matchedPattern = null;
for (const rule of codeowners) {
const pattern = rule.pattern;
if (pattern === '*') {
matchedOwner = rule.owners[0];
matchedPattern = pattern;
} else if (pattern.endsWith('/')) {
const dir = pattern.replace(/^\//, '').replace(/\/$/, '');
if (file.startsWith(dir + '/') || file.startsWith(dir)) {
matchedOwner = rule.owners[0];
matchedPattern = pattern;
}
} else if (file === pattern || file === pattern.replace(/^\//, '')) {
matchedOwner = rule.owners[0];
matchedPattern = pattern;
}
}
if (matchedOwner) {
matchedOwners.add(matchedOwner);
console.log(` "${file}" -> matched "${matchedPattern}" -> @${matchedOwner}`);
} else {
console.log(` "${file}" -> NO MATCH`);
}
}
console.log('========== Setting Assignees ==========');
// Set assignees
const assignees = Array.from(matchedOwners);
console.log(`Assignees to set: [${assignees.join(', ')}]`);
if (assignees.length > 0) {
try {
await github.rest.issues.addAssignees({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: prNumber,
assignees: assignees
});
console.log('SUCCESS: Assignees set successfully!');
} catch (e) {
console.log('ERROR: Failed to set assignees:', e.message);
console.log('Error details:', JSON.stringify(e, null, 2));
}
} else {
console.log('WARNING: No assignees to set - no matching owners found.');
}
console.log('========== Done ==========');
.github/workflows/continuous_bench.yml
+33
View File
@@ -0,0 +1,33 @@
name: Continuous Benchmark
on:
# Temporarily disabled. Restore original triggers below to re-enable.
# push:
# branches: [ "main"]
# paths-ignore:
# - '**.md'
# workflow_dispatch:
workflow_dispatch:
inputs:
__disabled__:
description: 'This workflow is temporarily disabled.'
required: false
concurrency:
group: cb-${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
benchmark:
if: false
runs-on: vdbbench
steps:
- uses: actions/checkout@v7
- name: Run VectorDBBench
env:
DATABASE_URL: ${{ secrets.DATABASE_URL }}
run: |
bash .github/workflows/scripts/run_vdb.sh
.github/workflows/docker/Dockerfile.linux_x64_glibc228
+83
View File
@@ -0,0 +1,83 @@
# =============================================================================
# Dockerfile.linux_x64_glibc228
# Purpose: Ubuntu 18.10 gcc-9 + glibc 2.28 + CMake 3.30.0 + PyBind11 build environment
# Warning: ubuntu:18.10 is EOL; use only for glibc 2.28 compatibility testing.
# =============================================================================
# Use official Ubuntu 18.10 (Cosmic Cuttlefish)
# glibc version: 2.28 (confirmed via `ldd --version`)
FROM ubuntu:18.10
# Replace Ubuntu mirror with old-releases.ubuntu.com for older glibc compatibility
RUN sed -i 's|http://\(.*\)/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list && \
sed -i 's|http://security.ubuntu.com/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list
# Add Ubuntu 20.04 (focal) repo for GCC 9 ONLY
RUN echo "deb http://archive.ubuntu.com/ubuntu/ focal main universe" >> /etc/apt/sources.list && \
echo "deb http://security.ubuntu.com/ubuntu/ focal-security main universe" >> /etc/apt/sources.list
# Prevent interactive prompts & set non-root user
ENV DEBIAN_FRONTEND=noninteractive \
TZ=Etc/UTC
# Create non-root user for safety (optional but recommended)
RUN useradd -m -u 1000 builder && \
mkdir -p /workspace && chown builder:builder /workspace
# Install base system dependencies
RUN apt-get update && \
apt-get install -y --no-install-recommends \
build-essential \
gcc-9 g++-9 \
ninja-build git curl ca-certificates vim wget lcov gnupg clang-format-18\
rsync lsb-release \
uuid-dev zlib1g-dev libssl-dev libffi-dev \
pybind11-dev && \
update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 90 \
--slave /usr/bin/g++ g++ /usr/bin/g++-9 && \
rm -rf /var/lib/apt/lists/*
# Install Miniforge (Conda) as root, then assign to builder
ENV MINIFORGE_VERSION="latest"
ENV MINIFORGE_HOME="/opt/miniforge3"
RUN curl -sSL "https://github.com/conda-forge/miniforge/releases/${MINIFORGE_VERSION}/download/Miniforge3-Linux-x86_64.sh" -o miniforge.sh && \
bash miniforge.sh -b -p ${MINIFORGE_HOME} && \
rm miniforge.sh && \
chown -R builder:builder ${MINIFORGE_HOME}
# Switch to non-root user
USER builder
ENV PATH="${MINIFORGE_HOME}/bin:${PATH}"
WORKDIR /workspace
# Create conda envs for supported Python versions
RUN conda create -n py310 python=3.10 -y && \
conda create -n py311 python=3.11 -y && \
conda create -n py312 python=3.12 -y
RUN conda clean --all -f -y
# Install CMake 3.30.0 from Kitware official binary
# Ref: https://github.com/Kitware/CMake/releases/tag/v3.30.0
RUN mkdir -p /tmp/cmake && cd /tmp/cmake && \
curl -sSL -o cmake.tar.gz \
"https://github.com/Kitware/CMake/releases/download/v3.30.0/cmake-3.30.0-linux-x86_64.tar.gz" && \
tar -xzf cmake.tar.gz --strip-components=1 -C /tmp/cmake && \
mkdir -p /home/builder/.local && \
mv * /home/builder/.local/ && \
chown -R builder:builder /home/builder/.local && \
rm -rf /tmp/cmake
# Add CMake to PATH
ENV PATH="/home/builder/.local/bin:${PATH}"
# Verify installations
RUN cmake --version && \
conda info && \
conda env list && \
python --version && \
gcc --version && \
ldd --version | head -n1
# Final setup
WORKDIR /workspace
.github/workflows/issue-auto-assign.yml
+99
View File
@@ -0,0 +1,99 @@
name: Issue Auto Assign
on:
issues:
types: [opened]
jobs:
auto-assign:
runs-on: ubuntu-latest
permissions:
issues: write
steps:
- name: Parse issue and assign
uses: actions/github-script@v9
with:
script: |
const issue = context.payload.issue;
const issueBody = issue.body || '';
const issueLabels = (issue.labels || []).map(l => l.name);
// Default assignees per template type (based on labels)
const defaultAssignees = {
'bug': 'zhourrr',
'feature': 'feihongxu0824',
'benchmark': 'egolearner',
'enhancement': 'feihongxu0824',
'integration': 'chinaux',
'profile': 'richyreachy'
};
// Global fallback assignee
const fallbackAssignee = 'feihongxu0824';
// Parse user-selected assignee from issue body
// The input field renders as: "### Preferred Assignee\n\n<entered_value>"
let selectedAssignee = null;
const assigneeMatch = issueBody.match(/### Preferred Assignee\s*\n+([^\n#]+)/);
if (assigneeMatch) {
const selection = assigneeMatch[1].trim();
console.log(`Parsed assignee input: "${selection}"`);
// If user entered a valid GitHub username (not empty, not placeholder text)
if (selection &&
selection !== '_No response_' &&
selection !== 'None' &&
!selection.toLowerCase().includes('leave empty') &&
!selection.startsWith('e.g.,')) {
// Clean up the username (remove @ if present)
selectedAssignee = selection.replace(/^@/, '').trim();
}
}
// Determine final assignee
let finalAssignee = selectedAssignee;
// If no user selection, use default based on label
if (!finalAssignee && issueLabels.length > 0) {
for (const [label, assignee] of Object.entries(defaultAssignees)) {
if (issueLabels.includes(label)) {
finalAssignee = assignee;
console.log(`Matched label "${label}" -> assignee "${assignee}"`);
break;
}
}
}
// Fallback to default assignee if no match
if (!finalAssignee) {
finalAssignee = fallbackAssignee;
console.log(`No match found, using fallback assignee: ${fallbackAssignee}`);
}
console.log(`Issue #${issue.number}: Labels = [${issueLabels.join(', ')}]`);
console.log(`User selected assignee: ${selectedAssignee || 'None (Auto)'}`);
console.log(`Final assignee: ${finalAssignee}`);
// Assign the issue
try {
await github.rest.issues.addAssignees({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue.number,
assignees: [finalAssignee]
});
console.log(`Successfully assigned issue #${issue.number} to ${finalAssignee}`);
} catch (error) {
console.error(`Failed to assign issue: ${error.message}`);
// If assignment fails (user may not have permission), add a comment
try {
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue.number,
body: `⚠️ Auto-assignment to \`${finalAssignee}\` failed. Please assign manually.`
});
} catch (commentError) {
console.error(`Failed to create comment: ${commentError.message}`);
}
}
.github/workflows/nightly_coverage.yml
+102
View File
@@ -0,0 +1,102 @@
name: Nightly Coverage Report
on:
schedule:
# Runs daily at 00:00 CST (China Standard Time) = 16:00 UTC
- cron: '0 16 * * *'
workflow_dispatch:
permissions:
contents: read
jobs:
coverage:
name: Nightly Coverage Report
runs-on: ubuntu-24.04
strategy:
matrix:
python-version: ['3.10']
fail-fast: false
steps:
- name: Checkout code
uses: actions/checkout@v7
with:
ref: main # Always use main for nightly
submodules: recursive
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: ${{ matrix.python-version }}
cache: 'pip'
cache-dependency-path: 'pyproject.toml'
- name: Set up environment variables
run: |
# Set number of processors for parallel builds
NPROC=$(nproc 2>/dev/null || echo 2)
echo "NPROC=$NPROC" >> $GITHUB_ENV
echo "Using $NPROC parallel jobs for builds"
# Add Python user base bin to PATH for pip-installed CLI tools
echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH
shell: bash
- name: Install system dependencies
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
lcov libaio-dev
shell: bash
- name: Install dependencies
run: |
python -m pip install --upgrade pip \
pybind11==3.0 \
cmake==3.30.0 \
ninja==1.11.1 \
pytest \
pytest-cov \
pytest-xdist \
scikit-build-core \
setuptools_scm
shell: bash
- name: Build with COVERAGE config
run: |
cd "$GITHUB_WORKSPACE"
CMAKE_GENERATOR="Unix Makefiles" \
CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \
python -m pip install -v . \
--no-build-isolation \
--config-settings="cmake.build-type=COVERAGE"
shell: bash
- name: Run Python Tests with Coverage
run: |
cd "$GITHUB_WORKSPACE"
python -m pytest python/tests/ --cov=zvec --cov-report=xml \
--deselect=python/tests/test_gil_release.py::TestGILRelease::test_gil_released_during_query
shell: bash
- name: Run C++ Tests and Generate Coverage
run: |
cd "$GITHUB_WORKSPACE/build"
make unittest -j$NPROC
cd "$GITHUB_WORKSPACE"
# Ensure gcov.sh is executable
chmod +x scripts/gcov.sh
bash scripts/gcov.sh -k
shell: bash
- name: Upload Coverage to Codecov
uses: codecov/codecov-action@v7
with:
files: ./proxima-zvec-filtered.lcov.info,./coverage.xml
flags: python,cpp,nightly
name: nightly-linux-py${{ matrix.python-version }}
token: ${{ secrets.CODECOV_TOKEN }}
.github/workflows/scripts/run_vdb.sh
+88
View File
@@ -0,0 +1,88 @@
set -e
QUANTIZE_TYPE_LIST="int8 int4 fp16 fp32"
CASE_TYPE_LIST="Performance768D1M Performance768D10M Performance1536D500K" # respectively test cosine, ip # Performance960D1M l2 metrics
LOG_FILE="bench.log"
DATE=$(date +%Y-%m-%d_%H-%M-%S)
NPROC=$(nproc 2>/dev/null || getconf _NPROCESSORS_ONLN 2>/dev/null || echo 2)
# COMMIT_ID = branch-date-sha
COMMIT_ID=${GITHUB_REF_NAME}-"$DATE"-$(echo ${GITHUB_WORKFLOW_SHA} | cut -c1-8)
COMMIT_ID=$(echo "$COMMIT_ID" | sed 's/\//_/g')
echo "COMMIT_ID: $COMMIT_ID"
echo "GITHUB_WORKFLOW_SHA: $GITHUB_WORKFLOW_SHA"
echo "workspace: $GITHUB_WORKSPACE"
DB_LABEL_PREFIX="Zvec16c64g-$COMMIT_ID"
# install zvec
git submodule update --init
# for debug
#cd ..
#export SKBUILD_BUILD_DIR="$GITHUB_WORKSPACE/../build"
pwd
python3 -m venv .venv
source .venv/bin/activate
pip install cmake ninja psycopg2-binary loguru fire
pip install -e /opt/VectorDBBench
CMAKE_GENERATOR="Unix Makefiles" \
CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \
pip install -v "$GITHUB_WORKSPACE"
for CASE_TYPE in $CASE_TYPE_LIST; do
echo "Running VectorDBBench for $CASE_TYPE"
DATASET_DESC=""
if [ "$CASE_TYPE" == "Performance768D1M" ]; then
DATASET_DESC="Performance768D1M - Cohere Cosine"
elif [ "$CASE_TYPE" == "Performance768D10M" ]; then
DATASET_DESC="Performance768D10M - Cohere Cosine"
else
DATASET_DESC="Performance1536D500K - OpenAI IP"
fi
for QUANTIZE_TYPE in $QUANTIZE_TYPE_LIST; do
DB_LABEL="$DB_LABEL_PREFIX-$CASE_TYPE-$QUANTIZE_TYPE"
echo "Running VectorDBBench for $DB_LABEL"
VDB_PARAMS="--path ${DB_LABEL} --db-label ${DB_LABEL} --case-type ${CASE_TYPE} --num-concurrency 12,14,16,18,20"
if [ "$CASE_TYPE" == "Performance768D1M" ]; then
VDB_PARAMS="${VDB_PARAMS} --m 15 --ef-search 180"
elif [ "$CASE_TYPE" == "Performance768D10M" ]; then
VDB_PARAMS="${VDB_PARAMS} --m 50 --ef-search 118 --is-using-refiner"
else #Performance1536D500K using default params + refiner to monitor performance degradation
VDB_PARAMS="${VDB_PARAMS} --m 50 --ef-search 100 --is-using-refiner"
fi
if [ "$QUANTIZE_TYPE" == "fp32" ]; then
vectordbbench zvec ${VDB_PARAMS} 2>&1 | tee $LOG_FILE
else
vectordbbench zvec ${VDB_PARAMS} --quantize-type "${QUANTIZE_TYPE}" 2>&1 | tee $LOG_FILE
fi
RESULT_JSON_PATH=$(grep -o "/opt/VectorDBBench/.*\.json" $LOG_FILE)
QPS=$(jq -r '.results[0].metrics.qps' "$RESULT_JSON_PATH")
RECALL=$(jq -r '.results[0].metrics.recall' "$RESULT_JSON_PATH")
LATENCY_P99=$(jq -r '.results[0].metrics.serial_latency_p99' "$RESULT_JSON_PATH")
LOAD_DURATION=$(jq -r '.results[0].metrics.load_duration' "$RESULT_JSON_PATH")
#quote the var to avoid space in the label
label_list="case_type=\"${CASE_TYPE}\",dataset_desc=\"${DATASET_DESC}\",db_label=\"${DB_LABEL}\",commit=\"${COMMIT_ID}\",date=\"${DATE}\",quantize_type=\"${QUANTIZE_TYPE}\""
# replace `/` with `_` in label_list
label_list=$(echo "$label_list" | sed 's/\//_/g')
cat <<EOF > prom_metrics.txt
# TYPE vdb_bench_qps gauge
vdb_bench_qps{$label_list} $QPS
# TYPE vdb_bench_recall gauge
vdb_bench_recall{$label_list} $RECALL
# TYPE vdb_bench_latency_p99 gauge
vdb_bench_latency_p99{$label_list} $LATENCY_P99
# TYPE vdb_bench_load_duration gauge
vdb_bench_load_duration{$label_list} $LOAD_DURATION
EOF
echo "prom_metrics:"
cat prom_metrics.txt
curl --data-binary @prom_metrics.txt "http://47.93.34.27:9091/metrics/job/benchmarks-${CASE_TYPE}/case_type/${CASE_TYPE}/quantize_type/${QUANTIZE_TYPE}" -v
done
done
.gitignore
+68
View File
@@ -0,0 +1,68 @@
.*
*~
bazel-*
build*
bin/*
lib/*
var/*
venv*
tests/integration/conf/*
tests/de_integration/conf/*
**/__pycache__/*
tests/bench/log/*
tests/integration/integration
tests/integration/log
tests/integration/*.log
tests/de_integration/log
tests/de_integration/*.log
!.git*
!.clang-format
!.clang-tidy
!.circleci
!.drone.yml
sdk/python/dist/
compile_commands.json
dist
html
*.lcov.info
# Compiled extension modules must never live in the source package: they are
# built by CMake and installed into the wheel. A stray copy left here gets
# bundled verbatim by scikit-build's wheel.packages and breaks the wheel
# (e.g. a stale binary with the wrong macOS deployment target).
python/zvec/*.so
python/zvec/*.dylib
python/zvec/*.pyd
# Dependencies
/node_modules
# Production
/build
# Generated files
.docusaurus
.cache-loader
# Misc
.DS_Store
.env.local
.env.development.local
.env.test.local
.env.production.local
npm-debug.log*
yarn-debug.log*
yarn-error.log*
allure-*
!build_android.sh
!build_ios.sh
# congfig
doc/
config/
examples/python/
examples/c_api/
logs/
.gitmodules
+61
View File
@@ -0,0 +1,61 @@
[submodule "thirdparty/googletest/googletest-1.10.0"]
path = thirdparty/googletest/googletest-1.10.0
url = https://github.com/google/googletest.git
ignore = dirty
[submodule "thirdparty/sparsehash/sparsehash-2.0.4"]
path = thirdparty/sparsehash/sparsehash-2.0.4
url = https://github.com/sparsehash/sparsehash.git
ignore = dirty
[submodule "thirdparty/gflags/gflags-2.2.2"]
path = thirdparty/gflags/gflags-2.2.2
url = https://github.com/gflags/gflags.git
[submodule "thirdparty/rocksdb/rocksdb-8.1.1"]
path = thirdparty/rocksdb/rocksdb-8.1.1
url = https://github.com/facebook/rocksdb.git
ignore = dirty
[submodule "thirdparty/yaml-cpp/yaml-cpp-0.6.3"]
path = thirdparty/yaml-cpp/yaml-cpp-0.6.3
url = https://github.com/jbeder/yaml-cpp.git
[submodule "thirdparty/arrow/apache-arrow-21.0.0"]
path = thirdparty/arrow/apache-arrow-21.0.0
url = https://github.com/apache/arrow.git
ignore = dirty
[submodule "thirdparty/glog/glog-0.5.0"]
path = thirdparty/glog/glog-0.5.0
url = https://github.com/google/glog.git
ignore = dirty
[submodule "thirdparty/protobuf/protobuf-3.21.12"]
path = thirdparty/protobuf/protobuf-3.21.12
url = https://github.com/protocolbuffers/protobuf.git
[submodule "thirdparty/lz4/lz4-1.9.4"]
path = thirdparty/lz4/lz4-1.9.4
url = https://github.com/lz4/lz4.git
[submodule "thirdparty/antlr/antlr4"]
path = thirdparty/antlr/antlr4
url = https://github.com/antlr/antlr4.git
ignore = dirty
[submodule "thirdparty/magic_enum/magic_enum-0.9.7"]
path = thirdparty/magic_enum/magic_enum-0.9.7
url = https://github.com/Neargye/magic_enum.git
ignore = dirty
[submodule "thirdparty/RaBitQ-Library/RaBitQ-Library-0.1"]
path = thirdparty/RaBitQ-Library/RaBitQ-Library-0.1
url = https://github.com/VectorDB-NTU/RaBitQ-Library.git
[submodule "thirdparty/aio/libaio-0.3"]
path = thirdparty/aio/libaio-0.3
url = https://github.com/yugabyte/libaio.git
[submodule "thirdparty/cppjieba/cppjieba-5.6.7"]
path = thirdparty/cppjieba/cppjieba-5.6.7
url = https://github.com/yanyiwu/cppjieba.git
[submodule "thirdparty/FastPFOR/FastPFOR-0.4.0"]
path = thirdparty/FastPFOR/FastPFOR-0.4.0
url = https://github.com/fast-pack/FastPFOR.git
[submodule "thirdparty/limonp/limonp-v1.0.2"]
path = thirdparty/limonp/limonp-v1.0.2
url = https://github.com/yanyiwu/limonp.git
[submodule "thirdparty/utf8proc/utf8proc-2.11.3"]
path = thirdparty/utf8proc/utf8proc-2.11.3
url = https://github.com/JuliaStrings/utf8proc.git
[submodule "thirdparty/snowball/snowball-3.1.1"]
path = thirdparty/snowball/snowball-3.1.1
url = https://github.com/snowballstem/snowball.git
CMakeLists.txt
+211
View File
@@ -0,0 +1,211 @@
cmake_minimum_required(VERSION 3.26)
cmake_policy(SET CMP0077 NEW)
project(zvec)
set(CC_CXX_STANDARD 17)
if(MSVC)
set(INTTYPES_FORMAT VC7)
# Use /Z7 (embed debug info in .obj) instead of /Zi (separate .pdb) to enable sccache/ccache.
# Use CACHE FORCE so it propagates into thirdparty subdirectories that call project().
foreach(_cfg DEBUG RELWITHDEBINFO)
foreach(_lang C CXX)
string(REPLACE "/Zi" "/Z7" _flags "${CMAKE_${_lang}_FLAGS_${_cfg}}")
set(CMAKE_${_lang}_FLAGS_${_cfg} "${_flags}" CACHE STRING "" FORCE)
endforeach()
endforeach()
add_compile_options(/EHsc) # def c++ exception behavior
add_compile_options(/Zc:preprocessor /Zc:__cplusplus)
add_compile_options(/we4716) # -Werror=return-type
###### reduce output length to make vibe coding work better :) should be removed or solved later
# TODO(windows): fix&cleanup
add_compile_options(/wd4267 /wd4244 /wd4305 /wd4242) # conversion/truncation from 'size_t' to 'type', possible loss of data
add_compile_options(/wd4310) # cast truncates constant value
add_compile_options(/wd4146) # unary minus operator applied to unsigned type #usage: uint32_t seg_id_{-1U};
add_compile_options(/wd4245) # conversion from 'type1' to 'type2', signed/unsigned mismatch
add_compile_options(/wd4251) # class 'type' needs to have dll-interface to be used by clients of class 'type2'
add_compile_options(/wd4099) # type name first seen using 'object_type1' now seen using 'object_type2'
add_compile_options(/wd4200) # nonstandard extension used : zero-sized array in struct/union
add_compile_options(/wd4324) # structure was padded due to alignment specifier
add_compile_options(/wd4702) # unreachable code
add_compile_options(/wd4996) # wstring_convert are deprecated in C++17 # usage: in antlr4
add_compile_options(/wd4101) # unreferenced local variable
add_compile_options(/wd4996) # The POSIX name for this item is deprecated. Instead, use the ISO C and C++ conformant name: _lseek
# result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?) # usage: keeps |= (1 << k);
add_compile_options(/wd4334)
######
else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Werror=return-type")
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Werror=return-type")
endif()
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT IOS)
set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,--no-as-needed")
set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--no-as-needed")
endif()
if(NOT DEFINED PROJECT_ROOT_DIR OR NOT PROJECT_ROOT_DIR)
set(PROJECT_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} CACHE PATH "Root directory of the project" FORCE)
endif()
message(STATUS "PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}")
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
# iOS platform detection
if(NOT ANDROID AND NOT IOS AND CMAKE_SYSTEM_NAME STREQUAL "iOS")
set(IOS TRUE)
endif()
# iOS bundle properties for test executables
if(IOS)
set(MACOSX_BUNDLE_BUNDLE_VERSION "1")
set(MACOSX_BUNDLE_SHORT_VERSION_STRING "1.0")
set(CMAKE_MACOSX_BUNDLE_INFO_PLIST "${PROJECT_ROOT_DIR}/cmake/iOSBundleInfo.plist.in")
endif()
if (NOT ANDROID AND NOT IOS AND AUTO_DETECT_ARCH AND HOST_ARCH MATCHES "^(x86|x64)$")
setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16)
message(STATUS "best compiler march, sse: " ${MATH_MARCH_FLAG_SSE} ", avx2: " ${MATH_MARCH_FLAG_AVX2} ", avx512: " ${MATH_MARCH_FLAG_AVX512} ", avx512fp16: " ${MATH_MARCH_FLAG_AVX512FP16})
endif()
include_directories(${PROJECT_ROOT_DIR}/src/include)
include_directories(${PROJECT_ROOT_DIR}/src)
option(BUILD_ZVEC_SHARED "Build all-in-one C++ shared library libzvec" ON)
option(BUILD_ZVEC_AILEGO_SHARED "Build all-in-one zvec-ailego shared library libzvec_ailego" ON)
option(BUILD_ZVEC_CORE_SHARED "Build all-in-one zvec-core shared library libzvec_core" ON)
option(BUILD_PYTHON_BINDINGS "Build Python bindings using pybind11" OFF)
option(BUILD_C_BINDINGS "Build C bindings" ON)
option(BUILD_TOOLS "Build tools" ON)
message(STATUS "BUILD_ZVEC_SHARED:${BUILD_ZVEC_SHARED}")
message(STATUS "BUILD_ZVEC_AILEGO_SHARED:${BUILD_ZVEC_AILEGO_SHARED}")
message(STATUS "BUILD_ZVEC_CORE_SHARED:${BUILD_ZVEC_CORE_SHARED}")
message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}")
message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}")
message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}")
option(RABITQ_ENABLE_AVX512 "Compile RaBitQ with AVX-512 support" OFF)
if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64" AND NOT ANDROID AND NOT IOS)
include(CheckCCompilerFlag)
check_c_compiler_flag("-mavx2" COMPILER_SUPPORTS_AVX2)
check_c_compiler_flag("-mavx512f -mavx512bw -mavx512vl" COMPILER_SUPPORTS_AVX512)
if(COMPILER_SUPPORTS_AVX2 OR COMPILER_SUPPORTS_AVX512)
set(RABITQ_SUPPORTED ON)
add_definitions(-DRABITQ_SUPPORTED=1)
if(RABITQ_ENABLE_AVX512 AND COMPILER_SUPPORTS_AVX512)
add_definitions(-DRABITQ_COMPILED_AVX512=1)
set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX512}")
else()
set(RABITQ_ARCH_FLAG "${MATH_MARCH_FLAG_AVX2}")
endif()
else()
set(RABITQ_SUPPORTED OFF)
add_definitions(-DRABITQ_SUPPORTED=0)
message(STATUS "RaBitQ support disabled - compiler does not support AVX2 or AVX-512")
endif()
elseif(IOS)
set(RABITQ_SUPPORTED OFF)
add_definitions(-DRABITQ_SUPPORTED=0)
message(STATUS "RaBitQ support disabled - not supported on iOS")
else()
set(RABITQ_SUPPORTED OFF)
add_definitions(-DRABITQ_SUPPORTED=0)
message(STATUS "RaBitQ support disabled - only supported on Linux x86_64")
endif()
message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}")
# DiskAnn support (Linux x86_64 only, requires libaio)
if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS)
set(DISKANN_SUPPORTED ON)
add_definitions(-DDISKANN_SUPPORTED=1)
else()
set(DISKANN_SUPPORTED OFF)
add_definitions(-DDISKANN_SUPPORTED=0)
message(STATUS "DiskAnn support disabled - only supported on Linux x86_64")
endif()
message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}")
option(USE_OSS_MIRROR "Use OSS mirror for faster third-party downloads" OFF)
if(DEFINED ENV{USE_OSS_MIRROR} AND NOT "$ENV{USE_OSS_MIRROR}" STREQUAL "")
set(USE_OSS_MIRROR "$ENV{USE_OSS_MIRROR}" CACHE BOOL "Use OSS mirror for faster third-party downloads" FORCE)
endif()
message(STATUS "USE_OSS_MIRROR:${USE_OSS_MIRROR}")
cc_directory(thirdparty)
cc_directories(src)
cc_directories(tests)
add_custom_target(clang_tidy_deps DEPENDS zvec_proto ARROW.BUILD glog gflags Lz4.BUILD)
if(BUILD_TOOLS)
cc_directories(tools)
endif()
git_version(GIT_SRCS_VER ${PROJECT_ROOT_DIR})
set(CPACK_PACKAGE_VERSION ${GIT_SRCS_VER})
set(CPACK_PACKAGE_NAME zvec)
include(CPack)
if(BUILD_PYTHON_BINDINGS)
if(APPLE)
set(CMAKE_STRIP "")
message(STATUS "Disabled strip on macOS to preserve code signature")
endif()
include(GNUInstallDirs)
if(DEFINED SKBUILD_PLATLIB_DIR)
set(ZVEC_PY_INSTALL_DIR "${SKBUILD_PLATLIB_DIR}")
elseif(DEFINED Python_SITEARCH)
set(ZVEC_PY_INSTALL_DIR "${Python_SITEARCH}")
else()
set(ZVEC_PY_INSTALL_DIR "${CMAKE_INSTALL_LIBDIR}")
endif()
message(STATUS "Zvec install path: ${ZVEC_PY_INSTALL_DIR}")
# COMPONENT python: only these runtime artifacts are pulled into the wheel
# (see install.components in pyproject.toml). The cc_library(PACKED) SDK
# install rules use the default component and are excluded from the wheel.
# Install the extension inside the zvec package (zvec/_zvec*.so) rather than
# at the site-packages root, so it does not pollute the top-level namespace.
# The Python code imports it as `zvec._zvec` accordingly.
install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec
COMPONENT python)
# DiskAnn ships as a runtime-loaded shared module
# (libzvec_diskann_plugin.so) that is brought online implicitly the
# first time a DiskAnn index is created — users never call any load
# function. The Python extension resolves the module next to _zvec.so
# (see the $ORIGIN rpath in src/binding/python/CMakeLists.txt); the
# module must therefore be installed alongside _zvec.so, i.e. inside the
# zvec package directory as well.
#
# Gate on DISKANN_SUPPORTED, not on the target's existence: on unsupported
# platforms (e.g. macOS / ARM64) the core_knn_diskann target is still
# defined, but built from an empty stub (src/core/algorithm/CMakeLists.txt)
# with zero exported symbols and a runtime load path compiled out
# (#if DISKANN_SUPPORTED). Shipping that stub is pure dead weight, so it is
# only packaged where DiskAnn is real — currently Linux x86_64 with libaio.
if(DISKANN_SUPPORTED)
install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec
COMPONENT python)
endif()
# Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works
# out of the box. python/zvec/__init__.py resolves this directory via
# importlib.resources and registers it with set_default_jieba_dict_dir().
set(ZVEC_JIEBA_DICT_SRC
"${PROJECT_SOURCE_DIR}/thirdparty/cppjieba/cppjieba-5.6.7/dict")
install(FILES
"${ZVEC_JIEBA_DICT_SRC}/jieba.dict.utf8"
"${ZVEC_JIEBA_DICT_SRC}/hmm_model.utf8"
DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec/data/jieba_dict
COMPONENT python)
endif()
CODE_OF_CONDUCT.md
+41
View File
@@ -0,0 +1,41 @@
# Zvec Code of Conduct
## Our Pledge
We pledge to foster an open, respectful, and harassment-free environment for everyone, regardless of age, body size, disability, ethnicity, gender identity and expression, level of experience, education, socioeconomic status, nationality, personal appearance, race, religion, or sexual identity and orientation.
## Expected Behavior
- Use welcoming and inclusive language
- Respect differing viewpoints and experiences
- Gracefully accept constructive criticism
- Focus on what is best for the community
- Show empathy and kindness toward others
## Unacceptable Behavior
- Harassment, intimidation, or discriminatory conduct
- Trolling, insulting, or derogatory comments
- Public or private harassment
- Publishing others’ private information without consent
- Any conduct that would reasonably be considered inappropriate in a professional setting
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team
at **zvec@alibaba-inc.com** (replace with your contact). All complaints will be reviewed and investigated promptly
and fairly.
The project team is obligated to respect the privacy and security of the reporter.
Consequences may include:
- A formal warning
- Temporary or permanent ban from project spaces
- Removal of contributions (e.g. comments, PRs)
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 2.1, available at
https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
[homepage]: https://www.contributor-covenant.org
CONTRIBUTING.md
+116
View File
@@ -0,0 +1,116 @@
# Contributing to Zvec
First off, thank you for considering contributing to Zvec! 🙌
Whether you're reporting a bug, proposing a feature, improving documentation, or submitting code — every contribution helps make Zvec better.
## Code of Conduct
By participating, you agree to abide by our [Code of Conduct](CODE_OF_CONDUCT.md). Please be respectful, collaborative, and inclusive.
---
## Development Setup
> [!TIP]
> **Linux** is the recommended environment for development and performance benchmarking.
### Prerequisites
- Python 3.10 - 3.14 (64-bit only)
- CMake ≥ 3.26, < 4.0 (`cmake --version`)
- A C++17-compatible compiler (e.g., `g++-11+`, `clang++`, Apple Clang on macOS)
### Clone & Initialize
```bash
git clone --recursive https://github.com/alibaba/zvec.git
cd zvec
```
> 💡 **Tip**
> - Forgot `--recursive`? Run:
> ```bash
> git submodule update --init --recursive
> ```
> - Set up pre-commit hooks:
> ```bash
> pip install pre-commit && pre-commit install
> ```
### Build from Source (Editable Install)
```bash
pip install -e ".[dev]"
# This installs dev dependencies (pytest, ruff, etc.) and builds the C++ extension in-place
```
> ✅ Verify:
> ```bash
> python -c "import zvec; print('Success!')"
> ```
---
## Testing
### Run All Tests
```bash
pytest python/tests/ -v
```
### Run with Coverage (Debug/CI)
```bash
pytest python/tests/ --cov=zvec --cov-report=term-missing
```
> 🔎 See full rules in `[tool.ruff]` section of `pyproject.toml`.
---
## Build Customization
You can control build behavior via environment variables or `pyproject.toml`:
| Option | How to Set | Description |
|--------|------------|-------------|
| **Build Type** | `CMAKE_BUILD_TYPE=Debug` | `Debug`, `Release`, or `Coverage` (for gcov/lcov) |
| **Generator** | `CMAKE_GENERATOR="Unix Makefiles"` | Default: `Ninja`; use Make if preferred |
| **AVX-512** | `ENABLE_SKYLAKE_AVX512=ON` | Enable AVX-512 optimizations (x86_64 only) |
Example (Debug + Make):
```bash
CMAKE_BUILD_TYPE=Debug CMAKE_GENERATOR="Unix Makefiles" pip install -v .
```
---
## Submitting Changes
1. Fork the repo and create a feature branch (`feat/...`, `fix/...`, `docs/...`)
2. Write clear commit messages (e.g., `fix(query): handle null vector in dense_fp32`)
3. Ensure tests pass & linter is clean
4. Open a Pull Request to `main`
5. Link related issue (e.g., `Closes #123`)
✅ **PRs should include**:
- Test coverage for new behavior
- Updates to documentation (if applicable)
- Reasoning behind non-obvious design choices
---
## Documentation
- User guides: `docs/` (built with MkDocs)
- API reference: generated from docstrings (follow [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings))
- Build & deploy: `mkdocs serve` / `mkdocs build`
---
## Need Help
- Browse [existing issues](https://github.com/alibaba/zvec/issues)
- For sensitive/security issues: email `zvec@alibaba-inc.com`
---
✨ Thanks again for being part of Zvec!
LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
NOTICE
+96
View File
@@ -0,0 +1,96 @@
zvec
Copyright 2025-present the zvec project
This product is licensed under the Apache License, Version 2.0 (see the LICENSE
file). It includes third-party software components that are distributed under
their own licenses, as listed below.
================================================================================
Third-Party Components
================================================================================
--------------------------------------------------------------------------------
Unicode Character Database
--------------------------------------------------------------------------------
Project: Unicode Character Database
Homepage: https://www.unicode.org/
License: Unicode License V3
Used in: src/db/index/column/fts_column/tokenizer/standard_tokenizer_unicode.inc
The generated standard tokenizer lookup tables are derived from Unicode 17.0.0
data files: auxiliary/WordBreakProperty.txt, emoji/emoji-data.txt,
LineBreak.txt, and Scripts.txt.
Unicode License V3 copyright and permission notice:
UNICODE LICENSE V3 COPYRIGHT AND PERMISSION NOTICE
Copyright © 1991-2026 Unicode, Inc.
NOTICE TO USER: Carefully read the following legal agreement. BY DOWNLOADING,
INSTALLING, COPYING OR OTHERWISE USING DATA FILES, AND/OR SOFTWARE, YOU
UNEQUIVOCALLY ACCEPT, AND AGREE TO BE BOUND BY, ALL OF THE TERMS AND
CONDITIONS OF THIS AGREEMENT. IF YOU DO NOT AGREE, DO NOT DOWNLOAD, INSTALL,
COPY, DISTRIBUTE OR USE THE DATA FILES OR SOFTWARE.
Permission is hereby granted, free of charge, to any person obtaining a copy
of data files and any associated documentation (the "Data Files") or software
and any associated documentation (the "Software") to deal in the Data Files
or Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, and/or sell copies of the Data
Files or Software, and to permit persons to whom the Data Files or Software
are furnished to do so, provided that either (a) this copyright and
permission notice appear with all copies of the Data Files or Software, or
(b) this copyright and permission notice appear in associated Documentation.
THE DATA FILES AND SOFTWARE ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT OF
THIRD PARTY RIGHTS.
IN NO EVENT SHALL THE COPYRIGHT HOLDER OR HOLDERS INCLUDED IN THIS NOTICE BE
LIABLE FOR ANY CLAIM, OR ANY SPECIAL INDIRECT OR CONSEQUENTIAL DAMAGES, OR
ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER
IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THE DATA FILES OR SOFTWARE.
Except as contained in this notice, the name of a copyright holder shall not
be used in advertising or otherwise to promote the sale, use or other
dealings in these Data Files or Software without prior written authorization
of the copyright holder.
--------------------------------------------------------------------------------
pyglass
--------------------------------------------------------------------------------
Project: pyglass — Graph Library for Approximate Similarity Search
Homepage: https://github.com/zilliztech/pyglass
License: MIT License
Used in: src/core/utility/linear_pool.h
The LinearPool implementation (and the accompanying Neighbor / Bitset helpers)
in src/core/utility/linear_pool.h is adapted from pyglass, with modifications
(a BlockHeap-compatible reset()/push_block() interface and the use of
MemoryHelper for huge-page-backed allocation). The related BlockHeap design in
src/core/utility/block_heap.{h,cc} is also derived from pyglass.
Original license text:
MIT License
Copyright (c) 2023 zh Wang
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+137
View File
@@ -0,0 +1,137 @@
<p align="right">
English | <a href="./README_CN.md">中文</a>
</p>
<div align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://zvec.oss-cn-hongkong.aliyuncs.com/logo/github_log_2.svg" />
<img src="https://zvec.oss-cn-hongkong.aliyuncs.com/logo/github_logo_1.svg" width="400" alt="zvec logo" />
</picture>
</div>
<p align="center">
<a href="https://codecov.io/github/alibaba/zvec"><img src="https://codecov.io/github/alibaba/zvec/graph/badge.svg?token=O81CT45B66" alt="Code Coverage"/></a>
<a href="https://github.com/alibaba/zvec/actions/workflows/01-ci-pipeline.yml"><img src="https://github.com/alibaba/zvec/actions/workflows/01-ci-pipeline.yml/badge.svg?branch=main" alt="Main"/></a>
<a href="https://github.com/alibaba/zvec/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-Apache%202.0-blue.svg" alt="License"/></a>
<a href="https://pypi.org/project/zvec/"><img src="https://img.shields.io/pypi/v/zvec.svg" alt="PyPI Release"/></a>
<a href="https://pypi.org/project/zvec/"><img src="https://img.shields.io/badge/python-3.10%20~%203.14-blue.svg" alt="Python Versions"/></a>
<a href="https://www.npmjs.com/package/@zvec/zvec"><img src="https://img.shields.io/npm/v/@zvec/zvec.svg" alt="npm Release"/></a>
</p>
<p align="center">
<a href="https://trendshift.io/repositories/20830" target="_blank"><img src="https://trendshift.io/api/badge/repositories/20830" alt="alibaba%2Fzvec | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</p>
<p align="center">
<a href="https://zvec.org/en/docs/db/quickstart/">🚀 <strong>Quickstart</strong> </a> |
<a href="https://zvec.org/en/">🏠 <strong>Home</strong> </a> |
<a href="https://zvec.org/en/docs/db/">📚 <strong>Docs</strong> </a> |
<a href="https://zvec.org/en/docs/db/benchmarks/">📊 <strong>Benchmarks</strong> </a> |
<a href="https://deepwiki.com/alibaba/zvec">🔎 <strong>DeepWiki</strong> </a> |
<a href="https://discord.gg/rKddFBBu9z">🎮 <strong>Discord</strong> </a> |
<a href="https://x.com/ZvecAI">🐦 <strong>X (Twitter)</strong> </a>
</p>
**Zvec** is an open-source, in-process vector database — lightweight, lightning-fast, and designed to embed directly into applications. Battle-tested within Alibaba Group, it delivers production-grade, low-latency and scalable similarity search with minimal setup.
> [!Important]
> 🚀 **v0.5.0 (June 12, 2026)**
>
> - **Full-Text Search (FTS)**: Native full-text search — attach an FTS index to any string field and query it with natural-language or structured expressions, no external search engine required.
> - **Hybrid Retrieval**: Combine full-text and vector search in a single `MultiQuery` across dense vectors, sparse vectors, scalar filters, and text.
> - **DiskANN Index**: New on-disk index that keeps the bulk of the index on disk, drastically cutting memory usage for large-scale datasets.
> - **Ecosystem & Platforms**: New official [Go](https://github.com/zvec-ai/zvec-go) / [Rust](https://github.com/zvec-ai/zvec-rust) SDKs, the [Zvec Studio](https://github.com/zvec-ai/zvec-studio) visual tool, and RISC-V support.
>
> 👉 [Read the Release Notes](https://github.com/alibaba/zvec/releases/tag/v0.5.0) | [View Roadmap 📍](https://github.com/alibaba/zvec/issues/309)
## 💫 Features
- **Blazing Fast**: Searches billions of vectors in milliseconds.
- **Simple, Just Works**: [Install](#-installation) and start searching in seconds. Pure local, no servers, no config, no fuss.
- **Dense + Sparse Vectors**: Support dense and sparse embeddings, multi-vector queries, and a rich selection of [vector index types](https://zvec.org/en/docs/db/concepts/vector-index/#vector-index-types) that scale from memory to disk.
- **Full-Text Search (FTS)**: Native keyword-based full-text search — query string fields with natural-language or structured expressions.
- **Hybrid Search**: Fuse vector similarity, full-text search, and structured filters in a single query for precise results.
- **Durable Storage**: Write-ahead logging (WAL) guarantees persistence — data is never lost, even on process crash or power failure.
- **Concurrent Access**: Multiple processes can read the same collection simultaneously; writes are single-process exclusive.
- **Runs Anywhere**: As an in-process library, Zvec runs wherever your code runs — notebooks, servers, CLI tools, or even edge devices.
## 📦 Installation
Zvec offers official SDKs across multiple languages:
- **[Python](https://pypi.org/project/zvec/)**: `pip install zvec` (requires Python 3.10–3.14)
- **[Node.js](https://www.npmjs.com/package/@zvec/zvec)**: `npm install @zvec/zvec`
- **[Go](https://github.com/zvec-ai/zvec-go)**: High-performance Go bindings.
- **[Rust](https://github.com/zvec-ai/zvec-rust)**: High-performance Rust bindings.
- **[Dart/Flutter](https://pub.dev/packages/zvec)**: `flutter pub add zvec`
Prefer a visual tool? Try **[Zvec Studio](https://github.com/zvec-ai/zvec-studio)** to browse data and debug queries — no code required.
### ✅ Supported Platforms
- Linux (x86_64, ARM64)
- macOS (ARM64)
- Windows (x86_64)
### 🛠️ Building from Source
If you prefer to build Zvec from source, please check the [Building from Source](https://zvec.org/en/docs/db/build/) guide.
## ⚡ One-Minute Example
```python
import zvec
# Define collection schema
schema = zvec.CollectionSchema(
name="example",
vectors=zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4),
)
# Create collection
collection = zvec.create_and_open(path="./zvec_example", schema=schema)
# Insert documents
collection.insert([
zvec.Doc(id="doc_1", vectors={"embedding": [0.1, 0.2, 0.3, 0.4]}),
zvec.Doc(id="doc_2", vectors={"embedding": [0.2, 0.3, 0.4, 0.1]}),
])
# Search by vector similarity
results = collection.query(
zvec.Query(field_name="embedding", vector=[0.4, 0.3, 0.3, 0.1]),
topk=10
)
# Results: list of {'id': str, 'score': float, ...}, sorted by relevance
print(results)
```
## 📈 Performance at Scale
Zvec delivers exceptional speed and efficiency, making it ideal for demanding production workloads.
<img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qps_10M.svg" width="800" alt="Zvec Performance Benchmarks" />
For detailed benchmark methodology, configurations, and complete results, please see our [Benchmarks documentation](https://zvec.org/en/docs/db/benchmarks/).
## 🤝 Join Our Community
<div align="center">
<div align="center">
| 💬 DingTalk | 📱 WeChat | 🎮 Discord | X (Twitter) |
| :---: | :---: | :---: | :---: |
| <img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qrcode/dingding.png" width="150" alt="DingTalk QR Code"/> | <img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qrcode/wechat.png" width="150" alt="WeChat QR Code"/> | [![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/rKddFBBu9z) | [![X (formerly Twitter) Follow](https://img.shields.io/twitter/follow/ZvecAI)](<https://x.com/ZvecAI>) |
| Scan to join | Scan to join | Click to join | Click to follow |
</div>
</div>
## ❤️ Contributing
We welcome and appreciate contributions from the community! Whether you're fixing a bug, adding a feature, or improving documentation, your help makes Zvec better for everyone.
Check out our [Contributing Guide](./CONTRIBUTING.md) to get started!
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`alibaba/zvec`
- 原始仓库：https://github.com/alibaba/zvec
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
README_CN.md
+133
View File
@@ -0,0 +1,133 @@
<p align="right">
<a href="./README.md">English</a> | 中文
</p>
<div align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://zvec.oss-cn-hongkong.aliyuncs.com/logo/github_log_2.svg" />
<img src="https://zvec.oss-cn-hongkong.aliyuncs.com/logo/github_logo_1.svg" width="400" alt="zvec logo" />
</picture>
</div>
<p align="center">
<a href="https://codecov.io/github/alibaba/zvec"><img src="https://codecov.io/github/alibaba/zvec/graph/badge.svg?token=O81CT45B66" alt="代码覆盖率"/></a>
<a href="https://github.com/alibaba/zvec/actions/workflows/01-ci-pipeline.yml"><img src="https://github.com/alibaba/zvec/actions/workflows/01-ci-pipeline.yml/badge.svg?branch=main" alt="Main"/></a>
<a href="https://github.com/alibaba/zvec/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-Apache%202.0-blue.svg" alt="许可证"/></a>
<a href="https://pypi.org/project/zvec/"><img src="https://img.shields.io/pypi/v/zvec.svg" alt="PyPI 版本"/></a>
<a href="https://pypi.org/project/zvec/"><img src="https://img.shields.io/badge/python-3.10%20~%203.14-blue.svg" alt="Python 版本"/></a>
<a href="https://www.npmjs.com/package/@zvec/zvec"><img src="https://img.shields.io/npm/v/@zvec/zvec.svg" alt="npm 版本"/></a>
</p>
<p align="center">
<a href="https://trendshift.io/repositories/20830" target="_blank"><img src="https://trendshift.io/api/badge/repositories/20830" alt="alibaba%2Fzvec | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</p>
<p align="center">
<a href="https://zvec.org/zh/docs/db/quickstart/">🚀 <strong>快速开始</strong> </a> |
<a href="https://zvec.org/zh/">🏠 <strong>主页</strong> </a> |
<a href="https://zvec.org/zh/docs/db/">📚 <strong>文档</strong> </a> |
<a href="https://zvec.org/zh/docs/db/benchmarks/">📊 <strong>性能报告</strong> </a> |
<a href="https://deepwiki.com/alibaba/zvec">🔎 <strong>DeepWiki</strong> </a> |
<a href="https://discord.gg/rKddFBBu9z">🎮 <strong>Discord</strong> </a> |
<a href="https://x.com/ZvecAI">🐦 <strong>X (Twitter)</strong> </a>
</p>
**Zvec** 是一款开源的嵌入式(进程内)向量数据库 — 轻量、极速，可直接嵌入应用程序。以极简的配置提供生产级、低延迟、可扩展的向量检索能力。
> [!IMPORTANT]
> 🚀 **v0.5.0（2026 年 6 月 12 日）**
>
> - **全文检索（FTS）**：原生全文检索能力——可为任意字符串字段挂载 FTS 索引，使用自然语言或结构化表达式检索，无需外接搜索引擎。
> - **混合检索**：在单次 `MultiQuery` 中融合全文与向量检索，跨稠密向量、稀疏向量、标量过滤与文本。
> - **DiskANN 索引**：全新磁盘索引，将索引主体存于磁盘，大幅降低大规模数据集的内存占用。
> - **生态与平台**：全新官方 [Go](https://github.com/zvec-ai/zvec-go) / [Rust](https://github.com/zvec-ai/zvec-rust) SDK、可视化工具 [Zvec Studio](https://github.com/zvec-ai/zvec-studio)，以及 RISC-V 架构支持。
>
> 👉 [查看更新日志](https://github.com/alibaba/zvec/releases/tag/v0.5.0) | [查看路线图 📍](https://github.com/alibaba/zvec/issues/309)
## 💫 核心特性
- **极致性能**：毫秒级响应，轻松检索数十亿级向量。
- **开箱即用**：[安装](#-安装)后即刻开始搜索，纯本地运行，无需服务器、无需配置、零门槛。
- **稠密 + 稀疏向量**：支持稠密向量、稀疏向量与多向量查询，以及从内存到磁盘、丰富多样的[向量索引类型](https://zvec.org/zh/docs/db/concepts/vector-index/#向量索引类型)。
- **全文检索（FTS）**：原生的基于关键词的全文检索——使用自然语言或结构化表达式检索字符串字段。
- **混合检索**：在单次查询中融合向量语义、全文检索与标量过滤，获得精确结果。
- **持久化存储**：WAL 预写日志保障数据持久性 — 即使进程崩溃或意外断电，数据也不会丢失。
- **并发访问**：支持多进程同时读取同一个 Collection；写入为单进程独占模式。
- **进程内运行**：无需单独部署服务，纯进程内运行。Notebook、高性能服务器、CLI 工具、边缘设备 — 随处可用。
## 📦 安装
Zvec 提供多语言官方 SDK：
- **[Python](https://pypi.org/project/zvec/)**：`pip install zvec`（需 64 位 Python 3.10–3.14）
- **[Node.js](https://www.npmjs.com/package/@zvec/zvec)**：`npm install @zvec/zvec`
- **[Go](https://github.com/zvec-ai/zvec-go)**：高性能的 Go 绑定。
- **[Rust](https://github.com/zvec-ai/zvec-rust)**：高性能的 Rust 绑定。
- **[Dart/Flutter](https://pub.dev/packages/zvec)**：`flutter pub add zvec`
想要图形界面？试试 **[Zvec Studio](https://github.com/zvec-ai/zvec-studio)**，零代码浏览数据与调试查询。
### ✅ 支持的平台
- Linux (x86_64, ARM64)
- macOS (ARM64)
- Windows (x86_64)
### 🛠️ 源码构建
如需从源码构建 Zvec，请参考[源码构建指南](https://zvec.org/zh/docs/db/build/)。
## ⚡ 一分钟上手
```python
import zvec
# 定义 collection schema
schema = zvec.CollectionSchema(
name="example",
vectors=zvec.VectorSchema("embedding", zvec.DataType.VECTOR_FP32, 4),
)
# 创建 collection
collection = zvec.create_and_open(path="./zvec_example", schema=schema)
# 插入 documents
collection.insert([
zvec.Doc(id="doc_1", vectors={"embedding": [0.1, 0.2, 0.3, 0.4]}),
zvec.Doc(id="doc_2", vectors={"embedding": [0.2, 0.3, 0.4, 0.1]}),
])
# 向量相似度检索
results = collection.query(
zvec.Query(field_name="embedding", vector=[0.4, 0.3, 0.3, 0.1]),
topk=10
)
# 查询结果：按相关性排序的 {'id': str, 'score': float, ...} 列表
print(results)
```
## 📈 极致性能
Zvec 提供极致的速度和效率，能够轻松应对高要求的生产环境负载。
<img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qps_10M.svg" width="800" alt="Zvec 性能基准测试" />
有关具体的测试方法、配置及完整结果，请参阅[性能报告](https://zvec.org/zh/docs/db/benchmarks/)。
## 🤝 加入社区
<div align="center">
| 💬 钉钉群 | 📱 微信群 | 🎮 Discord | X (Twitter) |
| :---: | :---: | :---: | :---: |
| <img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qrcode/dingding.png" width="150" alt="钉钉二维码"/> | <img src="https://zvec.oss-cn-hongkong.aliyuncs.com/qrcode/wechat.png" width="150" alt="微信二维码"/> | [![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/rKddFBBu9z) | [![X (formerly Twitter) Follow](https://img.shields.io/twitter/follow/ZvecAI)](<https://x.com/ZvecAI>) |
| 扫码加入 | 扫码加入 | 点击加入 | 点击关注 |
</div>
## ❤️ 参与贡献
非常欢迎来自社区的每一份贡献！无论是修复 Bug、新增功能，还是完善文档，都将让 Zvec 变得更好。
请查阅我们的[贡献指南](./CONTRIBUTING.md)开始参与！
cmake/bazel.cmake
+2521
View File
File diff suppressed because it is too large Load Diff
cmake/iOSBundleInfo.plist.in
+22
View File
@@ -0,0 +1,22 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple Computer//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>CFBundleDevelopmentRegion</key>
<string>English</string>
<key>CFBundleExecutable</key>
<string>${MACOSX_BUNDLE_EXECUTABLE_NAME}</string>
<key>CFBundleIdentifier</key>
<string>com.zvec.${MACOSX_BUNDLE_EXECUTABLE_NAME}</string>
<key>CFBundleInfoDictionaryVersion</key>
<string>6.0</string>
<key>CFBundleName</key>
<string>${MACOSX_BUNDLE_EXECUTABLE_NAME}</string>
<key>CFBundlePackageType</key>
<string>APPL</string>
<key>CFBundleShortVersionString</key>
<string>${MACOSX_BUNDLE_SHORT_VERSION_STRING}</string>
<key>CFBundleVersion</key>
<string>${MACOSX_BUNDLE_BUNDLE_VERSION}</string>
</dict>
</plist>
cmake/option.cmake
+321
View File
@@ -0,0 +1,321 @@
## https://en.wikipedia.org/wiki/List_of_Intel_CPU_microarchitectures
## https://en.wikipedia.org/wiki/List_of_AMD_CPU_microarchitectures
## https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html
## Intel Microarchitectures
option(ENABLE_NEHALEM "Enable Intel Nehalem CPU microarchitecture" OFF)
option(ENABLE_SANDYBRIDGE "Enable Intel Sandy Bridge CPU microarchitecture" OFF)
option(ENABLE_HASWELL "Enable Intel Haswell CPU microarchitecture" OFF)
option(ENABLE_BROADWELL "Enable Intel Broadwell CPU microarchitecture" OFF)
option(ENABLE_SKYLAKE "Enable Intel Skylake CPU microarchitecture" OFF)
option(ENABLE_SKYLAKE_AVX512 "Enable Intel Skylake Server CPU microarchitecture" OFF)
option(ENABLE_ICELAKE "Enable Intel Icelake CPU microarchitecture" OFF)
option(ENABLE_SAPPHIRERAPIDS "Enable Intel Sapphire Rapids Server CPU microarchitecture" OFF)
option(ENABLE_EMERALDRAPIDS "Enable Intel Emerald Rapids Server CPU microarchitecture" OFF)
option(ENABLE_GRANITERAPIDS "Enable Intel Granite Rapids Server CPU microarchitecture" OFF)
option(ENABLE_NATIVE "Enable native CPU microarchitecture" OFF)
## AMD Microarchitectures
option(ENABLE_ZEN1 "Enable AMD Zen+ Family 17h CPU microarchitecture" OFF)
option(ENABLE_ZEN2 "Enable AMD Zen 2 Family 17h CPU microarchitecture" OFF)
option(ENABLE_ZEN3 "Enable AMD Zen 3 Family 19h CPU microarchitecture" OFF)
## ARM architectures
option(ENABLE_ARMV8A "Enable ARMv8-a architecture" OFF)
option(ENABLE_ARMV8.1A "Enable ARMv8.1-a architecture" OFF)
option(ENABLE_ARMV8.2A "Enable ARMv8.2-a architecture" OFF)
option(ENABLE_ARMV8.3A "Enable ARMv8.3-a architecture" OFF)
option(ENABLE_ARMV8.4A "Enable ARMv8.4-a architecture" OFF)
option(ENABLE_ARMV8.5A "Enable ARMv8.5-a architecture" OFF)
option(ENABLE_ARMV8.6A "Enable ARMv8.6-a architecture" OFF)
## OpenMP option
option(ENABLE_OPENMP "Enable OpenMP support" OFF)
## Treat warnings as errors (for CI)
option(ENABLE_WERROR "Treat compiler warnings as errors in strict targets" OFF)
set(ARCH_OPTIONS
ENABLE_NEHALEM ENABLE_SANDYBRIDGE ENABLE_HASWELL ENABLE_BROADWELL ENABLE_SKYLAKE
ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS ENABLE_GRANITERAPIDS
ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3
ENABLE_ARMV8A ENABLE_ARMV8.1A ENABLE_ARMV8.2A ENABLE_ARMV8.3A ENABLE_ARMV8.4A
ENABLE_ARMV8.5A ENABLE_ARMV8.6A
ENABLE_NATIVE
)
option(AUTO_DETECT_ARCH "Auto detect CPU microarchitecture" ON)
foreach(opt IN LISTS ARCH_OPTIONS)
if(${opt})
message(WARNING "opt ${opt}")
set(AUTO_DETECT_ARCH OFF)
break()
endif()
endforeach()
include(CheckCCompilerFlag)
function(_AppendFlags _RESULT _FLAG)
if(${_RESULT} AND NOT "${${_RESULT}}" MATCHES "${_FLAG}")
set(${_RESULT} "${${_RESULT}} ${_FLAG}" PARENT_SCOPE)
else()
set(${_RESULT} "${_FLAG}" PARENT_SCOPE)
endif()
endfunction()
macro(add_arch_flag FLAG VAR_NAME OPTION_NAME)
if(MSVC)
if(${ARGC} GREATER 3)
set(_effective_flag "${ARGV3}")
else()
set(_effective_flag "")
endif()
else()
set(_effective_flag "${FLAG}")
endif()
if(NOT "${_effective_flag}" STREQUAL "")
check_c_compiler_flag("${_effective_flag}" COMPILER_SUPPORT_${VAR_NAME})
if(COMPILER_SUPPORT_${VAR_NAME})
_AppendFlags(CMAKE_C_FLAGS "${_effective_flag}")
_AppendFlags(CMAKE_CXX_FLAGS "${_effective_flag}")
set(${VAR_NAME}_ENABLED ON)
else()
if(${OPTION_NAME})
message(FATAL_ERROR "Compiler does not support required flag: '${_effective_flag}' for ${OPTION_NAME}")
else()
set(${VAR_NAME}_ENABLED OFF)
endif()
endif()
else()
set(${VAR_NAME}_ENABLED ON)
endif()
endmacro()
function(_setup_armv8_march)
if(MSVC)
return()
endif()
set(_arch "armv8")
check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch})
if(_COMP_SUPP_${_arch})
_AppendFlags(CMAKE_C_FLAGS "-march=${_arch}")
_AppendFlags(CMAKE_CXX_FLAGS "-march=${_arch}")
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}" PARENT_SCOPE)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}" PARENT_SCOPE)
return()
else()
message(WARNING "No ARMv8 march flag supported by compiler.")
endif()
endfunction()
function(_setup_x86_march)
if(MSVC)
return()
endif()
set(_arch "x86-64")
check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch})
if(_COMP_SUPP_${_arch})
_AppendFlags(CMAKE_C_FLAGS "-march=${_arch}")
_AppendFlags(CMAKE_CXX_FLAGS "-march=${_arch}")
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}" PARENT_SCOPE)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}" PARENT_SCOPE)
return()
else()
message(WARNING "No known x86 march flag supported; falling back to generic.")
endif()
endfunction()
function(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512 VAR_NAME_AVX512FP16)
if(MSVC)
#sse
set(${VAR_NAME_SSE} "" PARENT_SCOPE)
#avx2
check_c_compiler_flag("/arch:AVX2" _COMP_SUPP_MSVC_AVX2)
if(_COMP_SUPP_MSVC_AVX2)
set(${VAR_NAME_AVX2} "/arch:AVX2" PARENT_SCOPE)
else()
set(${VAR_NAME_AVX2} "" PARENT_SCOPE)
endif()
#avx512
check_c_compiler_flag("/arch:AVX512" _COMP_SUPP_MSVC_AVX512)
if(_COMP_SUPP_MSVC_AVX512)
set(${VAR_NAME_AVX512} "/arch:AVX512" PARENT_SCOPE)
elseif(_COMP_SUPP_MSVC_AVX2)
set(${VAR_NAME_AVX512} "/arch:AVX2" PARENT_SCOPE)
else()
set(${VAR_NAME_AVX512} "" PARENT_SCOPE)
endif()
#avx512fp16
if(_COMP_SUPP_MSVC_AVX512)
set(${VAR_NAME_AVX512FP16} "/arch:AVX512" PARENT_SCOPE)
elseif(_COMP_SUPP_MSVC_AVX2)
set(${VAR_NAME_AVX512FP16} "/arch:AVX2" PARENT_SCOPE)
else()
set(${VAR_NAME_AVX512FP16} "" PARENT_SCOPE)
endif()
return()
endif()
#sse
set(${VAR_NAME_SSE} "-march=corei7" PARENT_SCOPE)
#avx 2
set(${VAR_NAME_AVX2} "-march=core-avx2" PARENT_SCOPE)
#avx512
set(_x86_flags_avx512 "icelake-server" "skylake-avx512" "core-avx2" "x86-64")
foreach(_arch_avx512 IN LISTS _x86_flags_avx512)
check_c_compiler_flag("-march=${_arch_avx512}" _COMP_SUPP_${_arch_avx512})
if(_COMP_SUPP_${_arch_avx512})
set(${VAR_NAME_AVX512} "-march=${_arch_avx512}" PARENT_SCOPE)
break()
endif()
endforeach()
#avx512fp16
set(_x86_flags_avx512fp16
"sapphirerapids" "icelake-server" "skylake-avx512" "core-avx2" "x86-64"
)
foreach(_arch_avx512fp16 IN LISTS _x86_flags_avx512fp16)
check_c_compiler_flag("-march=${_arch_avx512fp16}" _COMP_SUPP_${_arch_avx512fp16})
if(_COMP_SUPP_${_arch_avx512fp16})
set(${VAR_NAME_AVX512FP16} "-march=${_arch_avx512fp16}" PARENT_SCOPE)
break()
endif()
endforeach()
endfunction()
# iOS: Skip -march flags and OpenMP; architecture is controlled by CMAKE_OSX_ARCHITECTURES
if(IOS OR CMAKE_SYSTEM_NAME STREQUAL "iOS")
return()
endif()
if(NOT AUTO_DETECT_ARCH)
if(ENABLE_NATIVE)
if (NOT MSVC)
add_arch_flag("-march=native" NATIVE ENABLE_NATIVE)
endif ()
endif()
if(ENABLE_ZEN3)
add_arch_flag("-march=znver3" ZNVER3 ENABLE_ZEN3 "/arch:AVX2")
endif()
if(ENABLE_ZEN2)
add_arch_flag("-march=znver2" ZNVER2 ENABLE_ZEN2 "/arch:AVX2")
endif()
if(ENABLE_ZEN1)
add_arch_flag("-march=znver1" ZNVER1 ENABLE_ZEN1 "/arch:AVX2")
endif()
if(ENABLE_GRANITERAPIDS)
add_arch_flag("-march=graniterapids" GRANITERAPIDS ENABLE_GRANITERAPIDS "/arch:AVX512")
endif()
if(ENABLE_EMERALDRAPIDS)
add_arch_flag("-march=emeraldrapids" EMERALDRAPIDS ENABLE_EMERALDRAPIDS "/arch:AVX512")
endif()
if(ENABLE_SAPPHIRERAPIDS)
add_arch_flag("-march=sapphirerapids" SAPPHIRERAPIDS ENABLE_SAPPHIRERAPIDS "/arch:AVX512")
endif()
if(ENABLE_ICELAKE)
add_arch_flag("-march=icelake-server" ICELAKE ENABLE_ICELAKE "/arch:AVX512")
endif()
if(ENABLE_SKYLAKE_AVX512)
add_arch_flag("-march=skylake-avx512" SKYLAKE_AVX512 ENABLE_SKYLAKE_AVX512 "/arch:AVX512")
endif()
if(ENABLE_SKYLAKE)
add_arch_flag("-march=skylake" SKYLAKE ENABLE_SKYLAKE "/arch:AVX2")
endif()
if(ENABLE_BROADWELL)
add_arch_flag("-march=broadwell" BROADWELL ENABLE_BROADWELL "/arch:AVX2")
endif()
if(ENABLE_HASWELL)
add_arch_flag("-march=haswell" HASWELL ENABLE_HASWELL "/arch:AVX2")
endif()
if(ENABLE_SANDYBRIDGE)
add_arch_flag("-march=sandybridge" SANDYBRIDGE ENABLE_SANDYBRIDGE "/arch:AVX")
endif()
if(ENABLE_NEHALEM)
add_arch_flag("-march=nehalem" NEHALEM ENABLE_NEHALEM)
endif()
# ARM (newest first — allow multiple? usually only one)
# But GCC allows only one -march=, so honor highest enabled
if(ENABLE_ARMV8.6A)
add_arch_flag("-march=armv8.6-a" ARMV86A ENABLE_ARMV8.6A)
endif()
if(ENABLE_ARMV8.5A)
add_arch_flag("-march=armv8.5-a" ARMV85A ENABLE_ARMV8.5A)
endif()
if(ENABLE_ARMV8.4A)
add_arch_flag("-march=armv8.4-a" ARMV84A ENABLE_ARMV8.4A)
endif()
if(ENABLE_ARMV8.3A)
add_arch_flag("-march=armv8.3-a" ARMV83A ENABLE_ARMV8.3A)
endif()
if(ENABLE_ARMV8.2A)
add_arch_flag("-march=armv8.2-a" ARMV82A ENABLE_ARMV8.2A)
endif()
if(ENABLE_ARMV8.1A)
add_arch_flag("-march=armv8.1-a" ARMV81A ENABLE_ARMV8.1A)
endif()
if(ENABLE_ARMV8A)
add_arch_flag("-march=armv8-a" ARMV8A ENABLE_ARMV8A)
endif()
else()
# AUTO DETECT
# Heuristic: detect host architecture and probe appropriate flags
string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" SYSTEM_PROC_LOWER)
if(SYSTEM_PROC_LOWER MATCHES "^(x86_64|amd64|x64)$")
set(HOST_ARCH x64)
elseif(SYSTEM_PROC_LOWER MATCHES "^(i[3-6]86|x86)$")
set(HOST_ARCH x86)
elseif(SYSTEM_PROC_LOWER MATCHES "^(aarch64|arm64)$")
set(HOST_ARCH arm64)
elseif(SYSTEM_PROC_LOWER MATCHES "^(arm|armv7|armv7-a|armv7l)$")
set(HOST_ARCH arm)
else()
set(HOST_ARCH unknown)
message(WARNING "unknown host arch: ${CMAKE_SYSTEM_PROCESSOR}")
endif()
# message(STATUS "host arch: ${HOST_ARCH}")
if (HOST_ARCH MATCHES "^(arm|arm64)$")
_setup_armv8_march()
elseif (HOST_ARCH MATCHES "^(x86|x64)$")
_setup_x86_march()
else ()
message(WARNING "unknown host arch - no -march set")
endif ()
endif()
# -----------------------------
# OpenMP
# -----------------------------
if(ENABLE_OPENMP)
find_package(OpenMP REQUIRED)
if(OpenMP_C_FLAGS)
_AppendFlags(CMAKE_C_FLAGS "${OpenMP_C_FLAGS}")
endif()
if(OpenMP_CXX_FLAGS)
_AppendFlags(CMAKE_CXX_FLAGS "${OpenMP_CXX_FLAGS}")
endif()
endif()
cmake/utils.cmake
+47
View File
@@ -0,0 +1,47 @@
# Promote a target's INTERFACE_INCLUDE_DIRECTORIES to be treated as SYSTEM
# includes by consumers, suppressing warnings from third-party headers.
function(mark_target_includes_system)
foreach(_target ${ARGN})
if(NOT TARGET ${_target})
continue()
endif()
get_target_property(_aliased ${_target} ALIASED_TARGET)
if(_aliased)
set(_target ${_aliased})
endif()
get_target_property(_inc ${_target} INTERFACE_INCLUDE_DIRECTORIES)
if(_inc)
set_target_properties(${_target} PROPERTIES
INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${_inc}")
endif()
endforeach()
endfunction()
function(apply_patch_once patch_name target_dir patch_file)
set(mark_file "${target_dir}/.${patch_name}_patched")
if(EXISTS "${mark_file}")
#message(STATUS "Patch '${patch_name}' already applied to ${target_dir}, skipping.")
return()
endif()
if(NOT EXISTS "${patch_file}")
message(FATAL_ERROR "Patch file '${patch_file}' not found!")
endif()
#message(STATUS "Applying patch '${patch_name}' to ${target_dir} ...")
execute_process(
COMMAND git apply --ignore-space-change --ignore-whitespace "${patch_file}"
WORKING_DIRECTORY "${target_dir}"
RESULT_VARIABLE patch_result
OUTPUT_VARIABLE patch_stdout
ERROR_VARIABLE patch_stderr
)
if(NOT patch_result EQUAL 0)
message(FATAL_ERROR "Failed to apply patch '${patch_name}' to ${target_dir}:\n${patch_stderr}")
else()
#message(STATUS "Patch '${patch_name}' applied successfully:\n${patch_stdout}")
file(WRITE "${mark_file}" "patched")
endif()
endfunction()
examples/c++/CMakeLists.txt
+128
View File
@@ -0,0 +1,128 @@
cmake_minimum_required(VERSION 3.13)
cmake_policy(SET CMP0077 NEW)
project(zvec-example-c++)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
# Enable compile_commands.json
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
# --- Paths to Zvec ---
# Allow custom host build directory, default to "build"
if(NOT DEFINED HOST_BUILD_DIR)
set(HOST_BUILD_DIR "build")
endif()
get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE)
set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include)
set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib)
# Add include and library search paths
include_directories(${ZVEC_INCLUDE_DIR})
set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR})
# Support multi-config builds (MSVC puts libs in Debug/Release subdirectories)
if(CMAKE_BUILD_TYPE)
set(ZVEC_CONFIG_LIB_DIR ${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE})
if(EXISTS "${ZVEC_CONFIG_LIB_DIR}")
list(APPEND ZVEC_LIB_SEARCH_DIRS ${ZVEC_CONFIG_LIB_DIR})
endif()
endif()
if(WIN32)
set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>")
endif()
function(zvec_find_shared_library OUT_VAR LIB_NAME)
unset(${OUT_VAR} CACHE)
if(WIN32)
find_library(${OUT_VAR}
NAMES ${LIB_NAME}_shared ${LIB_NAME}
PATHS ${ZVEC_LIB_SEARCH_DIRS}
NO_DEFAULT_PATH
NO_CMAKE_FIND_ROOT_PATH
)
else()
set(ZVEC_ORIGINAL_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES})
if(APPLE)
set(CMAKE_FIND_LIBRARY_SUFFIXES ".dylib")
else()
set(CMAKE_FIND_LIBRARY_SUFFIXES ".so")
endif()
find_library(${OUT_VAR}
NAMES ${LIB_NAME}
PATHS ${ZVEC_LIB_SEARCH_DIRS}
NO_DEFAULT_PATH
NO_CMAKE_FIND_ROOT_PATH
)
set(CMAKE_FIND_LIBRARY_SUFFIXES "${ZVEC_ORIGINAL_LIBRARY_SUFFIXES}")
endif()
set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE)
endfunction()
function(zvec_require_shared_library OUT_VAR LIB_NAME)
zvec_find_shared_library(${OUT_VAR} ${LIB_NAME})
if(NOT ${OUT_VAR})
message(FATAL_ERROR
"lib${LIB_NAME} shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. "
"Build zvec first, or pass -DHOST_BUILD_DIR=<build-dir>.")
endif()
set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE)
endfunction()
zvec_require_shared_library(ZVEC_SHARED_LIBRARY zvec)
zvec_require_shared_library(ZVEC_AILEGO_SHARED_LIBRARY zvec_ailego)
zvec_require_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core)
# --- Create INTERFACE target for libzvec (all-in-one C++ shared library) ---
# libzvec.so/.dylib/.dll already bundles all zvec internal components
# (zvec, zvec_core, zvec_ailego, zvec_turbo), so no individual dependency
# libraries need to be specified by the consumer.
add_library(zvec-lib INTERFACE)
target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}")
# --- Create INTERFACE target for libzvec_ailego (ailego-only all-in-one library) ---
# The ailego example intentionally depends only on libzvec_ailego.
add_library(zvec-ailego-lib INTERFACE)
target_link_libraries(zvec-ailego-lib INTERFACE "${ZVEC_AILEGO_SHARED_LIBRARY}")
# --- Create INTERFACE target for libzvec_core (core-only all-in-one library) ---
# The core example intentionally depends only on libzvec_core.
add_library(zvec-core-lib INTERFACE)
target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}")
# --- Executables ---
set(ZVEC_EXAMPLE_TARGETS)
add_executable(db-example db/main.cc)
target_link_libraries(db-example PRIVATE zvec-lib)
if(ANDROID)
target_link_libraries(db-example PRIVATE log)
endif()
list(APPEND ZVEC_EXAMPLE_TARGETS db-example)
add_executable(ailego-example ailego/main.cc)
target_link_libraries(ailego-example PRIVATE zvec-ailego-lib)
list(APPEND ZVEC_EXAMPLE_TARGETS ailego-example)
add_executable(core-example core/main.cc)
target_link_libraries(core-example PRIVATE zvec-core-lib)
list(APPEND ZVEC_EXAMPLE_TARGETS core-example)
add_executable(external-vector-example core/external_vector_example.cc)
target_link_libraries(external-vector-example PRIVATE zvec-core-lib)
list(APPEND ZVEC_EXAMPLE_TARGETS external-vector-example)
# Strip symbols to reduce executable size
if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
foreach(ZVEC_EXAMPLE_TARGET ${ZVEC_EXAMPLE_TARGETS})
add_custom_command(TARGET ${ZVEC_EXAMPLE_TARGET} POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:${ZVEC_EXAMPLE_TARGET}>"
COMMENT "Stripping symbols from ${ZVEC_EXAMPLE_TARGET}")
endforeach()
endif()
# Optimize for size
if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID AND ZVEC_EXAMPLE_TARGETS)
set_property(TARGET ${ZVEC_EXAMPLE_TARGETS} PROPERTY COMPILE_FLAGS "-Os")
set_property(TARGET ${ZVEC_EXAMPLE_TARGETS} PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
endif()
examples/c++/ailego/main.cc
+11
View File
@@ -0,0 +1,11 @@
#include <iostream>
#include <string>
#include <zvec/ailego/utility/string_helper.h>
using namespace zvec;
int main() {
std::string a{"hello world"};
std::cout << ailego::StringHelper::StartsWith(a, "hello") << std::endl;
}
examples/c++/core/external_vector_example.cc
+219
View File
@@ -0,0 +1,219 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/// @file external_vector_example.cc
/// @brief Demonstrates using HNSW index in external-vector mode.
///
/// In external-vector mode the index does NOT store raw vectors internally.
/// Instead, a user-provided VectorSource is passed on every Add/Search call
/// so the index can fetch vectors on demand. This is useful when vectors are
/// already stored elsewhere (e.g. a columnar store, mmap file, remote storage)
/// and you want to avoid duplicating them inside the index.
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <numeric>
#include <vector>
#include <zvec/core/interface/index.h>
#include <zvec/core/interface/index_factory.h>
#include <zvec/core/interface/index_param.h>
#include <zvec/core/interface/index_param_builders.h>
#include <zvec/core/interface/vector_source.h>
using namespace zvec::core_interface;
// ---------------------------------------------------------------------------
// A simple VectorSource backed by an in-memory float matrix.
// In production this could be backed by mmap, a database, or remote storage.
// ---------------------------------------------------------------------------
class InMemoryVectorSource : public zvec::core::VectorSource {
public:
/// @param base Pointer to a contiguous float array of shape [n, dim].
/// @param dim Dimensionality of each vector.
InMemoryVectorSource(const float *base, uint32_t dim)
: base_(base), dim_(dim) {}
const void *get_vector(uint32_t node_id) const override {
return base_ + static_cast<size_t>(node_id) * dim_;
}
private:
const float *base_;
uint32_t dim_;
};
// ---------------------------------------------------------------------------
// Helper: generate random float vectors in [0, 1)
// ---------------------------------------------------------------------------
static std::vector<float> generate_random_vectors(uint32_t count,
uint32_t dim) {
std::vector<float> data(static_cast<size_t>(count) * dim);
for (auto &v : data) {
v = static_cast<float>(std::rand()) / static_cast<float>(RAND_MAX);
}
return data;
}
// ---------------------------------------------------------------------------
// Helper: brute-force kNN (L2) for recall verification
// ---------------------------------------------------------------------------
static std::vector<uint32_t> brute_force_knn(const float *base, uint32_t n,
uint32_t dim, const float *query,
uint32_t topk) {
std::vector<std::pair<float, uint32_t>> dists(n);
for (uint32_t i = 0; i < n; ++i) {
float dist = 0.0f;
for (uint32_t d = 0; d < dim; ++d) {
float diff = base[static_cast<size_t>(i) * dim + d] - query[d];
dist += diff * diff;
}
dists[i] = {dist, i};
}
std::partial_sort(dists.begin(), dists.begin() + topk, dists.end());
std::vector<uint32_t> result(topk);
for (uint32_t i = 0; i < topk; ++i) {
result[i] = dists[i].second;
}
return result;
}
int main() {
constexpr uint32_t kDimension = 32;
constexpr uint32_t kDocCount = 200;
constexpr uint32_t kTopK = 5;
const std::string index_path = "external_vector_example.index";
// Clean up any previous run
std::filesystem::remove_all(index_path);
// ------ Step 1: Generate random vector data (simulating external storage)
std::srand(42);
auto vectors = generate_random_vectors(kDocCount, kDimension);
// Wrap data in our VectorSource
InMemoryVectorSource source(vectors.data(), kDimension);
// ------ Step 2: Build HNSW index with external-vector mode enabled
auto param = HNSWIndexParamBuilder()
.WithMetricType(MetricType::kL2sq)
.WithDataType(DataType::DT_FP32)
.WithDimension(kDimension)
.WithIsSparse(false)
.WithUseExternalVector(true) // <-- key setting
.Build();
auto index = IndexFactory::CreateAndInitIndex(*param);
if (!index) {
std::cerr << "Failed to create index." << std::endl;
return 1;
}
int ret = index->Open(
index_path, StorageOptions{StorageOptions::StorageType::kMMAP, true});
if (ret != 0) {
std::cerr << "Failed to open index." << std::endl;
return 1;
}
// ------ Step 3: Add vectors using AddWithSource
for (uint32_t i = 0; i < kDocCount; ++i) {
VectorData vd;
vd.vector =
DenseVector{vectors.data() + static_cast<size_t>(i) * kDimension};
ret = index->AddWithSource(vd, i, source);
if (ret != 0) {
std::cerr << "Failed to add doc " << i << std::endl;
return 1;
}
}
std::cout << "[OK] Added " << kDocCount << " vectors in external mode."
<< std::endl;
// ------ Step 4: Search using SearchWithSource
auto query_param =
HNSWQueryParamBuilder().with_topk(kTopK).with_ef_search(64).build();
// Use the first vector as query
const float *query_vec = vectors.data();
VectorData query;
query.vector = DenseVector{query_vec};
SearchResult result;
ret = index->SearchWithSource(query, query_param, source, &result);
if (ret != 0) {
std::cerr << "Search failed." << std::endl;
return 1;
}
std::cout << "[OK] Search returned " << result.doc_list_.size() << " results."
<< std::endl;
// The closest vector to vectors[0] should be itself (doc_id=0)
if (!result.doc_list_.empty() && result.doc_list_[0].key() == 0) {
std::cout << "[OK] Nearest neighbor is doc_id=0 (self), score="
<< result.doc_list_[0].score() << std::endl;
}
// ------ Step 5: Verify recall against brute-force
auto gt =
brute_force_knn(vectors.data(), kDocCount, kDimension, query_vec, kTopK);
uint32_t hits = 0;
for (const auto &doc : result.doc_list_) {
if (std::find(gt.begin(), gt.end(), static_cast<uint32_t>(doc.key())) !=
gt.end()) {
++hits;
}
}
float recall = static_cast<float>(hits) / static_cast<float>(kTopK);
std::cout << "[OK] Recall@" << kTopK << " = " << recall * 100.0f << "%"
<< std::endl;
// ------ Step 6: Reopen index and search again (persistence verification)
index->Close();
std::cout << "[OK] Index closed." << std::endl;
// Must re-create index instance with same params before reopening
index = IndexFactory::CreateAndInitIndex(*param);
if (!index) {
std::cerr << "Failed to re-create index for reopen." << std::endl;
return 1;
}
ret = index->Open(index_path,
StorageOptions{StorageOptions::StorageType::kMMAP, false});
if (ret != 0) {
std::cerr << "Failed to reopen index." << std::endl;
return 1;
}
SearchResult result2;
ret = index->SearchWithSource(query, query_param, source, &result2);
if (ret != 0) {
std::cerr << "Search after reopen failed." << std::endl;
return 1;
}
std::cout << "[OK] After reopen: search returned " << result2.doc_list_.size()
<< " results, top1 doc_id=" << result2.doc_list_[0].key()
<< std::endl;
// Cleanup
index->Close();
std::filesystem::remove_all(index_path);
std::cout << "\n=== External Vector Example Complete ===" << std::endl;
return 0;
}
examples/c++/core/main.cc
+88
View File
@@ -0,0 +1,88 @@
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <zvec/core/interface/index.h>
#include <zvec/core/interface/index_factory.h>
#include <zvec/core/interface/index_param.h>
#include <zvec/core/interface/index_param_builders.h>
using namespace zvec::core_interface;
constexpr uint32_t kDimension = 64;
const std::string index_name{"test.index"};
Index::Pointer create_index(const BaseIndexParam::Pointer &param,
int doc_num = 10) {
auto index = IndexFactory::CreateAndInitIndex(*param);
if (!index) {
std::cout << "Failed to create index." << std::endl;
return nullptr;
}
int ret = index->Open(
index_name, StorageOptions{StorageOptions::StorageType::kMMAP, true});
if (ret != 0) {
std::cout << "Failed to open index." << std::endl;
return nullptr;
}
for (int i = 0; i < doc_num; ++i) {
std::vector<float> vector(kDimension, i / 10.0f + 0.1f);
VectorData vector_data;
vector_data.vector = DenseVector{vector.data()};
ret = index->Add(vector_data, i);
if (ret != 0) {
std::cout << "Failed to add to index." << std::endl;
return nullptr;
}
}
ret = index->Train();
if (ret != 0) {
std::cout << "Failed to train index." << std::endl;
return nullptr;
}
return index;
}
int main() {
std::filesystem::remove(index_name);
auto param = HNSWIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(kDimension)
.WithIsSparse(false)
.Build();
auto index = create_index(param, 1);
std::cout << "index stats: " << index->GetDocCount() << std::endl;
// query
auto query_param = HNSWQueryParamBuilder()
.with_topk(10)
.with_fetch_vector(true)
.with_ef_search(20)
.build();
SearchResult result;
VectorData query;
std::vector<float> vector(kDimension, 0.1f);
query.vector = DenseVector{vector.data()};
int ret = index->Search(query, query_param, &result);
if (ret != 0) {
std::cout << "Failed to search index." << std::endl;
return -1;
}
std::cout << "query results: " << result.doc_list_.size() << std::endl;
if (result.doc_list_.size() == 0) {
std::cout << "No results found." << std::endl;
return -1;
}
std::cout << "key: " << result.doc_list_[0].key()
<< ", score: " << result.doc_list_[0].score() << std::endl;
return 0;
}
examples/c++/db/main.cc
+264
View File
@@ -0,0 +1,264 @@
#include <cstdlib>
#include <filesystem>
#include <string>
#include <vector>
#include <zvec/db/collection.h>
#include <zvec/db/doc.h>
#include <zvec/db/schema.h>
#include <zvec/db/status.h>
using namespace zvec;
Doc create_doc(const uint64_t doc_id, const CollectionSchema &schema,
std::string pk = "") {
Doc new_doc;
if (pk.empty()) {
pk = "pk_" + std::to_string(doc_id);
}
new_doc.set_pk(pk);
for (auto &field : schema.fields()) {
switch (field->data_type()) {
case DataType::BINARY: {
std::string binary_str("binary_" + std::to_string(doc_id));
new_doc.set<std::string>(field->name(), binary_str);
break;
}
case DataType::BOOL:
new_doc.set<bool>(field->name(), doc_id % 10 == 0);
break;
case DataType::INT32:
new_doc.set<int32_t>(field->name(), (int32_t)doc_id);
break;
case DataType::INT64:
new_doc.set<int64_t>(field->name(), (int64_t)doc_id);
break;
case DataType::UINT32:
new_doc.set<uint32_t>(field->name(), (uint32_t)doc_id);
break;
case DataType::UINT64:
new_doc.set<uint64_t>(field->name(), (uint64_t)doc_id);
break;
case DataType::FLOAT:
new_doc.set<float>(field->name(), (float)doc_id);
break;
case DataType::DOUBLE:
new_doc.set<double>(field->name(), (double)doc_id);
break;
case DataType::STRING:
new_doc.set<std::string>(field->name(),
"value_" + std::to_string(doc_id));
break;
case DataType::ARRAY_BINARY: {
std::vector<std::string> bin_vec;
for (size_t i = 0; i < (doc_id % 10); i++) {
bin_vec.push_back("bin_" + std::to_string(i));
}
new_doc.set<std::vector<std::string>>(field->name(), bin_vec);
break;
}
case DataType::ARRAY_BOOL:
new_doc.set<std::vector<bool>>(field->name(),
std::vector<bool>(10, doc_id % 10 == 0));
break;
case DataType::ARRAY_INT32:
new_doc.set<std::vector<int32_t>>(
field->name(), std::vector<int32_t>(10, (int32_t)doc_id));
break;
case DataType::ARRAY_INT64:
new_doc.set<std::vector<int64_t>>(
field->name(), std::vector<int64_t>(10, (int64_t)doc_id));
break;
case DataType::ARRAY_UINT32:
new_doc.set<std::vector<uint32_t>>(
field->name(), std::vector<uint32_t>(10, (uint32_t)doc_id));
break;
case DataType::ARRAY_UINT64:
new_doc.set<std::vector<uint64_t>>(
field->name(), std::vector<uint64_t>(10, (uint64_t)doc_id));
break;
case DataType::ARRAY_FLOAT:
new_doc.set<std::vector<float>>(field->name(),
std::vector<float>(10, (float)doc_id));
break;
case DataType::ARRAY_DOUBLE:
new_doc.set<std::vector<double>>(
field->name(), std::vector<double>(10, (double)doc_id));
break;
case DataType::ARRAY_STRING:
new_doc.set<std::vector<std::string>>(
field->name(),
std::vector<std::string>(10, "value_" + std::to_string(doc_id)));
break;
case DataType::VECTOR_BINARY32:
new_doc.set<std::vector<uint32_t>>(
field->name(),
std::vector<uint32_t>(field->dimension(), uint32_t(doc_id + 0.1)));
break;
case DataType::VECTOR_BINARY64:
new_doc.set<std::vector<uint64_t>>(
field->name(),
std::vector<uint64_t>(field->dimension(), uint64_t(doc_id + 0.1)));
break;
case DataType::VECTOR_FP32:
new_doc.set<std::vector<float>>(
field->name(),
std::vector<float>(field->dimension(), float(doc_id + 0.1)));
break;
case DataType::VECTOR_FP64:
new_doc.set<std::vector<double>>(
field->name(),
std::vector<double>(field->dimension(), double(doc_id + 0.1)));
break;
case DataType::VECTOR_FP16:
new_doc.set<std::vector<zvec::float16_t>>(
field->name(), std::vector<zvec::float16_t>(
field->dimension(), static_cast<zvec::float16_t>(
float(doc_id + 0.1))));
break;
case DataType::VECTOR_INT8:
new_doc.set<std::vector<int8_t>>(
field->name(),
std::vector<int8_t>(field->dimension(), (int8_t)doc_id));
break;
case DataType::VECTOR_INT16:
new_doc.set<std::vector<int16_t>>(
field->name(),
std::vector<int16_t>(field->dimension(), (int16_t)doc_id));
break;
case DataType::SPARSE_VECTOR_FP16: {
std::vector<uint32_t> indices;
std::vector<zvec::float16_t> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(zvec::float16_t(float(doc_id + 0.1)));
}
std::pair<std::vector<uint32_t>, std::vector<zvec::float16_t>>
sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<
std::pair<std::vector<uint32_t>, std::vector<zvec::float16_t>>>(
field->name(), sparse_float_vec);
break;
}
case DataType::SPARSE_VECTOR_FP32: {
std::vector<uint32_t> indices;
std::vector<float> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(float(doc_id + 0.1));
}
std::pair<std::vector<uint32_t>, std::vector<float>> sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
field->name(), sparse_float_vec);
break;
}
default:
std::cout << "Unsupported data type: " << field->name() << std::endl;
throw std::runtime_error("Unsupported vector data type");
}
}
return new_doc;
}
CollectionSchema::Ptr create_schema() {
auto schema = std::make_shared<CollectionSchema>("demo");
schema->set_max_doc_count_per_segment(1000);
schema->add_field(std::make_shared<FieldSchema>(
"id", DataType::INT64, false, std::make_shared<InvertIndexParams>(true)));
schema->add_field(std::make_shared<FieldSchema>(
"name", DataType::STRING, false,
std::make_shared<InvertIndexParams>(false)));
schema->add_field(
std::make_shared<FieldSchema>("weight", DataType::FLOAT, true));
schema->add_field(std::make_shared<FieldSchema>(
"dense", DataType::VECTOR_FP32, 128, false,
std::make_shared<HnswIndexParams>(MetricType::IP)));
schema->add_field(std::make_shared<FieldSchema>(
"sparse", DataType::SPARSE_VECTOR_FP32, 0, false,
std::make_shared<HnswIndexParams>(MetricType::IP)));
return schema;
}
int main() {
std::string path = "./demo";
std::filesystem::remove_all(path);
auto schema = create_schema();
CollectionOptions options{false, true};
auto result = Collection::CreateAndOpen(path, *schema, options);
if (!result.has_value()) {
std::cout << result.error().message() << std::endl;
return -1;
}
std::cout << "init stats: " << result.value()->Stats().value().to_string()
<< std::endl;
auto coll = std::move(result).value();
// insert docs
{
auto doc1 = create_doc(0, *schema);
std::vector<Doc> docs{doc1};
auto res = coll->Insert(docs);
if (!res.has_value()) {
std::cout << res.error().message() << std::endl;
return -1;
}
std::cout << "after insert stats " << coll->Stats().value().to_string()
<< std::endl;
}
// optimize
{
auto res = coll->Optimize();
if (!res.ok()) {
std::cout << res.message() << std::endl;
return -1;
}
std::cout << "after optimize stats " << coll->Stats().value().to_string()
<< std::endl;
}
// query
{
SearchQuery query;
query.topk_ = 10;
query.target_.field_name_ = "dense";
query.include_vector_ = true;
std::vector<float> query_vector = std::vector<float>(128, 0.1);
query.target_.set_vector(std::string((char *)query_vector.data(),
query_vector.size() * sizeof(float)));
auto res = coll->Query(query);
if (!res.has_value()) {
std::cout << res.error().message() << std::endl;
return -1;
}
std::cout << "query result: doc_count[" << res.value().size() << "]"
<< std::endl;
std::cout << "first doc: " << res.value()[0]->to_detail_string()
<< std::endl;
}
// close and reopen
coll.reset();
options.read_only_ = true;
result = Collection::Open(path, options);
if (!result.has_value()) {
std::cout << result.error().message() << std::endl;
return -1;
}
std::cout << "reopen stats: " << result.value()->Stats().value().to_string()
<< std::endl;
return 0;
}
examples/c/CMakeLists.txt
+137
View File
@@ -0,0 +1,137 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
cmake_minimum_required(VERSION 3.13)
cmake_policy(SET CMP0077 NEW)
project(zvec-example-c)
set(CMAKE_C_STANDARD 99)
set(CMAKE_C_STANDARD_REQUIRED ON)
# Enable compile_commands.json
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
# --- Paths to Zvec and dependencies ---
# Allow custom host build directory, default to "build"
if(NOT DEFINED HOST_BUILD_DIR)
set(HOST_BUILD_DIR "build")
endif()
set(ZVEC_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../src/include)
set(ZVEC_GENERATED_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/src/generated)
set(ZVEC_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/lib)
# Add include and library search paths
include_directories(${ZVEC_INCLUDE_DIR} ${ZVEC_GENERATED_INCLUDE_DIR})
link_directories(${ZVEC_LIB_DIR})
# Support multi-config builds (MSVC puts libs in Debug/Release subdirectories)
if(CMAKE_BUILD_TYPE)
link_directories(${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE})
endif()
# Find required packages
find_package(Threads REQUIRED)
# Create INTERFACE target for zvec_c_api (fat shared library)
# No whole-archive flags needed — all symbols are already resolved in the .so/.dylib
add_library(zvec-c-api INTERFACE)
target_link_libraries(zvec-c-api INTERFACE
zvec_c_api
Threads::Threads
)
if(NOT WIN32)
target_link_libraries(zvec-c-api INTERFACE ${CMAKE_DL_LIBS})
endif()
if(APPLE)
target_link_options(zvec-c-api INTERFACE -Wl,-rpath,${ZVEC_LIB_DIR})
endif()
# Basic example
add_executable(c_api_basic_example basic_example.c)
target_link_libraries(c_api_basic_example PRIVATE
zvec-c-api
)
# Schema example
add_executable(c_api_collection_schema_example collection_schema_example.c)
target_link_libraries(c_api_collection_schema_example PRIVATE
zvec-c-api
)
# Struct document example
add_executable(c_api_doc_example doc_example.c)
target_link_libraries(c_api_doc_example PRIVATE
zvec-c-api
)
# Index example
add_executable(c_api_index_example index_example.c)
target_link_libraries(c_api_index_example PRIVATE
zvec-c-api
)
# Field schema example
add_executable(c_api_field_schema_example field_schema_example.c)
target_link_libraries(c_api_field_schema_example PRIVATE
zvec-c-api
)
# Optimized example
add_executable(c_api_optimized_example optimized_example.c)
target_link_libraries(c_api_optimized_example PRIVATE
zvec-c-api
)
# DiskANN example
add_executable(c_api_diskann_example diskann_example.c)
target_link_libraries(c_api_diskann_example PRIVATE
zvec-c-api
)
# Strip symbols to reduce executable size
if(CMAKE_BUILD_TYPE STREQUAL "Release" AND (ANDROID OR (CMAKE_SYSTEM_NAME STREQUAL "Linux")))
add_custom_command(TARGET c_api_basic_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_basic_example>"
COMMENT "Stripping symbols from c_api_basic_example")
add_custom_command(TARGET c_api_collection_schema_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_collection_schema_example>"
COMMENT "Stripping symbols from c_api_collection_schema_example")
add_custom_command(TARGET c_api_doc_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_doc_example>"
COMMENT "Stripping symbols from c_api_doc_example")
add_custom_command(TARGET c_api_index_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_index_example>"
COMMENT "Stripping symbols from c_api_index_example")
add_custom_command(TARGET c_api_field_schema_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_field_schema_example>"
COMMENT "Stripping symbols from c_api_field_schema_example")
add_custom_command(TARGET c_api_optimized_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_optimized_example>"
COMMENT "Stripping symbols from c_api_optimized_example")
add_custom_command(TARGET c_api_diskann_example POST_BUILD
COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:c_api_diskann_example>"
COMMENT "Stripping symbols from c_api_diskann_example")
endif()
# Optimize for size
if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
set_property(TARGET c_api_basic_example c_api_collection_schema_example c_api_doc_example
c_api_index_example c_api_field_schema_example c_api_optimized_example
c_api_diskann_example
PROPERTY COMPILE_FLAGS "-Os")
set_property(TARGET c_api_basic_example c_api_collection_schema_example c_api_doc_example
c_api_index_example c_api_field_schema_example c_api_optimized_example
c_api_diskann_example
PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
endif()
examples/c/basic_example.c
+267
View File
@@ -0,0 +1,267 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Create a simple test collection using CollectionSchema
*/
static zvec_error_code_t create_simple_test_collection(
zvec_collection_t **collection) {
// Create collection schema using C API
zvec_collection_schema_t *schema =
zvec_collection_schema_create("test_collection");
if (!schema) {
return ZVEC_ERROR_INTERNAL_ERROR;
}
zvec_error_code_t error = ZVEC_OK;
// Create index parameters using new API
zvec_index_params_t *invert_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params) {
zvec_collection_schema_destroy(schema);
return ZVEC_ERROR_RESOURCE_EXHAUSTED;
}
zvec_index_params_set_invert_params(invert_params, true, false);
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params) {
zvec_index_params_destroy(invert_params);
zvec_collection_schema_destroy(schema);
return ZVEC_ERROR_RESOURCE_EXHAUSTED;
}
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_COSINE);
zvec_index_params_set_hnsw_params(hnsw_params, 16, 200);
// Create and add ID field (primary key)
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
zvec_field_schema_set_index_params(id_field, invert_params);
error = zvec_collection_schema_add_field(schema, id_field);
if (error != ZVEC_OK) {
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
zvec_collection_schema_destroy(schema);
return error;
}
// Create text field (inverted index)
zvec_field_schema_t *text_field =
zvec_field_schema_create("text", ZVEC_DATA_TYPE_STRING, true, 0);
zvec_field_schema_set_index_params(text_field, invert_params);
error = zvec_collection_schema_add_field(schema, text_field);
if (error != ZVEC_OK) {
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
zvec_collection_schema_destroy(schema);
return error;
}
// Create embedding field (HNSW index)
zvec_field_schema_t *embedding_field = zvec_field_schema_create(
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 3);
zvec_field_schema_set_index_params(embedding_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, embedding_field);
if (error != ZVEC_OK) {
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
zvec_collection_schema_destroy(schema);
return error;
}
// Cleanup index parameters (they have been copied to the field schemas)
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
// Use default options
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
zvec_collection_schema_destroy(schema);
return ZVEC_ERROR_RESOURCE_EXHAUSTED;
}
// Create collection using the new API
error = zvec_collection_create_and_open("./test_collection", schema, options,
collection);
// Cleanup resources
zvec_collection_options_destroy(options);
zvec_collection_schema_destroy(schema);
return error;
}
/**
* @brief Basic C API usage example
*/
int main() {
printf("=== ZVec C API Basic Example ===\n\n");
zvec_error_code_t error;
// Create collection using simplified function
zvec_collection_t *collection = NULL;
error = create_simple_test_collection(&collection);
if (handle_error(error, "creating collection") != ZVEC_OK) {
return 1;
}
printf("✓ Collection created successfully\n");
// Prepare test data
float vector1[] = {0.1f, 0.2f, 0.3f};
float vector2[] = {0.4f, 0.5f, 0.6f};
zvec_doc_t *docs[2];
for (int i = 0; i < 2; ++i) {
docs[i] = zvec_doc_create();
if (!docs[i]) {
fprintf(stderr, "Failed to create document %d\n", i);
// Cleanup allocated resources
for (int j = 0; j < i; ++j) {
zvec_doc_destroy(docs[j]);
}
return ZVEC_ERROR_INTERNAL_ERROR;
}
}
// Manually add fields to document 1
zvec_doc_set_pk(docs[0], "doc1");
zvec_doc_add_field_by_value(docs[0], "id", ZVEC_DATA_TYPE_STRING, "doc1",
strlen("doc1"));
zvec_doc_add_field_by_value(docs[0], "text", ZVEC_DATA_TYPE_STRING,
"First document", strlen("First document"));
zvec_doc_add_field_by_value(docs[0], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
vector1, 3 * sizeof(float));
// Manually add fields to document 2
zvec_doc_set_pk(docs[1], "doc2");
zvec_doc_add_field_by_value(docs[1], "id", ZVEC_DATA_TYPE_STRING, "doc2",
strlen("doc2"));
zvec_doc_add_field_by_value(docs[1], "text", ZVEC_DATA_TYPE_STRING,
"Second document", strlen("Second document"));
zvec_doc_add_field_by_value(docs[1], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
vector2, 3 * sizeof(float));
// Insert documents
size_t success_count = 0;
size_t error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2,
&success_count, &error_count);
if (handle_error(error, "inserting documents") != ZVEC_OK) {
zvec_collection_destroy(collection);
return 1;
}
printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count,
error_count);
for (int i = 0; i < 2; ++i) {
zvec_doc_destroy(docs[i]);
}
// Flush collection
error = zvec_collection_flush(collection);
if (handle_error(error, "flushing collection") != ZVEC_OK) {
printf("Collection flush failed\n");
} else {
printf("✓ Collection flushed successfully\n");
}
// Get collection statistics
zvec_collection_stats_t *stats = NULL;
error = zvec_collection_get_stats(collection, &stats);
if (handle_error(error, "getting collection stats") == ZVEC_OK) {
printf("✓ Collection stats - Document count: %llu\n",
(unsigned long long)zvec_collection_stats_get_doc_count(stats));
// Free statistics memory
zvec_collection_stats_destroy(stats);
}
printf("Testing vector query...\n");
// Query documents
zvec_vector_query_t *query = zvec_vector_query_create();
if (!query) {
fprintf(stderr, "Failed to create vector query\n");
zvec_collection_destroy(collection);
return 1;
}
zvec_vector_query_set_field_name(query, "embedding");
zvec_vector_query_set_query_vector(query, vector1, 3 * sizeof(float));
zvec_vector_query_set_topk(query, 10);
zvec_vector_query_set_filter(query, "");
zvec_vector_query_set_include_vector(query, true);
zvec_vector_query_set_include_doc_id(query, true);
zvec_doc_t **results = NULL;
size_t result_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&results, &result_count);
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
printf("[ERROR] Query failed: %s\n",
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
zvec_vector_query_destroy(query);
goto cleanup;
}
zvec_vector_query_destroy(query);
printf("✓ Query successful - Returned %zu results\n", result_count);
// Process query results
for (size_t i = 0; i < result_count && i < 5; ++i) {
const zvec_doc_t *doc = results[i];
const char *pk = zvec_doc_get_pk_copy(doc);
printf(" Result %zu: PK=%s, DocID=%llu, Score=%.4f\n", i + 1,
pk ? pk : "NULL", (unsigned long long)zvec_doc_get_doc_id(doc),
zvec_doc_get_score(doc));
if (pk) {
zvec_free((void *)pk);
}
}
// Free query results memory
zvec_docs_free(results, result_count);
cleanup:
// Cleanup resources
zvec_collection_destroy(collection);
printf("✓ Example completed\n");
return 0;
}
examples/c/collection_schema_example.c
+255
View File
@@ -0,0 +1,255 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Collection schema creation and management example
*/
int main() {
printf("=== ZVec Collection Schema Example ===\n\n");
zvec_error_code_t error;
// 1. Create collection schema
zvec_collection_schema_t *schema =
zvec_collection_schema_create("schema_example_collection");
if (!schema) {
fprintf(stderr, "Failed to create collection schema\n");
return 1;
}
printf("✓ Collection schema created successfully\n");
// 2. Set schema properties
zvec_collection_schema_set_max_doc_count_per_segment(schema, 1000000);
printf("✓ Set max documents per segment: %llu\n",
(unsigned long long)
zvec_collection_schema_get_max_doc_count_per_segment(schema));
// 3. Create index parameters
zvec_index_params_t *invert_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params) {
fprintf(stderr, "Failed to create invert index parameters\n");
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_index_params_set_invert_params(invert_params, true, false);
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params) {
fprintf(stderr, "Failed to create HNSW index parameters\n");
zvec_index_params_destroy(invert_params);
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params, 16, 200);
// 4. Create and add ID field (primary key)
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
if (!id_field) {
fprintf(stderr, "Failed to create ID field\n");
zvec_collection_schema_destroy(schema);
return 1;
}
error = zvec_collection_schema_add_field(schema, id_field);
if (handle_error(error, "adding ID field") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return 1;
}
printf("✓ ID field added successfully\n");
// 5. Create and add text field with inverted index
zvec_field_schema_t *text_field =
zvec_field_schema_create("content", ZVEC_DATA_TYPE_STRING, true, 0);
if (!text_field) {
fprintf(stderr, "Failed to create text field\n");
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_field_schema_set_index_params(text_field, invert_params);
error = zvec_collection_schema_add_field(schema, text_field);
if (handle_error(error, "adding text field") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return 1;
}
printf("✓ Text field with inverted index added successfully\n");
// 6. Create and add vector field with HNSW index
zvec_field_schema_t *vector_field = zvec_field_schema_create(
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128);
if (!vector_field) {
fprintf(stderr, "Failed to create vector field\n");
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_field_schema_set_index_params(vector_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, vector_field);
if (handle_error(error, "adding vector field") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return 1;
}
printf("✓ Vector field with HNSW index added successfully\n");
// 7. Check field count
// Note: This function may not exist in current API, commenting out for now
// size_t field_count = zvec_collection_schema_get_field_count(schema);
// printf("✓ Total field count: %zu\n", field_count);
// 8. Create collection with schema
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
fprintf(stderr, "Failed to create collection options\n");
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open("./schema_example_collection", schema,
options, &collection);
if (handle_error(error, "creating collection with schema") != ZVEC_OK) {
zvec_collection_options_destroy(options);
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_collection_options_destroy(options);
printf("✓ Collection created successfully with schema\n");
// 9. Prepare test data
float vector1[128];
float vector2[128];
for (int i = 0; i < 128; i++) {
vector1[i] = (float)(i + 1) / 128.0f;
vector2[i] = (float)(i + 2) / 128.0f;
}
// 10. Create documents
zvec_doc_t *docs[2];
for (int i = 0; i < 2; i++) {
docs[i] = zvec_doc_create();
if (!docs[i]) {
fprintf(stderr, "Failed to create document %d\n", i);
// Cleanup
for (int j = 0; j < i; j++) {
zvec_doc_destroy(docs[j]);
}
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
return 1;
}
}
// Add fields to document 1
zvec_doc_set_pk(docs[0], "doc1");
zvec_doc_add_field_by_value(docs[0], "id", ZVEC_DATA_TYPE_STRING, "doc1",
strlen("doc1"));
zvec_doc_add_field_by_value(docs[0], "content", ZVEC_DATA_TYPE_STRING,
"First test document",
strlen("First test document"));
zvec_doc_add_field_by_value(docs[0], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
vector1, 128 * sizeof(float));
// Add fields to document 2
zvec_doc_set_pk(docs[1], "doc2");
zvec_doc_add_field_by_value(docs[1], "id", ZVEC_DATA_TYPE_STRING, "doc2",
strlen("doc2"));
zvec_doc_add_field_by_value(docs[1], "content", ZVEC_DATA_TYPE_STRING,
"Second test document",
strlen("Second test document"));
zvec_doc_add_field_by_value(docs[1], "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
vector2, 128 * sizeof(float));
// 11. Insert documents
size_t success_count = 0, error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2,
&success_count, &error_count);
if (handle_error(error, "inserting documents") != ZVEC_OK) {
// Cleanup
for (int i = 0; i < 2; i++) {
zvec_doc_destroy(docs[i]);
}
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
return 1;
}
printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count,
error_count);
// Cleanup documents
for (int i = 0; i < 2; i++) {
zvec_doc_destroy(docs[i]);
}
// 12. Flush collection
error = zvec_collection_flush(collection);
if (handle_error(error, "flushing collection") == ZVEC_OK) {
printf("✓ Collection flushed successfully\n");
}
// 13. Query test
zvec_vector_query_t *query = zvec_vector_query_create();
if (!query) {
fprintf(stderr, "Failed to create vector query\n");
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
return 1;
}
zvec_vector_query_set_field_name(query, "embedding");
zvec_vector_query_set_query_vector(query, vector1, 128 * sizeof(float));
zvec_vector_query_set_topk(query, 5);
zvec_vector_query_set_filter(query, "");
zvec_vector_query_set_include_vector(query, true);
zvec_vector_query_set_include_doc_id(query, true);
zvec_doc_t **results = NULL;
size_t result_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&results, &result_count);
if (error == ZVEC_OK) {
printf("✓ Vector query successful - Returned %zu results\n", result_count);
zvec_docs_free(results, result_count);
}
zvec_vector_query_destroy(query);
// 14. Cleanup resources
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
printf("✓ Schema example completed\n");
return 0;
}
examples/c/diskann_example.c
+346
View File
@@ -0,0 +1,346 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/**
* @file diskann_example.c
* @brief End-to-end example demonstrating DiskANN index usage via the C API.
*
* DiskANN is a disk-based approximate nearest neighbor search algorithm
* optimized for large-scale datasets that exceed available memory. It uses
* a Vamana graph structure combined with product quantization (PQ) to
* achieve high recall with efficient disk I/O.
*
* NOTE: DiskANN requires Linux x86_64 with libaio. On other platforms the
* example will compile but the runtime plugin will fail to load.
*
* Workflow demonstrated:
* 1. Create collection schema with DiskANN-indexed vector field
* 2. Insert documents with high-dimensional vectors
* 3. Flush collection (triggers PQ training + graph build)
* 4. Search using DiskANN query parameters (list_size controls recall)
* 5. Clean up all resources
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/* --------------------------------------------------------------------------
* Helpers
* -------------------------------------------------------------------------- */
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
#define VECTOR_DIM 64
#define NUM_DOCS 100
#define COLLECTION_DIR "./diskann_example_collection"
/* --------------------------------------------------------------------------
* Main
* -------------------------------------------------------------------------- */
int main(void) {
printf("=== ZVec DiskANN Index Example ===\n\n");
zvec_error_code_t error;
int i;
/* ------------------------------------------------------------------
* Step 1: Create collection schema
* ------------------------------------------------------------------ */
printf("[Step 1] Creating collection schema...\n");
zvec_collection_schema_t *schema =
zvec_collection_schema_create("diskann_example");
if (!schema) {
fprintf(stderr, "Failed to create schema\n");
return 1;
}
/* Index params — declared up-front and NULL-initialized so the
* cleanup_schema path never touches an uninitialized pointer even if an
* early field addition fails. */
zvec_index_params_t *invert_params = NULL;
zvec_index_params_t *diskann_params = NULL;
/* Scalar field with inverted index (for primary key / filtering) */
invert_params = zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
zvec_index_params_set_invert_params(invert_params, true, false);
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
zvec_field_schema_set_index_params(id_field, invert_params);
error = zvec_collection_schema_add_field(schema, id_field);
if (handle_error(error, "adding id field") != ZVEC_OK) {
goto cleanup_schema;
}
printf(" + id field (STRING, inverted index)\n");
/* Vector field with DiskANN index */
diskann_params = zvec_index_params_create(ZVEC_INDEX_TYPE_DISKANN);
if (!diskann_params) {
fprintf(stderr, "Failed to create DiskANN index parameters\n");
goto cleanup_schema;
}
zvec_index_params_set_metric_type(diskann_params, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_diskann_params(
diskann_params, 64, /* max_degree: graph connectivity */
100, /* list_size: build-time candidates */
8); /* pq_chunk_num: PQ chunks (0=auto) */
printf(
" DiskANN index params: max_degree=%d, list_size=%d, pq_chunk_num=%d\n",
zvec_index_params_get_diskann_max_degree(diskann_params),
zvec_index_params_get_diskann_list_size(diskann_params),
zvec_index_params_get_diskann_pq_chunk_num(diskann_params));
zvec_field_schema_t *embedding_field = zvec_field_schema_create(
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, VECTOR_DIM);
zvec_field_schema_set_index_params(embedding_field, diskann_params);
error = zvec_collection_schema_add_field(schema, embedding_field);
if (handle_error(error, "adding embedding field") != ZVEC_OK) {
goto cleanup_schema;
}
printf(" + embedding field (VECTOR_FP32, %dD, DiskANN index)\n", VECTOR_DIM);
/* Index params are copied into field schemas; safe to destroy now */
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(diskann_params);
invert_params = NULL;
diskann_params = NULL;
/* ------------------------------------------------------------------
* Step 2: Create and open collection
* ------------------------------------------------------------------ */
printf("\n[Step 2] Creating collection...\n");
zvec_collection_options_t *options = zvec_collection_options_create();
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open(COLLECTION_DIR, schema, options,
&collection);
zvec_collection_options_destroy(options);
if (handle_error(error, "creating collection") != ZVEC_OK) {
goto cleanup_schema;
}
printf(" Collection created at %s\n", COLLECTION_DIR);
/* ------------------------------------------------------------------
* Step 3: Generate and insert documents
* ------------------------------------------------------------------ */
printf("\n[Step 3] Inserting %d documents with %dD vectors...\n", NUM_DOCS,
VECTOR_DIM);
/* Allocate vector storage */
float(*vectors)[VECTOR_DIM] =
(float(*)[VECTOR_DIM])malloc(NUM_DOCS * VECTOR_DIM * sizeof(float));
if (!vectors) {
fprintf(stderr, "Failed to allocate vector storage\n");
goto cleanup_collection;
}
/* Generate deterministic vector data */
for (i = 0; i < NUM_DOCS; i++) {
for (int d = 0; d < VECTOR_DIM; d++) {
vectors[i][d] = (float)((i * VECTOR_DIM + d) % 1000) / 1000.0f;
}
}
/* Insert in batches */
int batch_size = 20;
size_t total_success = 0, total_error = 0;
for (int batch_start = 0; batch_start < NUM_DOCS; batch_start += batch_size) {
int count = batch_start + batch_size > NUM_DOCS ? NUM_DOCS - batch_start
: batch_size;
zvec_doc_t **docs =
(zvec_doc_t **)malloc((size_t)count * sizeof(zvec_doc_t *));
for (i = 0; i < count; i++) {
int idx = batch_start + i;
docs[i] = zvec_doc_create();
char pk[32];
snprintf(pk, sizeof(pk), "doc_%04d", idx);
zvec_doc_set_pk(docs[i], pk);
zvec_doc_add_field_by_value(docs[i], "id", ZVEC_DATA_TYPE_STRING, pk,
strlen(pk));
zvec_doc_add_field_by_value(docs[i], "embedding",
ZVEC_DATA_TYPE_VECTOR_FP32, vectors[idx],
VECTOR_DIM * sizeof(float));
}
size_t success_count = 0, error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)docs,
(size_t)count, &success_count, &error_count);
if (error != ZVEC_OK) {
handle_error(error, "inserting batch");
}
total_success += success_count;
total_error += error_count;
for (i = 0; i < count; i++) {
zvec_doc_destroy(docs[i]);
}
free(docs);
}
printf(" Inserted: %zu succeeded, %zu failed\n", total_success, total_error);
/* ------------------------------------------------------------------
* Step 4: Flush to trigger index build (PQ training + graph construction)
* ------------------------------------------------------------------ */
printf("\n[Step 4] Flushing collection (triggers DiskANN index build)...\n");
error = zvec_collection_flush(collection);
if (handle_error(error, "flushing collection") != ZVEC_OK) {
goto cleanup_vectors;
}
zvec_collection_stats_t *stats = NULL;
error = zvec_collection_get_stats(collection, &stats);
if (error == ZVEC_OK && stats) {
printf(" Document count after flush: %llu\n",
(unsigned long long)zvec_collection_stats_get_doc_count(stats));
zvec_collection_stats_destroy(stats);
}
/* ------------------------------------------------------------------
* Step 5: Search with DiskANN query parameters
* ------------------------------------------------------------------ */
printf("\n[Step 5] Searching with DiskANN query parameters...\n");
/* Create DiskANN query params — list_size controls the search frontier
* (beam width). Larger values improve recall at the cost of latency. */
zvec_diskann_query_params_t *da_qp = zvec_query_params_diskann_create(200);
if (!da_qp) {
fprintf(stderr, "Failed to create DiskANN query params\n");
goto cleanup_vectors;
}
printf(" DiskANN query params: list_size=%d\n",
zvec_query_params_diskann_get_list_size(da_qp));
/* Build the vector query */
zvec_vector_query_t *query = zvec_vector_query_create();
zvec_vector_query_set_field_name(query, "embedding");
zvec_vector_query_set_query_vector(query, vectors[0],
VECTOR_DIM * sizeof(float));
zvec_vector_query_set_topk(query, 10);
zvec_vector_query_set_include_vector(query, false);
zvec_vector_query_set_include_doc_id(query, true);
/* Attach DiskANN query params (ownership transfers to query) */
error = zvec_vector_query_set_diskann_params(query, da_qp);
if (handle_error(error, "setting DiskANN query params") != ZVEC_OK) {
zvec_vector_query_destroy(query);
goto cleanup_vectors;
}
/* da_qp is now owned by query — do NOT call diskann_destroy on it */
/* Execute the query */
zvec_doc_t **results = NULL;
size_t result_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&results, &result_count);
if (error != ZVEC_OK) {
handle_error(error, "executing DiskANN query");
printf(
" (This is expected on non-Linux platforms — DiskANN requires "
"libaio)\n");
} else {
printf(" Query returned %zu results:\n", result_count);
for (size_t r = 0; r < result_count && r < 5; r++) {
const char *pk = zvec_doc_get_pk_copy(results[r]);
printf(" [%zu] pk=%s doc_id=%llu score=%.6f\n", r + 1,
pk ? pk : "NULL",
(unsigned long long)zvec_doc_get_doc_id(results[r]),
zvec_doc_get_score(results[r]));
if (pk) {
zvec_free((void *)pk);
}
}
if (result_count > 5) {
printf(" ... and %zu more\n", result_count - 5);
}
zvec_docs_free(results, result_count);
}
zvec_vector_query_destroy(query);
/* ------------------------------------------------------------------
* Step 6: Demonstrate list_size tuning (higher recall vs. lower latency)
* ------------------------------------------------------------------ */
printf("\n[Step 6] Tuning list_size for recall/latency trade-off...\n");
int list_sizes[] = {50, 100, 300};
for (int li = 0; li < 3; li++) {
zvec_diskann_query_params_t *tune_qp =
zvec_query_params_diskann_create(list_sizes[li]);
zvec_vector_query_t *tune_query = zvec_vector_query_create();
zvec_vector_query_set_field_name(tune_query, "embedding");
zvec_vector_query_set_query_vector(tune_query, vectors[0],
VECTOR_DIM * sizeof(float));
zvec_vector_query_set_topk(tune_query, 10);
zvec_vector_query_set_include_doc_id(tune_query, true);
zvec_vector_query_set_diskann_params(tune_query, tune_qp);
zvec_doc_t **tune_results = NULL;
size_t tune_count = 0;
error = zvec_collection_query(collection,
(const zvec_vector_query_t *)tune_query,
&tune_results, &tune_count);
if (error == ZVEC_OK) {
printf(" list_size=%3d -> %zu results returned\n", list_sizes[li],
tune_count);
zvec_docs_free(tune_results, tune_count);
} else {
printf(" list_size=%3d -> query failed (expected on non-Linux)\n",
list_sizes[li]);
}
zvec_vector_query_destroy(tune_query);
}
/* ------------------------------------------------------------------
* Cleanup
* ------------------------------------------------------------------ */
cleanup_vectors:
free(vectors);
cleanup_collection:
zvec_collection_destroy(collection);
cleanup_schema:
zvec_collection_schema_destroy(schema);
if (invert_params) {
zvec_index_params_destroy(invert_params);
}
if (diskann_params) {
zvec_index_params_destroy(diskann_params);
}
printf("\n DiskANN index type string: %s\n",
zvec_index_type_to_string(ZVEC_INDEX_TYPE_DISKANN));
printf("=== Example completed ===\n");
return 0;
}
examples/c/doc_example.c
+539
View File
@@ -0,0 +1,539 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Create a test document with all data types
* @param doc_index Document index for generating unique data
* @return zvec_doc_t* Created document pointer
*/
static zvec_doc_t *create_full_type_test_doc(int doc_index) {
zvec_doc_t *doc = zvec_doc_create();
if (!doc) {
fprintf(stderr, "Failed to create document\n");
return NULL;
}
// Set primary key
char pk_buffer[32];
snprintf(pk_buffer, sizeof(pk_buffer), "doc_%d", doc_index);
zvec_doc_set_pk(doc, pk_buffer);
// Add Id field with inverted index
char id_buffer[32];
snprintf(id_buffer, sizeof(id_buffer), "id_%d", doc_index);
zvec_doc_add_field_by_value(doc, "id", ZVEC_DATA_TYPE_STRING, id_buffer,
strlen(id_buffer));
// Add scalar fields with different data types
// String field
char string_value[64];
snprintf(string_value, sizeof(string_value), "test_string_%d", doc_index);
zvec_doc_add_field_by_value(doc, "string_field", ZVEC_DATA_TYPE_STRING,
string_value, strlen(string_value));
// Boolean field
bool bool_value = (doc_index % 2 == 0);
zvec_doc_add_field_by_value(doc, "bool_field", ZVEC_DATA_TYPE_BOOL,
&bool_value, sizeof(bool_value));
// Integer fields
int32_t int32_value = doc_index * 1000;
zvec_doc_add_field_by_value(doc, "int32_field", ZVEC_DATA_TYPE_INT32,
&int32_value, sizeof(int32_value));
int64_t int64_value = (int64_t)doc_index * 1000000LL;
zvec_doc_add_field_by_value(doc, "int64_field", ZVEC_DATA_TYPE_INT64,
&int64_value, sizeof(int64_value));
// Floating point fields
float float_value = (float)doc_index * 1.5f;
zvec_doc_add_field_by_value(doc, "float_field", ZVEC_DATA_TYPE_FLOAT,
&float_value, sizeof(float_value));
double double_value = (double)doc_index * 2.718281828;
zvec_doc_add_field_by_value(doc, "double_field", ZVEC_DATA_TYPE_DOUBLE,
&double_value, sizeof(double_value));
// Vector fields with different dimensions
// FP32 vector (3D)
float fp32_vector[3] = {(float)doc_index, (float)doc_index * 2.0f,
(float)doc_index * 3.0f};
zvec_doc_add_field_by_value(doc, "vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32,
fp32_vector, 3 * sizeof(float));
// Larger FP32 vector (16D)
float large_vector[16];
for (int i = 0; i < 16; i++) {
large_vector[i] = (float)(doc_index * 16 + i) / 256.0f;
}
zvec_doc_add_field_by_value(doc, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
large_vector, 16 * sizeof(float));
return doc;
}
/**
* @brief Compare two documents for equality
*/
static bool compare_documents(const zvec_doc_t *doc1, const zvec_doc_t *doc2) {
if (!doc1 || !doc2) return false;
// Compare primary keys
const char *pk1 = zvec_doc_get_pk_pointer(doc1);
const char *pk2 = zvec_doc_get_pk_pointer(doc2);
if (!pk1 || !pk2 || strcmp(pk1, pk2) != 0) {
return false;
}
// TODO: Compare other fields and values
return true;
}
/**
* @brief Print document fields and their values
* @param doc The document to print
* @param doc_index Document index for identification
*/
static void print_doc(const zvec_doc_t *doc, int doc_index) {
if (!doc) {
printf("Document %d: NULL document\n", doc_index);
return;
}
printf("\n=== Document %d ===\n", doc_index);
// Print primary key
const char *pk = zvec_doc_get_pk_pointer(doc);
printf("Primary Key: %s\n", pk ? pk : "NULL");
// Print document ID
uint64_t doc_id = zvec_doc_get_doc_id(doc);
printf("Document ID: %llu\n", (unsigned long long)doc_id);
// Print score
float score = zvec_doc_get_score(doc);
printf("Score: %.6f\n", score);
// Print scalar fields
printf("\nScalar Fields:\n");
// ID field (using pointer function for strings)
const void *id_value = NULL;
size_t id_size = 0;
zvec_error_code_t error = zvec_doc_get_field_value_pointer(
doc, "id", ZVEC_DATA_TYPE_STRING, &id_value, &id_size);
if (error == ZVEC_OK && id_value) {
printf(" id: %.*s\n", (int)id_size, (const char *)id_value);
}
// String field (using pointer function for strings)
const void *string_value = NULL;
size_t string_size = 0;
error = zvec_doc_get_field_value_pointer(
doc, "string_field", ZVEC_DATA_TYPE_STRING, &string_value, &string_size);
if (error == ZVEC_OK && string_value) {
printf(" string_field: %.*s\n", (int)string_size,
(const char *)string_value);
}
// Boolean field
bool bool_value;
error = zvec_doc_get_field_value_basic(doc, "bool_field", ZVEC_DATA_TYPE_BOOL,
&bool_value, sizeof(bool_value));
if (error == ZVEC_OK) {
printf(" bool_field: %s\n", bool_value ? "true" : "false");
}
// Int32 field
int32_t int32_value;
error =
zvec_doc_get_field_value_basic(doc, "int32_field", ZVEC_DATA_TYPE_INT32,
&int32_value, sizeof(int32_value));
if (error == ZVEC_OK) {
printf(" int32_field: %d\n", int32_value);
}
// Int64 field
int64_t int64_value;
error =
zvec_doc_get_field_value_basic(doc, "int64_field", ZVEC_DATA_TYPE_INT64,
&int64_value, sizeof(int64_value));
if (error == ZVEC_OK) {
printf(" int64_field: %lld\n", (long long)int64_value);
}
// Float field
float float_value;
error =
zvec_doc_get_field_value_basic(doc, "float_field", ZVEC_DATA_TYPE_FLOAT,
&float_value, sizeof(float_value));
if (error == ZVEC_OK) {
printf(" float_field: %.6f\n", float_value);
}
// Double field
double double_value;
error =
zvec_doc_get_field_value_basic(doc, "double_field", ZVEC_DATA_TYPE_DOUBLE,
&double_value, sizeof(double_value));
if (error == ZVEC_OK) {
printf(" double_field: %.6f\n", double_value);
}
// Print vector fields (using copy function for complex types)
printf("\nVector Fields:\n");
// FP32 vector (3D)
void *fp32_vector = NULL;
size_t fp32_size = 0;
error = zvec_doc_get_field_value_copy(
doc, "vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32, &fp32_vector, &fp32_size);
if (error == ZVEC_OK && fp32_vector) {
const float *vec = (const float *)fp32_vector;
size_t dim = fp32_size / sizeof(float);
printf(" vector_fp32 (%zuD): [", dim);
for (size_t i = 0; i < dim && i < 10; i++) { // Limit to first 10 elements
printf("%.3f", vec[i]);
if (i < dim - 1 && i < 9) printf(", ");
}
if (dim > 10) printf(", ...");
printf("]\n");
zvec_free(fp32_vector); // Free the allocated memory
}
// Large vector (16D)
void *large_vector = NULL;
size_t large_size = 0;
error = zvec_doc_get_field_value_copy(doc, "large_vector",
ZVEC_DATA_TYPE_VECTOR_FP32,
&large_vector, &large_size);
if (error == ZVEC_OK && large_vector) {
const float *vec = (const float *)large_vector;
size_t dim = large_size / sizeof(float);
printf(" large_vector (%zuD): [", dim);
for (size_t i = 0; i < dim && i < 10; i++) { // Limit to first 10 elements
printf("%.3f", vec[i]);
if (i < dim - 1 && i < 9) printf(", ");
}
if (dim > 10) printf(", ...");
printf("]\n");
zvec_free(large_vector); // Free the allocated memory
}
printf("==================\n\n");
}
/**
* @brief Document creation, manipulation, and query example
*/
int main() {
printf("=== ZVec Document Example ===\n\n");
zvec_error_code_t error;
// 1. Create collection schema for document testing
zvec_collection_schema_t *schema =
zvec_collection_schema_create("doc_example_collection");
if (!schema) {
fprintf(stderr, "Failed to create collection schema\n");
return -1;
}
printf("✓ Collection schema created\n");
// 2. Create index parameters
zvec_index_params_t *invert_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params) {
fprintf(stderr, "Failed to create invert index parameters\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_invert_params(invert_params, true, false);
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params) {
fprintf(stderr, "Failed to create HNSW index parameters\n");
zvec_index_params_destroy(invert_params);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params, 16, 200);
// 3. Create fields for all data types
printf("Creating fields for all data types...\n");
// Id field with inverted index
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
if (id_field) {
zvec_field_schema_set_index_params(id_field, invert_params);
error = zvec_collection_schema_add_field(schema, id_field);
if (handle_error(error, "adding ID field") == ZVEC_OK) {
printf("✓ ID field with inverted index added\n");
}
}
// Scalar fields
zvec_field_schema_t *string_field =
zvec_field_schema_create("string_field", ZVEC_DATA_TYPE_STRING, true, 0);
zvec_field_schema_t *bool_field =
zvec_field_schema_create("bool_field", ZVEC_DATA_TYPE_BOOL, true, 0);
zvec_field_schema_t *int32_field =
zvec_field_schema_create("int32_field", ZVEC_DATA_TYPE_INT32, true, 0);
zvec_field_schema_t *int64_field =
zvec_field_schema_create("int64_field", ZVEC_DATA_TYPE_INT64, true, 0);
zvec_field_schema_t *float_field =
zvec_field_schema_create("float_field", ZVEC_DATA_TYPE_FLOAT, true, 0);
zvec_field_schema_t *double_field =
zvec_field_schema_create("double_field", ZVEC_DATA_TYPE_DOUBLE, true, 0);
if (string_field) zvec_collection_schema_add_field(schema, string_field);
if (bool_field) zvec_collection_schema_add_field(schema, bool_field);
if (int32_field) zvec_collection_schema_add_field(schema, int32_field);
if (int64_field) zvec_collection_schema_add_field(schema, int64_field);
if (float_field) zvec_collection_schema_add_field(schema, float_field);
if (double_field) zvec_collection_schema_add_field(schema, double_field);
// Vector fields
zvec_field_schema_t *vector_fp32_field = zvec_field_schema_create(
"vector_fp32", ZVEC_DATA_TYPE_VECTOR_FP32, false, 3);
zvec_field_schema_t *large_vector_field = zvec_field_schema_create(
"large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 16);
if (vector_fp32_field) {
zvec_field_schema_set_index_params(vector_fp32_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, vector_fp32_field);
if (handle_error(error, "adding vector FP32 field") == ZVEC_OK) {
printf("✓ Vector FP32 field with HNSW index added\n");
}
}
if (large_vector_field) {
zvec_field_schema_set_index_params(large_vector_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, large_vector_field);
if (handle_error(error, "adding large vector field") == ZVEC_OK) {
printf("✓ Large vector field with HNSW index added\n");
}
}
// 4. Create collection
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
fprintf(stderr, "Failed to create collection options\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open("./doc_example_collection", schema,
options, &collection);
zvec_collection_options_destroy(options);
if (handle_error(error, "creating collection") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return -1;
}
printf("✓ Collection created successfully\n");
// 5. Create and insert multiple test documents
printf("Creating and inserting test documents...\n");
#define DOC_COUNT 5
// Use dynamic allocation for MSVC compatibility (no VLA support)
zvec_doc_t **test_docs =
(zvec_doc_t **)malloc(DOC_COUNT * sizeof(zvec_doc_t *));
if (!test_docs) {
fprintf(stderr, "Failed to allocate test documents\n");
goto cleanup;
}
for (int i = 0; i < DOC_COUNT; i++) {
test_docs[i] = create_full_type_test_doc(i);
if (!test_docs[i]) {
fprintf(stderr, "Failed to create document %d\n", i);
// Cleanup
for (int j = 0; j < i; j++) {
zvec_doc_destroy(test_docs[j]);
}
goto cleanup;
}
printf("✓ Created document %d with PK: %s\n", i,
zvec_doc_get_pk_pointer(test_docs[i]));
}
// Print all documents before insertion
printf("\nDocuments before insertion:\n");
for (int i = 0; i < DOC_COUNT; i++) {
print_doc(test_docs[i], i);
}
// Insert documents
size_t success_count = 0, error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)test_docs,
DOC_COUNT, &success_count, &error_count);
if (handle_error(error, "inserting documents") == ZVEC_OK) {
printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count,
error_count);
}
// 6. Flush collection
error = zvec_collection_flush(collection);
if (handle_error(error, "flushing collection") != ZVEC_OK) {
printf("Warning: Collection flush failed\n");
} else {
printf("✓ Collection flushed successfully\n");
}
// Use the first document's vector for querying
float query_vector[] = {0.0f, 0.0f, 0.0f};
zvec_vector_query_t *query = zvec_vector_query_create();
if (!query) {
fprintf(stderr, "Failed to create vector query\n");
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_vector_query_set_field_name(query, "vector_fp32");
zvec_vector_query_set_query_vector(query, query_vector, 3 * sizeof(float));
zvec_vector_query_set_topk(query, 5);
zvec_vector_query_set_filter(query, "");
zvec_vector_query_set_include_vector(query, true);
zvec_vector_query_set_include_doc_id(query, true);
zvec_doc_t **query_results = NULL;
size_t result_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&query_results, &result_count);
if (handle_error(error, "querying documents") != ZVEC_OK) {
query_results = NULL;
result_count = 0;
}
printf("Query returned %zu results\n", result_count);
// Print query results
printf("\nQuery Results:\n");
for (size_t i = 0; i < result_count; i++) {
print_doc(query_results[i], i);
}
// Compare query results
for (size_t i = 0; i < result_count && i < DOC_COUNT; i++) {
const char *result_pk = zvec_doc_get_pk_pointer(query_results[i]);
printf("Comparing query result[%zu]: %s\n", i, result_pk);
// Find matching original document
bool found = false;
for (int j = 0; j < DOC_COUNT; j++) {
const char *original_pk = zvec_doc_get_pk_pointer(test_docs[j]);
if (strcmp(result_pk, original_pk) == 0) {
if (compare_documents(test_docs[j], query_results[i])) {
printf("✓ Query result %s matches original document\n", result_pk);
} else {
printf("✗ Query result %s does not match original document\n",
result_pk);
}
found = true;
break;
}
}
if (!found) {
printf("⚠ Original document not found for: %s\n", result_pk);
}
}
// 7. Filter query test
printf("\n=== Filter Query Test ===\n");
// Create filtered query
zvec_vector_query_set_filter(query, "string_field = 'string_field_0'");
zvec_doc_t **filtered_results = NULL;
size_t filtered_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&filtered_results, &filtered_count);
if (handle_error(error, "filtered querying") == ZVEC_OK) {
printf("Filtered query returned %zu results\n", filtered_count);
// Verify filter results
bool filter_correct = true;
for (size_t i = 0; i < filtered_count; i++) {
// Note: Field value access may require different API
// For now, we'll just check that we got results
const char *pk = zvec_doc_get_pk_pointer(filtered_results[i]);
if (strstr(pk, "doc_") == NULL) {
filter_correct = false;
break;
}
}
if (filter_correct) {
printf("✓ Filter query results are correct\n");
} else {
printf("✗ Filter query results are incorrect\n");
}
if (filtered_results) {
zvec_docs_free(filtered_results, filtered_count);
}
}
// 8. Cleanup query results
if (query_results) {
zvec_docs_free(query_results, result_count);
}
// 9. Cleanup documents
for (int i = 0; i < DOC_COUNT; i++) {
zvec_doc_destroy(test_docs[i]);
}
free(test_docs); // Free the array itself
// 10. Final cleanup
cleanup:
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
printf("✓ Document example completed\n");
return 0;
}
examples/c/field_schema_example.c
+312
View File
@@ -0,0 +1,312 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Field schema creation and management example
*/
int main() {
printf("=== ZVec Field Schema Example ===\n\n");
zvec_error_code_t error;
// 1. Create collection schema
zvec_collection_schema_t *schema =
zvec_collection_schema_create("field_example_collection");
if (!schema) {
fprintf(stderr, "Failed to create collection schema\n");
return -1;
}
printf("✓ Collection schema created successfully\n");
// 2. Create different types of index parameters
zvec_index_params_t *invert_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params) {
fprintf(stderr, "Failed to create invert index parameters\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_invert_params(invert_params, true, false);
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params) {
fprintf(stderr, "Failed to create HNSW index parameters\n");
zvec_index_params_destroy(invert_params);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_COSINE);
zvec_index_params_set_hnsw_params(hnsw_params, 16, 200);
zvec_index_params_t *flat_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT);
if (!flat_params) {
fprintf(stderr, "Failed to create Flat index parameters\n");
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(flat_params, ZVEC_METRIC_TYPE_L2);
if (!invert_params || !hnsw_params || !flat_params) {
fprintf(stderr, "Failed to create index parameters\n");
zvec_index_params_destroy(invert_params);
zvec_index_params_destroy(hnsw_params);
zvec_index_params_destroy(flat_params);
zvec_collection_schema_destroy(schema);
return -1;
}
// 3. Create scalar fields with different data types
printf("Creating scalar fields...\n");
// String field with inverted index
zvec_field_schema_t *name_field =
zvec_field_schema_create("name", ZVEC_DATA_TYPE_STRING, false, 0);
if (name_field) {
zvec_field_schema_set_index_params(name_field, invert_params);
error = zvec_collection_schema_add_field(schema, name_field);
if (handle_error(error, "adding name field") == ZVEC_OK) {
printf("✓ String field 'name' with inverted index added\n");
}
}
// Integer field
zvec_field_schema_t *age_field =
zvec_field_schema_create("age", ZVEC_DATA_TYPE_INT32, true, 0);
if (age_field) {
error = zvec_collection_schema_add_field(schema, age_field);
if (handle_error(error, "adding age field") == ZVEC_OK) {
printf("✓ Integer field 'age' added\n");
}
}
// Float field
zvec_field_schema_t *score_field =
zvec_field_schema_create("score", ZVEC_DATA_TYPE_FLOAT, true, 0);
if (score_field) {
error = zvec_collection_schema_add_field(schema, score_field);
if (handle_error(error, "adding score field") == ZVEC_OK) {
printf("✓ Float field 'score' added\n");
}
}
// Boolean field
zvec_field_schema_t *active_field =
zvec_field_schema_create("active", ZVEC_DATA_TYPE_BOOL, false, 0);
if (active_field) {
error = zvec_collection_schema_add_field(schema, active_field);
if (handle_error(error, "adding active field") == ZVEC_OK) {
printf("✓ Boolean field 'active' added\n");
}
}
// 4. Create vector fields with different dimensions and indexes
printf("Creating vector fields...\n");
// Small dimension vector with HNSW index
zvec_field_schema_t *small_vector_field = zvec_field_schema_create(
"small_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 32);
if (small_vector_field) {
zvec_field_schema_set_index_params(small_vector_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, small_vector_field);
if (handle_error(error, "adding small vector field") == ZVEC_OK) {
printf(
"✓ Small vector field 'small_vector' (32D) with HNSW index added\n");
}
}
// Medium dimension vector with Flat index
zvec_field_schema_t *medium_vector_field = zvec_field_schema_create(
"medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128);
if (medium_vector_field) {
zvec_field_schema_set_index_params(medium_vector_field, flat_params);
error = zvec_collection_schema_add_field(schema, medium_vector_field);
if (handle_error(error, "adding medium vector field") == ZVEC_OK) {
printf(
"✓ Medium vector field 'medium_vector' (128D) with Flat index "
"added\n");
}
}
// Large dimension vector with HNSW index
zvec_field_schema_t *large_vector_field = zvec_field_schema_create(
"large_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 512);
if (large_vector_field) {
zvec_field_schema_set_index_params(large_vector_field, hnsw_params);
error = zvec_collection_schema_add_field(schema, large_vector_field);
if (handle_error(error, "adding large vector field") == ZVEC_OK) {
printf(
"✓ Large vector field 'large_vector' (512D) with HNSW index added\n");
}
}
// 5. Create collection with the schema
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
fprintf(stderr, "Failed to create collection options\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open("./field_example_collection", schema,
options, &collection);
zvec_collection_options_destroy(options);
if (handle_error(error, "creating collection") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return -1;
}
printf("✓ Collection created successfully\n");
// 6. Create test documents with various field types
printf("Creating test documents...\n");
zvec_doc_t *doc1 = zvec_doc_create();
zvec_doc_t *doc2 = zvec_doc_create();
if (!doc1 || !doc2) {
fprintf(stderr, "Failed to create documents\n");
goto cleanup;
}
// Document 1
zvec_doc_set_pk(doc1, "user1");
zvec_doc_add_field_by_value(doc1, "name", ZVEC_DATA_TYPE_STRING,
"Alice Johnson", strlen("Alice Johnson"));
int32_t age1 = 28;
zvec_doc_add_field_by_value(doc1, "age", ZVEC_DATA_TYPE_INT32, &age1,
sizeof(age1));
float score1 = 87.5f;
zvec_doc_add_field_by_value(doc1, "score", ZVEC_DATA_TYPE_FLOAT, &score1,
sizeof(score1));
bool active1 = true;
zvec_doc_add_field_by_value(doc1, "active", ZVEC_DATA_TYPE_BOOL, &active1,
sizeof(active1));
// Add vector data
float small_vec1[32];
float medium_vec1[128];
float large_vec1[512];
for (int i = 0; i < 32; i++) small_vec1[i] = (float)i / 32.0f;
for (int i = 0; i < 128; i++) medium_vec1[i] = (float)i / 128.0f;
for (int i = 0; i < 512; i++) large_vec1[i] = (float)i / 512.0f;
zvec_doc_add_field_by_value(doc1, "small_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
small_vec1, 32 * sizeof(float));
zvec_doc_add_field_by_value(doc1, "medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
medium_vec1, 128 * sizeof(float));
zvec_doc_add_field_by_value(doc1, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
large_vec1, 512 * sizeof(float));
// Document 2
zvec_doc_set_pk(doc2, "user2");
zvec_doc_add_field_by_value(doc2, "name", ZVEC_DATA_TYPE_STRING, "Bob Smith",
strlen("Bob Smith"));
int32_t age2 = 35;
zvec_doc_add_field_by_value(doc2, "age", ZVEC_DATA_TYPE_INT32, &age2,
sizeof(age2));
float score2 = 92.0f;
zvec_doc_add_field_by_value(doc2, "score", ZVEC_DATA_TYPE_FLOAT, &score2,
sizeof(score2));
bool active2 = false;
zvec_doc_add_field_by_value(doc2, "active", ZVEC_DATA_TYPE_BOOL, &active2,
sizeof(active2));
// Add vector data
float small_vec2[32];
float medium_vec2[128];
float large_vec2[512];
for (int i = 0; i < 32; i++) small_vec2[i] = (float)(32 - i) / 32.0f;
for (int i = 0; i < 128; i++) medium_vec2[i] = (float)(128 - i) / 128.0f;
for (int i = 0; i < 512; i++) large_vec2[i] = (float)(512 - i) / 512.0f;
zvec_doc_add_field_by_value(doc2, "small_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
small_vec2, 32 * sizeof(float));
zvec_doc_add_field_by_value(doc2, "medium_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
medium_vec2, 128 * sizeof(float));
zvec_doc_add_field_by_value(doc2, "large_vector", ZVEC_DATA_TYPE_VECTOR_FP32,
large_vec2, 512 * sizeof(float));
// 7. Insert documents
zvec_doc_t *docs[] = {doc1, doc2};
size_t success_count = 0, error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 2,
&success_count, &error_count);
if (handle_error(error, "inserting documents") == ZVEC_OK) {
printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count,
error_count);
}
// 8. Flush and test queries
zvec_collection_flush(collection);
printf("✓ Collection flushed\n");
// Test vector query on medium vector field
zvec_vector_query_t *query = zvec_vector_query_create();
if (!query) {
fprintf(stderr, "Failed to create vector query\n");
goto cleanup;
}
zvec_vector_query_set_field_name(query, "medium_vector");
zvec_vector_query_set_query_vector(query, medium_vec1, 128 * sizeof(float));
zvec_vector_query_set_topk(query, 2);
zvec_vector_query_set_filter(query, "");
zvec_vector_query_set_include_vector(query, false);
zvec_vector_query_set_include_doc_id(query, true);
zvec_doc_t **results = NULL;
size_t result_count = 0;
error = zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&results, &result_count);
if (error == ZVEC_OK) {
printf("✓ Vector query successful - Found %zu results\n", result_count);
zvec_docs_free(results, result_count);
}
zvec_vector_query_destroy(query);
// 9. Cleanup
cleanup:
if (doc1) zvec_doc_destroy(doc1);
if (doc2) zvec_doc_destroy(doc2);
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
printf("✓ Field schema example completed\n");
return 0;
}
examples/c/index_example.c
+398
View File
@@ -0,0 +1,398 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Index creation and management example
*/
int main() {
printf("=== ZVec Index Example ===\n\n");
zvec_error_code_t error;
// 1. Create collection schema
zvec_collection_schema_t *schema =
zvec_collection_schema_create("index_example_collection");
if (!schema) {
fprintf(stderr, "Failed to create collection schema\n");
return -1;
}
printf("✓ Collection schema created successfully\n");
// 2. Create different index parameter configurations
printf("Creating index parameters...\n");
// Inverted index parameters
zvec_index_params_t *invert_params_standard =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params_standard) {
fprintf(stderr, "Failed to create invert index parameters (standard)\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_invert_params(invert_params_standard, true, false);
zvec_index_params_t *invert_params_extended =
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
if (!invert_params_extended) {
fprintf(stderr, "Failed to create invert index parameters (extended)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_invert_params(invert_params_extended, true, true);
// HNSW index parameters with different configurations
zvec_index_params_t *hnsw_params_fast =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params_fast) {
fprintf(stderr, "Failed to create HNSW index parameters (fast)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_index_params_destroy(invert_params_extended);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params_fast, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params_fast, 16, 100);
// Demonstrate INT8 quantization with random rotation preprocessing
// (enable_rotate rotates vectors before INT8 quantization to reduce error)
zvec_index_params_set_quantize_type(hnsw_params_fast,
ZVEC_QUANTIZE_TYPE_INT8);
zvec_index_params_set_quantizer_enable_rotate(hnsw_params_fast, true);
zvec_index_params_t *hnsw_params_balanced =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params_balanced) {
fprintf(stderr, "Failed to create HNSW index parameters (balanced)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_index_params_destroy(invert_params_extended);
zvec_index_params_destroy(hnsw_params_fast);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params_balanced,
ZVEC_METRIC_TYPE_COSINE);
zvec_index_params_set_hnsw_params(hnsw_params_balanced, 32, 200);
zvec_index_params_t *hnsw_params_accurate =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params_accurate) {
fprintf(stderr, "Failed to create HNSW index parameters (accurate)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_index_params_destroy(invert_params_extended);
zvec_index_params_destroy(hnsw_params_fast);
zvec_index_params_destroy(hnsw_params_balanced);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params_accurate, ZVEC_METRIC_TYPE_IP);
zvec_index_params_set_hnsw_params(hnsw_params_accurate, 64, 400);
// Flat index parameters
zvec_index_params_t *flat_params_l2 =
zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT);
if (!flat_params_l2) {
fprintf(stderr, "Failed to create Flat index parameters (L2)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_index_params_destroy(invert_params_extended);
zvec_index_params_destroy(hnsw_params_fast);
zvec_index_params_destroy(hnsw_params_balanced);
zvec_index_params_destroy(hnsw_params_accurate);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(flat_params_l2, ZVEC_METRIC_TYPE_L2);
zvec_index_params_t *flat_params_cosine =
zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT);
if (!flat_params_cosine) {
fprintf(stderr, "Failed to create Flat index parameters (cosine)\n");
zvec_index_params_destroy(invert_params_standard);
zvec_index_params_destroy(invert_params_extended);
zvec_index_params_destroy(hnsw_params_fast);
zvec_index_params_destroy(hnsw_params_balanced);
zvec_index_params_destroy(hnsw_params_accurate);
zvec_index_params_destroy(flat_params_l2);
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(flat_params_cosine,
ZVEC_METRIC_TYPE_COSINE);
// 3. Create fields with different index types
printf("Creating fields with various index types...\n");
// Fields with inverted indexes
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
if (id_field) {
zvec_field_schema_set_index_params(id_field, invert_params_standard);
error = zvec_collection_schema_add_field(schema, id_field);
if (handle_error(error, "adding ID field") == ZVEC_OK) {
printf("✓ ID field with standard inverted index added\n");
}
}
zvec_field_schema_t *category_field =
zvec_field_schema_create("category", ZVEC_DATA_TYPE_STRING, true, 0);
if (category_field) {
zvec_field_schema_set_index_params(category_field, invert_params_extended);
error = zvec_collection_schema_add_field(schema, category_field);
if (handle_error(error, "adding category field") == ZVEC_OK) {
printf("✓ Category field with extended inverted index added\n");
}
}
// Vector fields with HNSW indexes (different configurations)
zvec_field_schema_t *fast_search_field = zvec_field_schema_create(
"fast_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 64);
if (fast_search_field) {
zvec_field_schema_set_index_params(fast_search_field, hnsw_params_fast);
error = zvec_collection_schema_add_field(schema, fast_search_field);
if (handle_error(error, "adding fast search field") == ZVEC_OK) {
printf("✓ Fast search vector field (64D) with HNSW index added\n");
}
}
zvec_field_schema_t *balanced_field = zvec_field_schema_create(
"balanced_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128);
if (balanced_field) {
zvec_field_schema_set_index_params(balanced_field, hnsw_params_balanced);
error = zvec_collection_schema_add_field(schema, balanced_field);
if (handle_error(error, "adding balanced field") == ZVEC_OK) {
printf("✓ Balanced vector field (128D) with HNSW index added\n");
}
}
zvec_field_schema_t *accurate_field = zvec_field_schema_create(
"accurate_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 256);
if (accurate_field) {
zvec_field_schema_set_index_params(accurate_field, hnsw_params_accurate);
error = zvec_collection_schema_add_field(schema, accurate_field);
if (handle_error(error, "adding accurate field") == ZVEC_OK) {
printf("✓ Accurate vector field (256D) with HNSW index added\n");
}
}
// Vector field with Flat index
zvec_field_schema_t *exact_field = zvec_field_schema_create(
"exact_vector", ZVEC_DATA_TYPE_VECTOR_FP32, false, 32);
if (exact_field) {
zvec_field_schema_set_index_params(exact_field, flat_params_l2);
error = zvec_collection_schema_add_field(schema, exact_field);
if (handle_error(error, "adding exact field") == ZVEC_OK) {
printf("✓ Exact search vector field (32D) with Flat index added\n");
}
}
// 4. Create collection
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
fprintf(stderr, "Failed to create collection options\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open("./index_example_collection", schema,
options, &collection);
zvec_collection_options_destroy(options);
if (handle_error(error, "creating collection") != ZVEC_OK) {
zvec_collection_schema_destroy(schema);
return -1;
}
printf("✓ Collection created successfully\n");
// 5. Create test data
printf("Creating test documents...\n");
zvec_doc_t *docs[3];
for (int i = 0; i < 3; i++) {
docs[i] = zvec_doc_create();
if (!docs[i]) {
fprintf(stderr, "Failed to create document %d\n", i);
// Cleanup
for (int j = 0; j < i; j++) {
zvec_doc_destroy(docs[j]);
}
goto cleanup;
}
}
// Prepare vector data
float fast_vec[3][64];
float balanced_vec[3][128];
float accurate_vec[3][256];
float exact_vec[3][32];
// Generate different vector patterns for testing
for (int doc_idx = 0; doc_idx < 3; doc_idx++) {
for (int i = 0; i < 64; i++) {
fast_vec[doc_idx][i] = (float)(doc_idx * 64 + i) / (64.0f * 3.0f);
}
for (int i = 0; i < 128; i++) {
balanced_vec[doc_idx][i] = (float)(doc_idx * 128 + i) / (128.0f * 3.0f);
}
for (int i = 0; i < 256; i++) {
accurate_vec[doc_idx][i] = (float)(doc_idx * 256 + i) / (256.0f * 3.0f);
}
for (int i = 0; i < 32; i++) {
exact_vec[doc_idx][i] = (float)(doc_idx * 32 + i) / (32.0f * 3.0f);
}
}
// Populate documents
for (int i = 0; i < 3; i++) {
char pk[16];
snprintf(pk, sizeof(pk), "doc%d", i + 1);
zvec_doc_set_pk(docs[i], pk);
char id_val[16];
snprintf(id_val, sizeof(id_val), "ID_%d", i + 1);
zvec_doc_add_field_by_value(docs[i], "id", ZVEC_DATA_TYPE_STRING, id_val,
strlen(id_val));
char category_val[16];
snprintf(category_val, sizeof(category_val), "cat_%d", (i % 2) + 1);
zvec_doc_add_field_by_value(docs[i], "category", ZVEC_DATA_TYPE_STRING,
category_val, strlen(category_val));
zvec_doc_add_field_by_value(docs[i], "fast_vector",
ZVEC_DATA_TYPE_VECTOR_FP32, fast_vec[i],
64 * sizeof(float));
zvec_doc_add_field_by_value(docs[i], "balanced_vector",
ZVEC_DATA_TYPE_VECTOR_FP32, balanced_vec[i],
128 * sizeof(float));
zvec_doc_add_field_by_value(docs[i], "accurate_vector",
ZVEC_DATA_TYPE_VECTOR_FP32, accurate_vec[i],
256 * sizeof(float));
zvec_doc_add_field_by_value(docs[i], "exact_vector",
ZVEC_DATA_TYPE_VECTOR_FP32, exact_vec[i],
32 * sizeof(float));
}
// 6. Insert documents
size_t success_count = 0, error_count = 0;
error = zvec_collection_insert(collection, (const zvec_doc_t **)docs, 3,
&success_count, &error_count);
if (handle_error(error, "inserting documents") == ZVEC_OK) {
printf("✓ Documents inserted - Success: %zu, Failed: %zu\n", success_count,
error_count);
}
// Cleanup documents
for (int i = 0; i < 3; i++) {
zvec_doc_destroy(docs[i]);
}
// 7. Flush collection to build indexes
error = zvec_collection_flush(collection);
if (handle_error(error, "flushing collection") == ZVEC_OK) {
printf("✓ Collection flushed - indexes built\n");
}
// 8. Test different query types
printf("Testing various index queries...\n");
// Test HNSW query (balanced)
zvec_vector_query_t *hnsw_query = zvec_vector_query_create();
if (!hnsw_query) {
fprintf(stderr, "Failed to create HNSW query\n");
goto cleanup;
}
zvec_vector_query_set_field_name(hnsw_query, "balanced_vector");
zvec_vector_query_set_query_vector(hnsw_query, balanced_vec[0],
128 * sizeof(float));
zvec_vector_query_set_topk(hnsw_query, 2);
zvec_vector_query_set_filter(hnsw_query, "");
zvec_vector_query_set_include_vector(hnsw_query, false);
zvec_vector_query_set_include_doc_id(hnsw_query, true);
zvec_doc_t **hnsw_results = NULL;
size_t hnsw_result_count = 0;
error =
zvec_collection_query(collection, (const zvec_vector_query_t *)hnsw_query,
&hnsw_results, &hnsw_result_count);
if (error == ZVEC_OK) {
printf("✓ HNSW query successful - Found %zu results\n", hnsw_result_count);
zvec_docs_free(hnsw_results, hnsw_result_count);
}
zvec_vector_query_destroy(hnsw_query);
// Test Flat query (exact)
zvec_vector_query_t *flat_query = zvec_vector_query_create();
if (!flat_query) {
fprintf(stderr, "Failed to create Flat query\n");
goto cleanup;
}
zvec_vector_query_set_field_name(flat_query, "exact_vector");
zvec_vector_query_set_query_vector(flat_query, exact_vec[0],
32 * sizeof(float));
zvec_vector_query_set_topk(flat_query, 2);
zvec_vector_query_set_filter(flat_query, "");
zvec_vector_query_set_include_vector(flat_query, false);
zvec_vector_query_set_include_doc_id(flat_query, true);
zvec_doc_t **flat_results = NULL;
size_t flat_result_count = 0;
error =
zvec_collection_query(collection, (const zvec_vector_query_t *)flat_query,
&flat_results, &flat_result_count);
if (error == ZVEC_OK) {
printf("✓ Flat (exact) query successful - Found %zu results\n",
flat_result_count);
zvec_docs_free(flat_results, flat_result_count);
}
zvec_vector_query_destroy(flat_query);
// 9. Performance comparison information
printf("\nIndex Performance Characteristics:\n");
printf("- Inverted Index: Fast text search, supports filtering\n");
printf(
"- HNSW Index: Approximate nearest neighbor search, good balance of "
"speed/accuracy\n");
printf("- Flat Index: Exact search, slower but 100%% accurate\n");
printf(
"- Trade-off: Speed vs Accuracy - choose based on your requirements\n");
// 10. Cleanup
cleanup:
zvec_collection_destroy(collection);
zvec_collection_schema_destroy(schema);
// Cleanup index parameters
printf("✓ Index example completed\n");
return 0;
}
examples/c/optimized_example.c
+310
View File
@@ -0,0 +1,310 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "zvec/c_api.h"
/**
* @brief Print error message and return error code
*/
static zvec_error_code_t handle_error(zvec_error_code_t error,
const char *context) {
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
fprintf(stderr, "Error in %s: %d - %s\n", context, error,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
}
return error;
}
/**
* @brief Create test vector data
*/
static float *create_test_vector(size_t dimension) {
float *vector = malloc(dimension * sizeof(float));
if (!vector) {
return NULL;
}
for (size_t i = 0; i < dimension; i++) {
vector[i] = (float)rand() / RAND_MAX;
}
return vector;
}
/**
* @brief Optimized C API usage example with performance considerations
*/
int main() {
printf("=== ZVec Optimized C API Example ===\n\n");
// Get version information
const char *version = zvec_get_version();
printf("ZVec Version: %s\n\n", version ? version : "Unknown");
zvec_error_code_t error;
// 1. Create optimized collection schema
zvec_collection_schema_t *schema =
zvec_collection_schema_create("optimized_example_collection");
if (!schema) {
fprintf(stderr, "Failed to create collection schema\n");
return -1;
}
printf("✓ Collection schema created\n");
// 2. Create optimized index parameters
zvec_index_params_t *hnsw_params =
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
if (!hnsw_params) {
fprintf(stderr, "Failed to create HNSW index parameters\n");
zvec_collection_schema_destroy(schema);
return -1;
}
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2);
zvec_index_params_set_hnsw_params(hnsw_params, 32, 200);
// 3. Create fields with optimized configuration
zvec_field_schema_t *id_field =
zvec_field_schema_create("id", ZVEC_DATA_TYPE_STRING, false, 0);
zvec_field_schema_t *text_field =
zvec_field_schema_create("text", ZVEC_DATA_TYPE_STRING, true, 0);
zvec_field_schema_t *embedding_field = zvec_field_schema_create(
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, 128);
if (!id_field || !text_field || !embedding_field) {
fprintf(stderr, "Failed to create field schemas\n");
goto cleanup_params;
}
// Set indexes
zvec_field_schema_set_index_params(embedding_field, hnsw_params);
// Add fields to schema
error = zvec_collection_schema_add_field(schema, id_field);
if (handle_error(error, "adding ID field") != ZVEC_OK) goto cleanup_fields;
error = zvec_collection_schema_add_field(schema, text_field);
if (handle_error(error, "adding text field") != ZVEC_OK) goto cleanup_fields;
error = zvec_collection_schema_add_field(schema, embedding_field);
if (handle_error(error, "adding embedding field") != ZVEC_OK)
goto cleanup_fields;
printf("✓ Fields configured with indexes\n");
// 4. Create collection with optimized options
zvec_collection_options_t *options = zvec_collection_options_create();
if (!options) {
fprintf(stderr, "Failed to create collection options\n");
goto cleanup_fields;
}
zvec_collection_options_set_enable_mmap(
options, true); // Enable memory mapping for better performance
zvec_collection_t *collection = NULL;
error = zvec_collection_create_and_open("./optimized_example_collection",
schema, options, &collection);
zvec_collection_options_destroy(options);
if (handle_error(error, "creating collection") != ZVEC_OK) {
goto cleanup_fields;
}
printf("✓ Collection created with optimized settings\n");
// 5. Bulk insert test data
const size_t DOC_COUNT = 1000;
const size_t BATCH_SIZE = 100;
printf("Inserting %zu documents in batches of %zu...\n", DOC_COUNT,
BATCH_SIZE);
clock_t start_time = clock();
for (size_t batch_start = 0; batch_start < DOC_COUNT;
batch_start += BATCH_SIZE) {
size_t current_batch_size = (batch_start + BATCH_SIZE > DOC_COUNT)
? DOC_COUNT - batch_start
: BATCH_SIZE;
zvec_doc_t **batch_docs = malloc(current_batch_size * sizeof(zvec_doc_t *));
if (!batch_docs) {
fprintf(stderr, "Failed to allocate batch documents\n");
break;
}
// Create batch documents
for (size_t i = 0; i < current_batch_size; i++) {
batch_docs[i] = zvec_doc_create();
if (!batch_docs[i]) {
fprintf(stderr, "Failed to create document\n");
// Cleanup previous documents in batch
for (size_t j = 0; j < i; j++) {
zvec_doc_destroy(batch_docs[j]);
}
free(batch_docs);
goto cleanup_collection;
}
size_t doc_id = batch_start + i;
char pk[32];
snprintf(pk, sizeof(pk), "doc_%zu", doc_id);
zvec_doc_set_pk(batch_docs[i], pk);
// Add ID field
char id_str[32];
snprintf(id_str, sizeof(id_str), "ID_%zu", doc_id);
zvec_doc_add_field_by_value(batch_docs[i], "id", ZVEC_DATA_TYPE_STRING,
id_str, strlen(id_str));
// Add text field
char text_str[64];
snprintf(text_str, sizeof(text_str),
"Document number %zu with sample text", doc_id);
zvec_doc_add_field_by_value(batch_docs[i], "text", ZVEC_DATA_TYPE_STRING,
text_str, strlen(text_str));
// Add vector field
float *vector = create_test_vector(128);
if (vector) {
zvec_doc_add_field_by_value(batch_docs[i], "embedding",
ZVEC_DATA_TYPE_VECTOR_FP32, vector,
128 * sizeof(float));
free(vector);
}
}
// Insert batch
size_t success_count, error_count;
error = zvec_collection_insert(collection, (const zvec_doc_t **)batch_docs,
current_batch_size, &success_count,
&error_count);
if (handle_error(error, "inserting batch") != ZVEC_OK) {
// Cleanup batch documents
for (size_t i = 0; i < current_batch_size; i++) {
zvec_doc_destroy(batch_docs[i]);
}
free(batch_docs);
goto cleanup_collection;
}
printf(" Batch %zu-%zu: %zu successful, %zu failed\n", batch_start,
batch_start + current_batch_size - 1, success_count, error_count);
// Cleanup batch documents
for (size_t i = 0; i < current_batch_size; i++) {
zvec_doc_destroy(batch_docs[i]);
}
free(batch_docs);
}
clock_t insert_end_time = clock();
double insert_time =
((double)(insert_end_time - start_time)) / CLOCKS_PER_SEC;
printf("✓ Bulk insertion completed in %.3f seconds (%.0f docs/sec)\n",
insert_time, DOC_COUNT / insert_time);
// 6. Flush and optimize collection
printf("Flushing and optimizing collection...\n");
zvec_collection_flush(collection);
zvec_collection_optimize(collection);
printf("✓ Collection optimized\n");
// 7. Performance query test
printf("Testing query performance...\n");
float *query_vector = create_test_vector(128);
if (!query_vector) {
fprintf(stderr, "Failed to create query vector\n");
goto cleanup_collection;
}
zvec_vector_query_t *query = zvec_vector_query_create();
if (!query) {
fprintf(stderr, "Failed to create vector query\n");
free(query_vector);
goto cleanup_collection;
}
zvec_vector_query_set_field_name(query, "embedding");
zvec_vector_query_set_query_vector(query, query_vector, 128 * sizeof(float));
zvec_vector_query_set_topk(query, 10);
zvec_vector_query_set_filter(query, "");
zvec_vector_query_set_include_vector(query, false);
zvec_vector_query_set_include_doc_id(query, true);
const int QUERY_COUNT = 100;
start_time = clock();
for (int q = 0; q < QUERY_COUNT; q++) {
zvec_doc_t **results = NULL;
size_t result_count = 0;
error =
zvec_collection_query(collection, (const zvec_vector_query_t *)query,
&results, &result_count);
if (error != ZVEC_OK) {
char *error_msg = NULL;
zvec_get_last_error(&error_msg);
printf("Query %d failed: %s\n", q,
error_msg ? error_msg : "Unknown error");
zvec_free(error_msg);
continue;
}
if (results) {
zvec_docs_free(results, result_count);
}
}
clock_t query_end_time = clock();
double query_time = ((double)(query_end_time - start_time)) / CLOCKS_PER_SEC;
double avg_query_time = (query_time * 1000) / QUERY_COUNT;
printf("✓ Performance test completed\n");
printf(" Average query time: %.2f ms\n", avg_query_time);
printf(" Queries per second: %.0f\n", 1000.0 / avg_query_time);
free(query_vector);
zvec_vector_query_destroy(query);
// 8. Memory usage information
zvec_collection_stats_t *stats = NULL;
error = zvec_collection_get_stats(collection, &stats);
if (error == ZVEC_OK && stats) {
printf("Collection Statistics:\n");
printf(" Document count: %llu\n",
(unsigned long long)zvec_collection_stats_get_doc_count(stats));
zvec_collection_stats_destroy(stats);
}
// 9. Cleanup
cleanup_collection:
zvec_collection_destroy(collection);
cleanup_fields:
// Field schemas are managed by the collection schema, no need to destroy
// individually
cleanup_params:
zvec_collection_schema_destroy(schema);
printf("✓ Optimized example completed\n");
return 0;
}
pyproject.toml
+299
View File
@@ -0,0 +1,299 @@
######################################################################################################
# Zvec: High-Performance Vector Database with PyBind11 & C++ Backend
######################################################################################################
[project]
name = "zvec"
dynamic = ["version"]
description = "A high-performance vector database engine with native C++ backend and Python bindings"
readme = "README.md"
license = { text = "Apache-2.0" }
authors = [
{ name = "zvec", email = "zvec@alibaba-inc.com" },
]
maintainers = [
{ name = "Zvec Core Team", email = "zvec@alibaba-inc.com" },
]
requires-python = ">=3.9" # NOTE: Only 64-bit Python interpreters are supported.
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
"Intended Audience :: Education",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: Apache Software License",
"Operating System :: POSIX :: Linux",
"Operating System :: MacOS",
"Programming Language :: C++",
"Programming Language :: Python :: 3 :: Only",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
"Topic :: Database",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development :: Libraries :: Python Modules",
]
keywords = [
"vector-database", "ann", "nearest-neighbor"
]
dependencies = [
"numpy >=1.23",
]
[project.urls]
Homepage = "https://github.com/alibaba/zvec"
Repository = "https://github.com/alibaba/zvec"
"Bug Tracker" = "https://github.com/alibaba/zvec/issues"
"Documentation" = "https://zvec.org"
[project.optional-dependencies]
test = [
"pytest >=8.0",
"pytest-cov >=4.1",
"pytest-mock >=3.12",
"pytest-xdist >=3.5",
"cibuildwheel == 3.4.0",
]
docs = [
"mkdocs >=1.5",
"mkdocs-material >=9.5",
"mkdocstrings[python] >=0.24",
]
dev = [
"ruff >=0.4",
"black >=24.0",
"mypy >=1.8",
"pre-commit >=3.6",
"build >=1.0",
"twine >=4.0",
"numpy >=1.23",
# Inherit test deps
"pytest >=8.0",
"pytest-cov >=4.1",
"pytest-mock >=3.12",
"cibuildwheel == 3.4.0",
# Inherit docs deps
"mkdocs >=1.5",
"mkdocs-material >=9.5",
"mkdocstrings[python] >=0.24",
"pybind11-stubgen>=2.5.5",
"pybind11 >=3.0",
]
######################################################################################################
# BUILD SYSTEM CONFIGURATION (scikit-build-core)
######################################################################################################
[build-system]
requires = [
"scikit-build-core >=0.11",
"pybind11 >=3.0",
"setuptools_scm>=8.0",
"cmake>=3.26,<4.0",
"ninja>=1.11",
]
build-backend = "scikit_build_core.build"
[tool.scikit-build]
# Core settings
minimum-version = "0.11"
metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
# CMake configuration
cmake.version = ">=3.26,<4.0"
ninja.version = ">=1.11"
cmake.build-type = "Release"
install.strip = true # Strip symbols in release builds to reduce wheel size
# Only pull the runtime artifacts tagged COMPONENT "python" into the wheel.
# This excludes the C++ SDK install rules (static/shared libs under lib/ and
# public headers under include/) that cc_library(PACKED) registers for a
# standalone `make install`, which would otherwise bloat the wheel.
install.components = ["python"]
# Build directory
build-dir = "build"
# Platform support
wheel.expand-macos-universal-tags = true
wheel.packages = ["python/zvec"]
# Source distribution
sdist.include = [
"README.md",
"LICENSE",
"pyproject.toml",
"CMakeLists.txt",
"src/**/*",
"stub/zvec/**/*",
"python/zvec/py.typed",
]
# CMake defines (env-overridable)
[tool.scikit-build.cmake.define]
BUILD_TOOLS = "OFF"
BUILD_PYTHON_BINDINGS = "ON"
#CMAKE_VERBOSE_MAKEFILE = "ON"
# Setuptools config for test pypi
[tool.setuptools_scm]
local_scheme = "no-local-version"
version_scheme = "guess-next-dev"
fallback_version = "0.0.0"
######################################################################################################
# TESTING & QUALITY
######################################################################################################
[tool.pytest.ini_options]
minversion = "8.0"
addopts = [
"-ra",
"--showlocals",
"--strict-markers",
"--strict-config",
"--tb=short",
"-n=auto",
]
xfail_strict = true
log_cli_level = "INFO"
filterwarnings = [
"error",
"ignore::pytest.PytestCacheWarning",
# Ignore numpy deprecation warnings in tests (if any)
"ignore:.*numpy.*:DeprecationWarning",
]
testpaths = ["python/tests"]
markers = [
"title: Custom marker for test title/description",
# "slow: marks tests as slow",
]
######################################################################################################
# BUILD WHEEL
######################################################################################################
[tool.cibuildwheel]
build = [
"cp310-*",
"cp311-*",
"cp312-*",
"cp313-*",
"cp314-*",
]
build-frontend = "build"
test-requires = ["pytest", "pytest-xdist", "numpy"]
test-command = "cd {project} && pytest python/tests -v --tb=short"
build-verbosity = 1
[tool.cibuildwheel.linux]
# libaio is required by the C++ backend; install it inside the manylinux
# container (manylinux_2_28 is AlmaLinux 8 based, so use dnf/libaio-devel).
# libaio-devel lives in BaseOS; disable EPEL so a flaky EPEL mirror cannot
# break metadata refresh.
before-all = "dnf install -y --disablerepo=epel libaio-devel"
archs = ["auto"]
environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" }
manylinux-x86_64-image = "manylinux_2_28"
manylinux-aarch64-image = "manylinux_2_28"
# Skip 32-bit builds and musllinux
skip = ["*-manylinux_i686", "*-musllinux*"]
[tool.cibuildwheel.macos]
archs = ["arm64"]
# Inherits CMAKE_GENERATOR and CMAKE_BUILD_PARALLEL_LEVEL from [tool.cibuildwheel] won't work;
# platform-level environment overrides the top-level entirely, so all vars must be listed here
environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16", MACOSX_DEPLOYMENT_TARGET = "11.0" }
[tool.cibuildwheel.windows]
archs = ["auto64"]
test-command = "cd /d {project} && pytest python/tests -v --tb=short --basetemp=./.pytest_tmp"
environment = { CMAKE_GENERATOR = "Ninja" }
#config-settings = { build-dir = "build/{wheel_tag}" }
#&& if exist .pytest_tmp rmdir /s /q .pytest_tmp"
before-build = [
"python -c \"import shutil, glob, os; [shutil.rmtree(p, ignore_errors=True) for p in ['build', '.pytest_tmp']];\""
]
######################################################################################################
# CODE QUALITY & FORMATTING (Ruff)
######################################################################################################
[tool.ruff]
target-version = "py310"
line-length = 88
exclude = [
"build/",
"dist/",
".git/",
".venv/",
"venv/",
"thirdparty",
]
[tool.ruff.lint]
extend-select = [
"B", # flake8-bugbear
"I", # isort
"ARG", # flake8-unused-arguments
"C4", # flake8-comprehensions
"EM", # flake8-errmsg
"ICN", # flake8-import-conventions
"G", # flake8-logging-format
"PGH", # pygrep-hooks
"PIE", # flake8-pie
"PL", # pylint
"PT", # flake8-pytest-style
"PTH", # flake8-use-pathlib
"RET", # flake8-return
"RUF", # Ruff-specific
"SIM", # flake8-simplify
"T20", # flake8-print
"UP", # pyupgrade
"YTT", # flake8-2020
"EXE", # flake8-executable
"NPY", # NumPy-specific
"PD", # pandas-vet
]
ignore = [
"PLR0913", # Too many arguments (common in bindings)
"PLR2004", # Magic value used in comparison
"UP045", "UP007", # Use list() instead of [] (breaks C++ init)
"EM101", "EM102", # Exception messages as literals (ok in tests/utils)
"B008", # Mutable default args (cautiously allowed in config)
"E731", # Lambda assignment (used in callbacks)
"B019", # `functools.lru_cache` on methods (handled manually)
"PLR0912", # Too many branches
"PLC0105", # Ignore contravariant
"RUF002", # Ignore Unicode
]
fixable = ["ALL"]
unfixable = []
# Ignore all errors in docstrings
[tool.ruff.lint.pydocstyle]
convention = "google" # or "numpy", "pep257"
ignore-decorators = ["typing.overload"]
[tool.ruff.lint.flake8-type-checking]
# Don't check code examples in docstrings
quote-annotations = true
[tool.ruff.lint.isort]
required-imports = ["from __future__ import annotations"]
known-first-party = ["zvec"]
[tool.ruff.lint.per-file-ignores]
"python/tests/**" = ["ALL"]
"bench/core/**" = ["ALL"]
"python/zvec/__init__.py" = [
"F401", # Unused import (for __all__)
"E402", # Module level import not at top (C++ module init order)
"PLE0605", # Invalid format for __all__
"RUF022", # __all__ is not sorted
]
"python/zvec/model/doc.py" = [
"RUF023", # Unused sort (for __slot__)
]
"python/zvec/extension/**" = [
"PLC0415", # Import outside top-level (dynamic imports in _get_model)
]
[tool.ruff.format]
indent-style = "space"
quote-style = "double"
line-ending = "lf"
skip-magic-trailing-comma = false
python/tests/detail/distance_helper.py
+392
View File
@@ -0,0 +1,392 @@
import logging
import math
import numpy as np
from zvec import (
MetricType,
DataType,
QuantizeType,
Doc,
CollectionSchema,
FieldSchema,
VectorSchema,
)
from typing import Dict
def is_float_equal(actual, expected, rel_tol=1e-5, abs_tol=1e-8):
if actual is None and expected is None:
return True
return math.isclose(actual, expected, rel_tol=rel_tol, abs_tol=abs_tol)
def is_dense_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):
"""Compare two dense vectors with tolerance."""
return np.allclose(vec1, vec2, rtol=rtol, atol=atol)
def is_sparse_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):
"""Compare two sparse vectors with tolerance."""
# Check if they have the same keys
if set(vec1.keys()) != set(vec2.keys()):
return False
# Check if all values are close
for key in vec1:
if not math.isclose(vec1[key], vec2[key], rel_tol=rtol, abs_tol=atol):
return False
return True
def is_float_array_equal(arr1, arr2, rtol=1e-5, atol=1e-8):
"""Compare two float arrays with tolerance."""
return np.allclose(arr1, arr2, rtol=rtol, atol=atol)
def is_double_array_equal(arr1, arr2, rtol=1e-9, atol=1e-12):
"""Compare two double arrays with tolerance."""
return np.allclose(arr1, arr2, rtol=rtol, atol=atol)
def is_int_array_equal(arr1, arr2):
"""Compare two integer arrays with exact equality."""
return np.array_equal(arr1, arr2)
def cosine_distance_dense(
vec1,
vec2,
dtype: DataType = DataType.VECTOR_FP32,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
# More stable conversion to float16 to avoid numerical issues
vec1 = [float(np.float16(a)) for a in vec1]
vec2 = [float(np.float16(b)) for b in vec2]
elif dtype == DataType.VECTOR_INT8:
# For INT8 vectors, convert to integers for proper calculation
vec1 = [
int(round(min(max(val, -128), 127))) for val in vec1
] # Clamp to valid INT8 range
vec2 = [
int(round(min(max(val, -128), 127))) for val in vec2
] # Clamp to valid INT8 range
dot_product = sum(a * b for a, b in zip(vec1, vec2))
magnitude1 = math.sqrt(sum(a * a for a in vec1))
magnitude2 = math.sqrt(sum(b * b for b in vec2))
if magnitude1 == 0 or magnitude2 == 0:
return 1.0 # Zero vector case - maximum distance
cosine_similarity = dot_product / (magnitude1 * magnitude2)
# Clamp to [-1, 1] range to handle floating-point precision errors
cosine_similarity = max(-1.0, min(1.0, cosine_similarity))
# For identical vectors (within floating point precision), ensure cosine distance is 0.0
# This is especially important for low-precision types which have limited precision
if (
dtype == DataType.VECTOR_FP16
or quantize_type == QuantizeType.FP16
or dtype == DataType.VECTOR_INT8
):
if (
abs(cosine_similarity - 1.0) < 1e-3
): # Handle precision issues for low-precision types
cosine_similarity = 1.0
# Return cosine distance (1 - cosine similarity) to maintain compatibility
# with system internal processing and existing test expectations
return 1.0 - cosine_similarity
def dp_distance_dense(
vec1,
vec2,
dtype: DataType = DataType.VECTOR_FP32,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
# More stable computation to avoid numerical issues
products = [
float(np.float16(a)) * float(np.float16(b)) for a, b in zip(vec1, vec2)
]
return sum(products)
elif dtype == DataType.VECTOR_INT8:
# For INT8 vectors, convert to integers for proper calculation
products = [
int(round(min(max(a, -128), 127))) * int(round(min(max(b, -128), 127)))
for a, b in zip(vec1, vec2)
]
return sum(products)
return sum(a * b for a, b in zip(vec1, vec2))
def euclidean_distance_dense(
vec1,
vec2,
dtype: DataType = DataType.VECTOR_FP32,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
# Convert to float16 and compute squared differences safely
# Use a more stable computation to avoid overflow
squared_diffs = []
for a, b in zip(vec1, vec2):
diff = np.float16(a) - np.float16(b)
squared_diff = float(diff) * float(
diff
) # Convert to float for multiplication
squared_diffs.append(squared_diff)
squared_distance = sum(squared_diffs)
elif dtype == DataType.VECTOR_INT8:
# For INT8 vectors, convert to integers and handle potential scaling
# INT8 values might be treated differently in the library implementation
vec1_int = [
int(round(min(max(val, -128), 127))) for val in vec1
] # Clamp to valid INT8 range
vec2_int = [
int(round(min(max(val, -128), 127))) for val in vec2
] # Clamp to valid INT8 range
# Use float type to prevent overflow when summing large squared differences
squared_distance = sum(float(a - b) ** 2 for a, b in zip(vec1_int, vec2_int))
else:
squared_distance = sum((a - b) ** 2 for a, b in zip(vec1, vec2))
return squared_distance # Return squared distance for INT8
def distance_dense(
vec1,
vec2,
metric: MetricType,
data_type: DataType = DataType.VECTOR_FP32,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
if metric == MetricType.COSINE:
return cosine_distance_dense(vec1, vec2, data_type, quantize_type)
elif metric == MetricType.L2:
return euclidean_distance_dense(vec1, vec2, data_type, quantize_type)
elif metric == MetricType.IP:
return dp_distance_dense(vec1, vec2, data_type, quantize_type)
else:
raise ValueError("Unsupported metric type")
def dp_distance_sparse(
vec1,
vec2,
data_type: DataType = DataType.SPARSE_VECTOR_FP32,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
dot_product = 0.0
for dim in set(vec1.keys()) & set(vec2.keys()):
print("dim,vec1,vec2:\n")
print(dim, vec1, vec2)
if (
data_type == DataType.SPARSE_VECTOR_FP16
or quantize_type == QuantizeType.FP16
):
vec1[dim] = np.float16(vec1[dim])
vec2[dim] = np.float16(vec2[dim])
dot_product += vec1[dim] * vec2[dim]
return dot_product
def distance(
vec1,
vec2,
metric: MetricType,
data_type: DataType,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
is_sparse = (
data_type == DataType.SPARSE_VECTOR_FP32
or data_type == DataType.SPARSE_VECTOR_FP16
)
if is_sparse:
if metric != MetricType.IP:
raise ValueError("Unsupported metric type for sparse vectors")
if is_sparse:
return dp_distance_sparse(vec1, vec2, data_type, quantize_type)
else:
return distance_dense(vec1, vec2, metric, data_type, quantize_type)
def distance_recall(
vec1,
vec2,
metric: MetricType,
data_type: DataType,
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
):
is_sparse = (
data_type == DataType.SPARSE_VECTOR_FP32
or data_type == DataType.SPARSE_VECTOR_FP16
)
if is_sparse:
return dp_distance_sparse(vec1, vec2, data_type, quantize_type)
else:
if data_type in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]:
return distance_dense(vec1, vec2, metric, data_type, quantize_type)
elif data_type in [DataType.VECTOR_INT8] and metric in [
MetricType.L2,
MetricType.IP,
]:
return distance_dense(vec1, vec2, metric, data_type, quantize_type)
else:
return dp_distance_dense(vec1, vec2, data_type, quantize_type)
def calculate_rrf_score(rank, k=60):
return 1.0 / (k + rank + 1)
def calculate_multi_vector_rrf_scores(query_results: Dict[str, Doc], k=60):
rrf_scores = {}
for vector_name, docs in query_results.items():
for rank, doc in enumerate(docs):
doc_id = doc.id
rrf_score = calculate_rrf_score(rank, k)
if doc_id in rrf_scores:
rrf_scores[doc_id] += rrf_score
else:
rrf_scores[doc_id] = rrf_score
return rrf_scores
def calculate_multi_vector_weighted_scores(
query_results: Dict[str, Doc], weights: Dict[str, float], metric: MetricType
):
def _normalize_score(score: float, metric: MetricType) -> float:
if metric == MetricType.L2:
return 1.0 - 2 * math.atan(score) / math.pi
if metric == MetricType.IP:
return 0.5 + math.atan(score) / math.pi
if metric == MetricType.COSINE:
return 1.0 - score / 2.0
raise ValueError("Unsupported metric type")
weighted_scores = {}
for vector_name, docs in query_results.items():
weight = weights.get(vector_name, 1.0)
for doc in docs:
doc_id = doc.id
weighted_score = (_normalize_score(doc.score, metric)) * weight
if doc_id in weighted_scores:
weighted_scores[doc_id] += weighted_score
else:
weighted_scores[doc_id] = weighted_score
return weighted_scores
def is_field_equal(field1, field2, schema: FieldSchema) -> bool:
if field1 is None and field2 is None:
return True
if field1 is None or field2 is None:
return False
if schema.data_type == DataType.ARRAY_FLOAT:
return is_float_array_equal(field1, field2)
elif schema.data_type == DataType.ARRAY_DOUBLE:
return is_double_array_equal(field1, field2)
elif schema.data_type in [
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_BOOL,
DataType.ARRAY_STRING,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_INT64,
]:
return is_int_array_equal(field1, field2)
elif schema.data_type in [DataType.FLOAT, DataType.DOUBLE]:
return is_float_equal(field1, field2)
return field1 == field2
def is_vector_equal(vec1, vec2, schema: VectorSchema) -> bool:
if (
schema.data_type == DataType.SPARSE_VECTOR_FP16
or schema.data_type == DataType.VECTOR_FP16
):
# skip fp16 vector equal
return True
is_sparse = (
schema.data_type == DataType.SPARSE_VECTOR_FP32
or schema.data_type == DataType.SPARSE_VECTOR_FP16
)
if is_sparse:
return is_sparse_vector_equal(vec1, vec2)
else:
return is_dense_vector_equal(vec1, vec2)
def is_doc_equal(
doc1: Doc,
doc2: Doc,
schema: CollectionSchema,
except_score: bool = True,
include_vector: bool = True,
):
if doc1.id != doc2.id:
logging.error("doc ids are not equal")
return False
reduce_field_names = set(doc1.field_names() + doc2.field_names())
reduce_vector_names = set(doc1.vector_names() + doc2.vector_names())
is_doc1_fields_empty = doc1.fields is None or doc1.fields == {}
is_doc2_fields_empty = doc2.fields is None or doc2.fields == {}
if is_doc1_fields_empty or is_doc2_fields_empty:
if is_doc1_fields_empty != is_doc2_fields_empty:
return False
else:
for field_name in reduce_field_names:
field_schema = schema.field(field_name)
if field_schema is None:
return False
if is_field_equal(
doc1.field(field_name), doc2.field(field_name), field_schema
):
continue
else:
logging.error(f"{field_name} are not equal")
return False
if include_vector:
is_doc1_vectors_empty = doc1.vectors is None or doc1.vectors == {}
is_doc2_vectors_empty = doc2.vectors is None or doc2.vectors == {}
if is_doc1_vectors_empty or is_doc2_vectors_empty:
if is_doc1_fields_empty != is_doc2_vectors_empty:
return False
else:
for vector_name in reduce_vector_names:
vector_schema = schema.vector(vector_name)
if vector_schema is None:
return False
if is_vector_equal(
doc1.vector(vector_name), doc2.vector(vector_name), vector_schema
):
continue
else:
return False
return True
python/tests/detail/doc_helper.py
+464
View File
@@ -0,0 +1,464 @@
from zvec import CollectionSchema, Doc
from support_helper import *
import numpy as np
from typing import Literal, Optional, Union, Tuple
import random
import string
import math
def generate_constant_vector(
i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32"
):
if dtype == "int8":
vec = [(i % 127)] * dimension
vec[i % dimension] = (i + 1) % 127
else:
base_val = (i % 1000) / 256.0
special_val = ((i + 1) % 1000) / 256.0
vec = [base_val] * dimension
vec[i % dimension] = special_val
return vec
def generate_constant_vector_recall(
i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32"
):
if dtype == "int8":
vec = [(i % 127)] * dimension
vec[i % dimension] = (i + 1) % 127
else:
base_val = math.sin((i) * 1000) / 256.0
special_val = math.sin((i + 1) * 1000) / 256.0
vec = [base_val] * dimension
vec[i % dimension] = special_val
return vec
def generate_sparse_vector(i: int):
return {i: i + 0.1}
def generate_vectordict(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = i % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
i, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
i, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
i,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_vectordict_recall(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = i % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector_recall(
i, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector_recall(
i, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector_recall(
i,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_vectordict_update(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = (i + 1) % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i + 1) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i + 1) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i + 1}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [(i + 1) % 2 == 0, (i + 1) % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 1.1), float(i + 2.1), float(i + 3.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 1.11), float(i + 2.11), float(i + 3.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i + 1}", f"test_{i + 2}", f"test_{i + 3}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
i + 1, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
i + 1, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
i + 1,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i + 1)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i + 1)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_doc(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_doc_recall(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict_recall(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_update_doc(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict_update(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_doc_random(i, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
random.seed(i)
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = random.choice([True, False])
elif field.data_type == DataType.INT32:
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = random.randint(0, 4294967295)
elif field.data_type == DataType.INT64:
doc_fields[field.name] = random.randint(
-9223372036854775808, 9223372036854775807
)
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = random.randint(0, 18446744073709551615)
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = random.uniform(
-1.7976931348623157e308, 1.7976931348623157e308
)
elif field.data_type == DataType.STRING:
length = random.randint(1, 999)
doc_fields[field.name] = "".join(
random.choices(string.ascii_letters + string.digits, k=length)
)
elif field.data_type == DataType.ARRAY_BOOL:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.choice([True, False]) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-2147483648, 2147483647) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 4294967295) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-9223372036854775808, 9223372036854775807)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 18446744073709551615) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_FLOAT:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_DOUBLE:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_STRING:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
"".join(
random.choices(
string.ascii_letters + string.digits, k=random.randint(1, 100)
)
)
for _ in range(array_length)
]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "int8"
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
doc = Doc(id=i, fields=doc_fields, vectors=doc_vectors)
return doc
def generate_vectordict_random(schema: CollectionSchema):
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = random.choice([True, False])
elif field.data_type == DataType.INT32:
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = random.randint(0, 4294967295)
elif field.data_type == DataType.INT64:
doc_fields[field.name] = random.randint(
-9223372036854775808, 9223372036854775807
)
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = random.randint(0, 18446744073709551615)
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = random.uniform(
-1.7976931348623157e308, 1.7976931348623157e308
)
elif field.data_type == DataType.STRING:
length = random.randint(1, 999)
doc_fields[field.name] = "".join(
random.choices(string.ascii_letters + string.digits, k=length)
)
elif field.data_type == DataType.ARRAY_BOOL:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.choice([True, False]) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-2147483648, 2147483647) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 4294967295) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-9223372036854775808, 9223372036854775807)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 18446744073709551615) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_FLOAT:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_DOUBLE:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_STRING:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
"".join(
random.choices(
string.ascii_letters + string.digits, k=random.randint(1, 100)
)
)
for _ in range(array_length)
]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "int8"
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
python/tests/detail/fixture_helper.py
+652
View File
@@ -0,0 +1,652 @@
import pytest
import logging
import platform
DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in (
"x86_64",
"AMD64",
"i686",
"i386",
)
from typing import Any, Generator
from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
import zvec
# Cache the DiskAnn plugin preload status so we pay the load cost once per
# test session. The plugin normally auto-loads on first DiskAnn use, but we
# preload it explicitly here so a missing libaio / misplaced plugin .so
# surfaces as a clear pytest skip instead of a confusing
# "Create vector column indexer failed" deep inside the collection code path.
_DISKANN_PRELOAD_REASON: str | None = None
_DISKANN_PRELOAD_DONE: bool = False
def _ensure_diskann_runtime_or_reason() -> str | None:
"""Preload the DiskAnn plugin and return None on success or a human-readable
skip reason on failure. Idempotent across calls."""
global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
if _DISKANN_PRELOAD_DONE:
return _DISKANN_PRELOAD_REASON
_DISKANN_PRELOAD_DONE = True
if not DISKANN_SUPPORTED:
_DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
return _DISKANN_PRELOAD_REASON
if not zvec.is_libaio_available():
_DISKANN_PRELOAD_REASON = (
"libaio is not available on this host; DiskAnn cannot run. "
"Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry."
)
return _DISKANN_PRELOAD_REASON
status = zvec.load_diskann_plugin()
if status != zvec.DISKANN_PLUGIN_OK:
_DISKANN_PRELOAD_REASON = (
f"Failed to load DiskAnn plugin (status={status}); "
"check that libzvec_diskann_plugin.so is installed alongside "
"_zvec.so in the Python site-packages directory."
)
return _DISKANN_PRELOAD_REASON
_DISKANN_PRELOAD_REASON = None
return None
from zvec import (
CollectionOption,
InvertIndexParam,
HnswIndexParam,
FlatIndexParam,
IVFIndexParam,
FieldSchema,
VectorSchema,
CollectionSchema,
Collection,
Doc,
Query,
)
from support_helper import *
@pytest.fixture(scope="session")
def basic_schema(collection_name="test_collection") -> CollectionSchema:
return CollectionSchema(
name=collection_name if len(collection_name) > 0 else "test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name", DataType.STRING, nullable=False, index_param=InvertIndexParam()
),
FieldSchema("weight", DataType.FLOAT, nullable=True),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
@pytest.fixture(scope="session")
def full_schema(
nullable: bool = False,
has_index: bool = False,
) -> CollectionSchema:
scalar_index_param = None
vector_index_param = None
if has_index:
scalar_index_param = InvertIndexParam(enable_range_optimization=True)
vector_index_param = HnswIndexParam()
fields = []
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
fields.append(
FieldSchema(
v,
k,
nullable=nullable,
index_param=scalar_index_param,
)
)
vetors = []
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
vetors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
return CollectionSchema(
name="full_collection",
fields=fields,
vectors=vetors,
)
@pytest.fixture(scope="function")
def full_schema_new(request) -> CollectionSchema:
if hasattr(request, "param"):
nullable, has_index, vector_index = request.param
else:
nullable, has_index, vector_index = True, False, HnswIndexParam()
# Skip DiskAnn tests on unsupported platforms or when the runtime cannot
# be brought up (missing libaio, plugin .so not installed, etc.).
from zvec.model.param import DiskAnnIndexParam
if isinstance(vector_index, DiskAnnIndexParam):
skip_reason = _ensure_diskann_runtime_or_reason()
if skip_reason is not None:
pytest.skip(skip_reason)
scalar_index_param = None
vector_index_param = None
if has_index:
scalar_index_param = InvertIndexParam(enable_range_optimization=True)
vector_index_param = vector_index
fields = []
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
fields.append(
FieldSchema(
v,
k,
nullable=nullable,
index_param=scalar_index_param,
)
)
vectors = []
if vector_index_param in [
HnswIndexParam(),
FlatIndexParam(),
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
),
FlatIndexParam(
metric_type=MetricType.IP,
),
]:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
elif vector_index_param in [
IVFIndexParam(),
IVFIndexParam(
metric_type=MetricType.IP,
n_list=100,
n_iters=10,
use_soar=False,
),
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
(
IVFIndexParam(
metric_type=MetricType.COSINE,
n_list=150,
n_iters=15,
use_soar=False,
)
),
(
HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
)
),
(
HnswIndexParam(
metric_type=MetricType.L2,
m=32,
ef_construction=200,
)
),
(
FlatIndexParam(
metric_type=MetricType.COSINE,
)
),
(
FlatIndexParam(
metric_type=MetricType.L2,
)
),
]:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
if v in ["vector_fp16_field", "vector_fp32_field"]:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
elif v in ["vector_int8_field"] and vector_index_param in [
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
(
HnswIndexParam(
metric_type=MetricType.L2,
m=32,
ef_construction=200,
)
),
(
FlatIndexParam(
metric_type=MetricType.L2,
)
),
]:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
else:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=HnswIndexParam(),
)
)
else:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
if v in ["vector_fp16_field", "vector_fp32_field"]:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
else:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=HnswIndexParam(),
)
)
return CollectionSchema(
name="full_collection_new",
fields=fields,
vectors=vectors,
)
@pytest.fixture(scope="function")
def full_schema_ivf(request) -> CollectionSchema:
if hasattr(request, "param"):
nullable, has_index, vector_index = request.param
else:
nullable, has_index, vector_index = True, False, IVFIndexParam()
scalar_index_param = None
vector_index_param = None
if has_index:
scalar_index_param = InvertIndexParam(enable_range_optimization=True)
vector_index_param = vector_index
fields = []
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
fields.append(
FieldSchema(
v,
k,
nullable=nullable,
index_param=scalar_index_param,
)
)
vectors = []
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
if v in ["vector_fp16_field", "vector_fp32_field"]:
vectors.append(
VectorSchema(
v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
return CollectionSchema(
name="full_collection_ivf",
fields=fields,
vectors=vectors,
)
@pytest.fixture(scope="function")
def full_schema_1024(request) -> CollectionSchema:
if hasattr(request, "param"):
nullable, has_index, vector_index = request.param
else:
nullable, has_index, vector_index = True, False, HnswIndexParam()
scalar_index_param = None
vector_index_param = None
if has_index:
scalar_index_param = InvertIndexParam(enable_range_optimization=True)
vector_index_param = vector_index
fields = []
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
fields.append(
FieldSchema(
v,
k,
nullable=nullable,
index_param=scalar_index_param,
)
)
vectors = []
if vector_index_param in [
HnswIndexParam(),
FlatIndexParam(),
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
),
FlatIndexParam(
metric_type=MetricType.IP,
),
]:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
vectors.append(
VectorSchema(
v,
k,
dimension=VECTOR_DIMENSION_1024,
index_param=vector_index_param,
)
)
elif vector_index_param in [
IVFIndexParam(),
IVFIndexParam(
metric_type=MetricType.IP,
n_list=100,
n_iters=10,
use_soar=False,
),
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
IVFIndexParam(
metric_type=MetricType.COSINE,
n_list=150,
n_iters=15,
use_soar=False,
),
]:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
if v in ["vector_fp16_field", "vector_fp32_field"]:
vectors.append(
VectorSchema(
v,
k,
dimension=VECTOR_DIMENSION_1024,
index_param=vector_index_param,
)
)
elif v in ["vector_int8_field"] and vector_index_param in [
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
IVFIndexParam(
metric_type=MetricType.COSINE,
n_list=150,
n_iters=15,
use_soar=False,
),
]:
vectors.append(
VectorSchema(
v,
k,
dimension=DVECTOR_DIMENSION_1024,
index_param=vector_index_param,
)
)
else:
vectors.append(
VectorSchema(
v,
k,
dimension=VECTOR_DIMENSION_1024,
index_param=HnswIndexParam(),
)
)
else:
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
if v in ["vector_fp16_field", "vector_fp32_field", "vector_int8_field"]:
vectors.append(
VectorSchema(
v,
k,
dimension=VECTOR_DIMENSION_1024,
index_param=vector_index_param,
)
)
else:
vectors.append(
VectorSchema(
v,
k,
dimension=VECTOR_DIMENSION_1024,
index_param=HnswIndexParam(),
)
)
return CollectionSchema(
name="full_collection_new",
fields=fields,
vectors=vectors,
)
@pytest.fixture(scope="function")
def single_vector_schema(
data_type: DataType,
) -> CollectionSchema:
vector_schema = [
VectorSchema(
DEFAULT_VECTOR_FIELD_NAME[data_type],
data_type,
DEFAULT_VECTOR_DIMENSION,
)
]
return CollectionSchema(
name="full_collection",
vectors=vector_schema,
)
@pytest.fixture(scope="function")
def single_vector_schema_with_index_param(
data_type: DataType, index_param
) -> CollectionSchema:
vector_schema = [
VectorSchema(
DEFAULT_VECTOR_FIELD_NAME[data_type],
data_type,
DEFAULT_VECTOR_DIMENSION,
index_param,
)
]
return CollectionSchema(
name="full_collection",
vectors=vector_schema,
)
def create_collection_fixture(
collection_temp_dir, schema: CollectionSchema, collection_option: CollectionOption
) -> Generator[Any, Any, Collection]:
"""Common helper function to create and manage collection fixtures."""
coll = zvec.create_and_open(
path=str(collection_temp_dir),
schema=schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open collection"
assert coll.path == str(collection_temp_dir)
assert coll.schema.name == schema.name
assert list(coll.schema.fields) == list(schema.fields)
assert list(coll.schema.vectors) == list(schema.vectors)
assert coll.option.read_only == collection_option.read_only
assert coll.option.enable_mmap == collection_option.enable_mmap
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
logging.warning(f"Warning: failed to destroy collection: {e}")
@pytest.fixture(scope="function")
def basic_collection(
collection_temp_dir, basic_schema, collection_option
) -> Generator[Any, Any, Collection]:
yield from create_collection_fixture(
collection_temp_dir, basic_schema, collection_option
)
@pytest.fixture(scope="function")
def collection_option():
return CollectionOption(read_only=False, enable_mmap=True)
@pytest.fixture(scope="function")
def collection_temp_dir(tmp_path_factory):
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection_path"
return str(collection_path)
@pytest.fixture(scope="function")
def full_collection(
collection_temp_dir,
full_schema,
collection_option,
nullable: bool = True,
has_index: bool = False,
) -> Generator[Any, Any, Collection]:
yield from create_collection_fixture(
collection_temp_dir, full_schema, collection_option
)
@pytest.fixture(scope="function")
def full_collection_new(
collection_temp_dir, full_schema_new, collection_option
) -> Generator[Any, Any, Collection]:
yield from create_collection_fixture(
collection_temp_dir, full_schema_new, collection_option
)
@pytest.fixture(scope="function")
def full_collection_ivf(
collection_temp_dir, full_schema_ivf, collection_option
) -> Generator[Any, Any, Collection]:
yield from create_collection_fixture(
collection_temp_dir, full_schema_ivf, collection_option
)
@pytest.fixture(scope="function")
def full_collection_1024(
collection_temp_dir, full_schema_1024, collection_option
) -> Generator[Any, Any, Collection]:
yield from create_collection_fixture(
collection_temp_dir, full_schema_1024, collection_option
)
@pytest.fixture
def sample_field_list(nullable: bool = True, scalar_index_param=None, name_prefix=""):
field_list = []
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
field_list.append(
FieldSchema(
f"{name_prefix}_{v}" if len(name_prefix) > 0 else v,
k,
nullable=nullable,
index_param=scalar_index_param,
)
)
return field_list
@pytest.fixture
def sample_vector_list(vector_index_param=None, name_prefix=""):
vector_list = []
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
vector_list.append(
VectorSchema(
f"{name_prefix}_{v}" if len(name_prefix) > 0 else v,
k,
dimension=DEFAULT_VECTOR_DIMENSION,
index_param=vector_index_param,
)
)
return vector_list
python/tests/detail/params_helper.py
+209
View File
@@ -0,0 +1,209 @@
from zvec import (
CollectionOption,
IndexOption,
OptimizeOption,
InvertIndexParam,
HnswIndexParam,
IVFIndexParam,
FlatIndexParam,
AlterColumnOption,
AddColumnOption,
DataType,
MetricType,
QuantizeType,
)
VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {
DataType.VECTOR_FP32: [
HnswIndexParam(),
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
quantize_type=QuantizeType.INT8,
),
HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
quantize_type=QuantizeType.INT4,
),
HnswIndexParam(
metric_type=MetricType.L2,
m=32,
ef_construction=200,
quantize_type=QuantizeType.FP16,
),
FlatIndexParam(),
FlatIndexParam(metric_type=MetricType.IP, quantize_type=QuantizeType.INT4),
FlatIndexParam(metric_type=MetricType.L2, quantize_type=QuantizeType.INT8),
FlatIndexParam(metric_type=MetricType.COSINE, quantize_type=QuantizeType.FP16),
IVFIndexParam(),
IVFIndexParam(
metric_type=MetricType.IP,
quantize_type=QuantizeType.INT4,
n_list=100,
n_iters=10,
use_soar=False,
),
IVFIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
n_list=200,
n_iters=20,
use_soar=True,
),
IVFIndexParam(
metric_type=MetricType.COSINE,
quantize_type=QuantizeType.FP16,
n_list=150,
n_iters=15,
use_soar=False,
),
],
DataType.VECTOR_FP16: [
HnswIndexParam(),
FlatIndexParam(),
# IVFIndexParam(),
],
DataType.VECTOR_INT8: [
HnswIndexParam(),
FlatIndexParam(),
# IVFIndexParam(),
],
DataType.SPARSE_VECTOR_FP32: [
HnswIndexParam(),
FlatIndexParam(),
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
quantize_type=QuantizeType.FP16,
),
],
DataType.SPARSE_VECTOR_FP16: [
HnswIndexParam(),
FlatIndexParam(),
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
),
],
}
VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [
(data_type, param)
for data_type, params in VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()
for param in params
]
INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {
DataType.VECTOR_FP32: [
InvertIndexParam(),
],
DataType.VECTOR_FP16: [
InvertIndexParam(),
],
DataType.VECTOR_INT8: [
InvertIndexParam(),
],
DataType.SPARSE_VECTOR_FP32: [
HnswIndexParam(metric_type=MetricType.L2),
FlatIndexParam(metric_type=MetricType.COSINE),
IVFIndexParam(),
InvertIndexParam(),
],
DataType.SPARSE_VECTOR_FP16: [
HnswIndexParam(metric_type=MetricType.L2),
FlatIndexParam(metric_type=MetricType.COSINE),
IVFIndexParam(),
InvertIndexParam(),
],
}
INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [
(data_type, param)
for data_type, params in INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()
for param in params
]
COLLECTION_NAME_MAX_LENGTH = 64
COLLECTION_NAME_VALID_LIST = [
"col",
"C0llECTION",
"Collection1",
"collection_2",
"123collection-",
"a" * COLLECTION_NAME_MAX_LENGTH,
]
COLLECTION_NAME_INVALID_LIST = [
"l",
"1C",
"",
" ",
None,
"abcdefghijklmnopqrstuvwxzy123456abcdefghijklmnopqrstuvwxzy1234561",
"test/",
"!@#$%^&*()test",
]
FIELD_NAME_VALID_LIST = [
"1",
"12",
"col",
"ID",
"name1",
"Weigt_12-",
"123age",
"name_with_underscores",
"123numeric_start",
"name-with-dashes",
]
FIELD_NAME_INVALID_LIST = [
"",
" ",
None,
"abcdefghijklmnopqrstuvwxzy1234561",
"test/",
"!@#$%^&*()test",
"name@with#special$chars",
"name with spaces",
]
FIELD_LIST_MAX_LENGTH = 1024
VECTOR_LIST_MAX_LENGTH = 5
DENSE_VECTOR_MAX_DIMENSION = 20000
SPARSE_VECTOR_MAX_DIMENSION = 4096
FIELD_VECTOR_LIST_DIMENSION_VALID_LIST = [
# field_list_len, vector_list_len, dimension
(1, 1, 1),
(2, 2, 512),
(512, 3, 1024),
(1024, 4, 20000),
]
FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST = [
# field_list_len, vector_list_len, dimension
(1, 1, 0),
(1, 1, -1),
(1, 1, "1"),
(1, 1, 20001),
]
INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG = "incompatible constructor arguments"
SCHEMA_VALIDATE_ERROR_MSG = "schema validate failed"
CREATE_READ_ONLY_ERROR_MSG = "Unable to create collection with read-only mode"
INCOMPATIBLE_FUNCTION_ERROR_MSG = "incompatible function arguments"
INVALID_PATH_ERROR_MSG = "path validate failed"
INDEX_NON_EXISTENT_COLUMN_ERROR_MSG = "not found in schema"
ACCESS_DESTROYED_COLLECTION_ERROR_MSG = "is already destroyed"
COLLECTION_PATH_NOT_EXIST_ERROR_MSG = "not exist"
NOT_SUPPORT_ADD_COLUMN_ERROR_MSG = "Only support basic numeric data type"
NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG = "Column not exists"
python/tests/detail/support_helper.py
+126
View File
@@ -0,0 +1,126 @@
from zvec import (
CollectionOption,
IndexOption,
OptimizeOption,
InvertIndexParam,
HnswIndexParam,
IVFIndexParam,
FlatIndexParam,
DataType,
IndexType,
QuantizeType,
)
SUPPORT_SCALAR_DATA_TYPES = [
DataType.BOOL,
DataType.FLOAT,
DataType.DOUBLE,
DataType.INT32,
DataType.INT64,
DataType.UINT32,
DataType.UINT64,
DataType.STRING,
DataType.ARRAY_BOOL,
DataType.ARRAY_FLOAT,
DataType.ARRAY_DOUBLE,
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_STRING,
]
DEFAULT_SCALAR_FIELD_NAME = {
DataType.BOOL: "bool_field",
DataType.FLOAT: "float_field",
DataType.DOUBLE: "double_field",
DataType.INT32: "int32_field",
DataType.INT64: "int64_field",
DataType.UINT32: "uint32_field",
DataType.UINT64: "uint64_field",
DataType.STRING: "string_field",
DataType.ARRAY_BOOL: "array_bool_field",
DataType.ARRAY_FLOAT: "array_float_field",
DataType.ARRAY_DOUBLE: "array_double_field",
DataType.ARRAY_INT32: "array_int32_field",
DataType.ARRAY_INT64: "array_int64_field",
DataType.ARRAY_UINT32: "array_uint32_field",
DataType.ARRAY_UINT64: "array_uint64_field",
DataType.ARRAY_STRING: "array_string_field",
}
SUPPORT_SCALAR_INDEX_TYPES = [
IndexType.INVERT,
]
SUPPORT_VECTOR_DATA_TYPES = [
DataType.VECTOR_FP16,
DataType.VECTOR_FP32,
DataType.VECTOR_INT8,
DataType.SPARSE_VECTOR_FP32,
DataType.SPARSE_VECTOR_FP16,
]
SUPPORT_VECTOR_INDEX_TYPES = [
IndexType.FLAT,
IndexType.HNSW,
IndexType.IVF,
]
DEFAULT_VECTOR_FIELD_NAME = {
DataType.VECTOR_FP16: "vector_fp16_field",
DataType.VECTOR_FP32: "vector_fp32_field",
DataType.VECTOR_INT8: "vector_int8_field",
DataType.SPARSE_VECTOR_FP32: "sparse_vector_fp32_field",
DataType.SPARSE_VECTOR_FP16: "sparse_vector_fp16_field",
}
DEFAULT_VECTOR_DIMENSION = 128
VECTOR_DIMENSION_1024 = 4
SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP = {
DataType.VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],
DataType.VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],
DataType.VECTOR_INT8: [IndexType.FLAT, IndexType.HNSW],
DataType.SPARSE_VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW],
DataType.SPARSE_VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW],
}
SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS = [
(data_type, index_type)
for data_type, index_types in SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP.items()
for index_type in index_types
]
DEFAULT_INDEX_PARAMS = {
IndexType.FLAT: FlatIndexParam(),
IndexType.HNSW: HnswIndexParam(),
IndexType.IVF: IVFIndexParam(),
IndexType.INVERT: InvertIndexParam(),
}
SUPPORT_VECTOR_DATA_TYPE_QUANT_MAP = {
DataType.VECTOR_FP32: [QuantizeType.FP16, QuantizeType.INT8, QuantizeType.INT4],
DataType.SPARSE_VECTOR_FP32: [QuantizeType.FP16],
}
SUPPORT_ADD_COLUMN_DATA_TYPE = [
DataType.INT32,
DataType.UINT32,
DataType.INT64,
DataType.UINT64,
DataType.FLOAT,
DataType.DOUBLE,
]
NOT_SUPPORT_ADD_COLUMN_DATA_TYPE = [
DataType.BOOL,
DataType.STRING,
DataType.ARRAY_BOOL,
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_FLOAT,
DataType.ARRAY_DOUBLE,
DataType.ARRAY_STRING,
]
python/tests/detail/test_collection_concurrency.py
+429
View File
@@ -0,0 +1,429 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import pytest
import threading
import numpy as np
import zvec
from zvec import (
CollectionOption,
InvertIndexParam,
HnswIndexParam,
Collection,
Doc,
DataType,
FieldSchema,
VectorSchema,
)
class TestCollectionConcurrency:
@pytest.fixture(scope="function")
def test_collection(self, tmp_path_factory):
"""Fixture to create a test collection"""
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
FieldSchema("weight", DataType.FLOAT, nullable=True),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open collection"
yield coll
# Clean up
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_concurrent_read_write(self, test_collection: Collection):
results = []
def insert_docs(thread_id):
try:
docs = [
Doc(
id=f"{thread_id}_{i}",
fields={
"id": int(f"{thread_id}{i}"),
"name": f"thread_{thread_id}_doc_{i}",
"weight": float(i),
},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i), 2: float(i * 2)},
},
)
for i in range(5)
]
result = test_collection.insert(docs)
results.append((thread_id, "insert", len(result)))
except Exception as e:
results.append((thread_id, "insert_exception", str(e)))
def query_docs(thread_id):
try:
result = test_collection.query(filter="id > 0", topk=10)
results.append((thread_id, "query", len(result)))
except Exception as e:
results.append((thread_id, "query_exception", str(e)))
# Create threads for concurrent operations
threads = []
# Start insert threads
for i in range(3):
thread = threading.Thread(target=insert_docs, args=(i,))
threads.append(thread)
thread.start()
# Start query threads
for i in range(3):
thread = threading.Thread(target=query_docs, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Analyze results
insert_results = [r for r in results if r[1] == "insert"]
query_results = [r for r in results if r[1] == "query"]
logging.info(
f"Concurrent read/write results - Inserts: {len(insert_results)}, Queries: {len(query_results)}"
)
# At least some operations should succeed
assert len(insert_results) + len(query_results) > 0
def test_concurrent_query(self, test_collection: Collection):
# First insert some data
docs = [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i), 2: float(i * 2)},
},
)
for i in range(20)
]
insert_result = test_collection.insert(docs)
assert len(insert_result) == 20
results = []
def query_operation(thread_id):
"""Perform query operation from a thread"""
try:
result = test_collection.query(filter=f"id > {thread_id}", topk=5)
results.append((thread_id, "query", len(result)))
except Exception as e:
results.append((thread_id, "query_exception", str(e)))
# Create multiple threads for concurrent queries
threads = []
for i in range(5):
thread = threading.Thread(target=query_operation, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Analyze results
query_results = [r for r in results if r[1] == "query"]
logging.info(f"Concurrent query results - Queries: {len(query_results)}")
# All query operations should succeed
assert len(query_results) == 5
def test_concurrent_modifications(self, test_collection: Collection):
# First insert some data
docs = [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i), 2: float(i * 2)},
},
)
for i in range(10)
]
insert_result = test_collection.insert(docs)
assert len(insert_result) == 10
results = []
def update_operation(thread_id):
"""Perform update operation from a thread"""
try:
# Each thread updates different documents
update_docs = [
Doc(
id=f"{i}",
fields={
"id": i,
"name": f"updated_by_thread_{thread_id}",
"weight": float(i + thread_id),
},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i) + 0.5, 2: float(i * 2) + 0.5},
},
)
for i in range(thread_id * 2, thread_id * 2 + 2)
]
result = test_collection.update(update_docs)
results.append((thread_id, "update", len(result)))
except Exception as e:
results.append((thread_id, "update_exception", str(e)))
def delete_operation(thread_id):
"""Perform delete operation from a thread"""
try:
# Each thread deletes different documents
delete_ids = [f"{thread_id * 2 + 2}", f"{thread_id * 2 + 3}"]
result = test_collection.delete(delete_ids)
results.append((thread_id, "delete", len(result)))
except Exception as e:
results.append((thread_id, "delete_exception", str(e)))
# Create threads for concurrent operations
threads = []
# Start update threads
for i in range(3):
thread = threading.Thread(target=update_operation, args=(i,))
threads.append(thread)
thread.start()
# Start delete threads
for i in range(2):
thread = threading.Thread(target=delete_operation, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Analyze results
update_results = [r for r in results if r[1] == "update"]
delete_results = [r for r in results if r[1] == "delete"]
logging.info(
f"Concurrent modification results - Updates: {len(update_results)}, Deletes: {len(delete_results)}"
)
# At least some operations should succeed
assert len(update_results) + len(delete_results) > 0
def test_read_write_locking(self, test_collection: Collection):
# Perform operations that should be thread-safe
docs = [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i), 2: float(i * 2)},
},
)
for i in range(5)
]
# Insert data
insert_result = test_collection.insert(docs)
assert len(insert_result) == 5
# Concurrent operations should not cause data corruption
results = []
def mixed_operation(thread_id):
"""Perform mixed operations from a thread"""
try:
# Mix of read and write operations
if thread_id % 2 == 0:
# Read operation
result = test_collection.fetch([f"{thread_id % 5}"])
results.append((thread_id, "read", len(result)))
else:
# Write operation
doc = Doc(
id=f"{thread_id % 5}",
fields={
"id": thread_id % 5,
"name": f"mixed_op_{thread_id}",
"weight": float(thread_id),
},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(thread_id), 2: float(thread_id * 2)},
},
)
result = test_collection.upsert(doc)
results.append((thread_id, "write", len(result)))
except Exception as e:
results.append((thread_id, "exception", str(e)))
# Create multiple threads
threads = []
for i in range(10):
thread = threading.Thread(target=mixed_operation, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Verify that the collection is still in a consistent state
final_result = test_collection.query()
assert len(final_result) >= 0 # Should not crash or return corrupted data
def test_race_condition_detection(self, test_collection: Collection):
# Insert initial data
docs = [
Doc(
id=f"{i}",
fields={"id": i, "name": f"initial_{i}", "weight": float(i)},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(i), 2: float(i * 2)},
},
)
for i in range(10)
]
insert_result = test_collection.insert(docs)
assert len(insert_result) == 10
# Perform many rapid concurrent operations
operation_count = 100
results = []
def rapid_operation(op_id):
"""Perform rapid operations"""
try:
# Alternate between different types of operations
if op_id % 4 == 0:
# Insert
doc = Doc(
id=f"rapid_{op_id}",
fields={
"id": op_id,
"name": f"rapid_{op_id}",
"weight": float(op_id),
},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(op_id), 2: float(op_id * 2)},
},
)
result = test_collection.insert(doc)
results.append(("insert", len(result)))
elif op_id % 4 == 1:
# Update
doc = Doc(
id=f"{op_id % 10}",
fields={
"id": op_id % 10,
"name": f"rapid_update_{op_id}",
"weight": float(op_id),
},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: float(op_id), 2: float(op_id * 2)},
},
)
result = test_collection.update(doc)
results.append(("update", len(result)))
elif op_id % 4 == 2:
# Query
result = test_collection.query(filter=f"id > {op_id % 5}", topk=3)
results.append(("query", len(result)))
else:
# Fetch
result = test_collection.fetch([f"{op_id % 10}"])
results.append(("fetch", len(result)))
except Exception as e:
results.append(("exception", str(e)))
# Create many threads for rapid concurrent operations
threads = []
for i in range(operation_count):
thread = threading.Thread(target=rapid_operation, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Verify collection is still functional
final_query = test_collection.query()
assert len(final_query) >= 0 # Should not be corrupted
logging.info(
f"Rapid concurrent operations completed - Total operations: {len(results)}"
)
python/tests/detail/test_collection_create_and_open.py
+791
View File
@@ -0,0 +1,791 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import threading
import os
from distance_helper import *
from fixture_helper import *
from doc_helper import *
from params_helper import *
def check_collection_info(
coll: Collection, schema: CollectionSchema, option: CollectionOption, path: str
):
assert coll is not None, "Failed to create and open collection"
assert coll.path == path
assert coll.schema.name == schema.name
assert list(coll.schema.fields) == list(schema.fields)
assert list(coll.schema.vectors) == list(schema.vectors)
assert coll.option.read_only == option.read_only
assert coll.option.enable_mmap == option.enable_mmap
def check_collection_basic(coll: Collection, optimize: bool = False):
schema = coll.schema
docs = [generate_doc(i, schema) for i in range(10)]
results = coll.insert(docs=docs)
assert len(results) == len(docs)
for result in results:
assert result.ok()
assert coll.stats.doc_count == len(docs)
def check_fetch_query():
results = coll.fetch([str(i) for i in range(len(docs))])
assert len(results) == len(docs)
for i in range(len(docs)):
assert str(i) in results
results = coll.query()
assert len(results) == len(docs)
check_fetch_query()
if optimize:
coll.optimize()
check_fetch_query()
def check_collection_full(coll: Collection):
test_doc = generate_doc(1, coll.schema)
insert_result = coll.insert(test_doc)
assert insert_result.ok()
stats = coll.stats
assert stats.doc_count == 1
fetched_docs = coll.fetch(ids=["1"])
assert len(fetched_docs) == 1
assert "1" in fetched_docs
assert fetched_docs["1"] is not None
assert is_doc_equal(fetched_docs["1"], test_doc, coll.schema)
query_result = coll.query()
assert len(query_result) == 1
updated_doc = Doc(
id="1",
fields={"int32_field": 1},
vectors={"vector_fp32_field": [0.2] * 128},
)
update_result = coll.update(updated_doc)
assert update_result.ok()
upserted_doc = generate_doc(1, coll.schema)
upsert_result = coll.upsert(upserted_doc)
assert upsert_result.ok()
# 8. Delete document
delete_result = coll.delete("1")
assert delete_result.ok()
# Verify document was deleted
stats = coll.stats
assert stats.doc_count == 0
valid_collection_options = [
# (read_only, enable_mmap)
(False, True),
(False, False),
]
invalid_collection_options = [
# (read_only, enable_mmap)
(True, True),
(True, False),
]
duplicate_names_test = [
("field1", "field1", "vector1", "vector2"),
("field1", "field2", "vector1", "vector1"),
(
"shared_name1",
"shared_name2",
"shared_name1",
"shared_name2",
),
]
long_names = [
"a" * 100, # 100 characters
"b" * 200, # 200 characters
]
valid_path_list = [
"/tmp/nonexistent/directory/test_collection",
"test/collection/with/slashes",
"test/collection/with/slashes/哈哈",
]
invalid_path_list = [
"invalid\0path",
"",
]
class TestCreateAndOpen:
@pytest.mark.parametrize("collection_name", COLLECTION_NAME_VALID_LIST)
def test_valid_collection_name(
self,
collection_temp_dir,
collection_name,
collection_option,
sample_field_list,
sample_vector_list,
):
collection_schema = zvec.CollectionSchema(
name=collection_name,
fields=sample_field_list,
vectors=sample_vector_list,
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
check_collection_info(
coll, collection_schema, collection_option, collection_temp_dir
)
check_collection_basic(coll)
coll.destroy()
@pytest.mark.parametrize("collection_name", COLLECTION_NAME_INVALID_LIST)
def test_invalid_collection_name(
self,
collection_temp_dir,
collection_name,
collection_option,
sample_field_list,
sample_vector_list,
):
with pytest.raises(Exception) as exc_info:
collection_schema = zvec.CollectionSchema(
name=collection_name,
fields=sample_field_list,
vectors=sample_vector_list,
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize("name_prefix", FIELD_NAME_VALID_LIST)
def test_valid_field_vector_name(
self,
collection_temp_dir,
collection_option,
name_prefix,
sample_field_list,
sample_vector_list,
):
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=sample_field_list,
vectors=sample_vector_list,
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
check_collection_info(
coll, collection_schema, collection_option, collection_temp_dir
)
check_collection_basic(coll)
coll.destroy()
@pytest.mark.parametrize("field_name", FIELD_NAME_INVALID_LIST)
def test_invalid_field_name(
self, collection_temp_dir, collection_option, field_name
):
with pytest.raises(Exception) as exc_info:
field_list = [FieldSchema(field_name, DataType.STRING)]
vector_list = [
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
]
collection_schema = zvec.CollectionSchema(
name="collection_name", fields=field_list, vectors=vector_list
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize("vector_name", FIELD_NAME_INVALID_LIST)
def test_invalid_vector_name(
self, collection_temp_dir, collection_option, vector_name
):
with pytest.raises(Exception) as exc_info:
field_list = [
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
)
]
vector_list = [
VectorSchema(vector_name, DataType.VECTOR_FP32, dimension=128)
]
collection_schema = zvec.CollectionSchema(
name="collection_name", fields=field_list, vectors=vector_list
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize(
"field_list_len,vector_list_len,dimension",
FIELD_VECTOR_LIST_DIMENSION_VALID_LIST,
)
def test_valid_field_vector_size_dimension(
self,
collection_temp_dir,
collection_option,
field_list_len,
vector_list_len,
dimension,
):
field_list = []
vector_list = []
for i in range(0, field_list_len):
field_list.append(
FieldSchema("id_" + str(i), DataType.INT64, nullable=True)
)
for i in range(0, vector_list_len):
vector_list.append(
VectorSchema(
"dense_vector_" + str(i),
DataType.VECTOR_FP32,
dimension=dimension,
index_param=HnswIndexParam(),
)
)
collection_schema = zvec.CollectionSchema(
name="test_dense_vector_list", fields=field_list, vectors=vector_list
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
check_collection_info(
coll, collection_schema, collection_option, collection_temp_dir
)
check_collection_basic(coll)
coll.destroy()
@pytest.mark.parametrize(
"field_list_len,vector_list_len,dimension",
FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST,
)
def test_invalid_field_vector_size_dimension(
self,
collection_temp_dir,
collection_option,
vector_list_len,
field_list_len,
dimension,
):
with pytest.raises(Exception) as exc_info:
field_list = []
vector_list = []
for i in range(0, field_list_len):
field_list.append(
FieldSchema(
"id_" + str(i),
DataType.INT64,
nullable=False,
)
)
for i in range(0, vector_list_len):
vector_list.append(
VectorSchema(
"dense_vector_" + str(i),
DataType.VECTOR_FP32,
dimension=dimension,
index_param=HnswIndexParam(),
)
)
collection_schema = zvec.CollectionSchema(
name="test_dense_vector_list", fields=field_list, vectors=vector_list
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
def test_valid_single_vector_field_construction(
self, collection_temp_dir, collection_option
):
field = FieldSchema(
"id",
DataType.INT64,
nullable=True,
index_param=InvertIndexParam(enable_range_optimization=True),
)
vector = VectorSchema(
"dense_vector",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
collection_schema = zvec.CollectionSchema(
name="test_single_dense_vector_non_list",
fields=field,
vectors=vector, # Non-list form
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
check_collection_info(
coll, collection_schema, collection_option, collection_temp_dir
)
check_collection_basic(coll)
coll.destroy()
def test_collection_concurrent_create(
self, collection_temp_dir, basic_schema, collection_option
):
results = []
errors = []
lock = threading.Lock()
# Function to be executed by each thread
def create_collection_thread(thread_id):
try:
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=basic_schema,
option=collection_option,
)
with lock:
results.append((thread_id, coll))
except Exception as e:
with lock:
errors.append((thread_id, str(e)))
threads = []
for i in range(5):
thread = threading.Thread(target=create_collection_thread, args=(i,))
threads.append(thread)
thread.start()
for thread in threads:
thread.join()
assert len(results) == 1, (
f"Expected exactly one successful creation, but got {len(results)}"
)
assert len(errors) == 4, (
f"Expected exactly four failures, but got {len(errors)}"
)
successful_thread_id, successful_collection = results[0]
assert successful_collection is not None, (
"Successful creation should return a valid collection"
)
assert successful_collection.path == collection_temp_dir, (
"Collection path mismatch"
)
def test_create_open_loop(
self, collection_temp_dir, collection_option, full_schema
):
for cycle in range(10):
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=full_schema,
option=collection_option,
)
assert coll is not None, (
f"Failed to create and open collection in cycle {cycle}"
)
assert coll.path == collection_temp_dir, (
f"Collection path mismatch in cycle {cycle}"
)
del coll
reopened_coll = zvec.open(
path=collection_temp_dir, option=collection_option
)
assert reopened_coll is not None, (
f"Failed to reopen collection in cycle {cycle}"
)
assert reopened_coll.path == collection_temp_dir, (
f"Reopened collection path mismatch in cycle {cycle}"
)
check_collection_full(reopened_coll)
reopened_coll.destroy()
@pytest.mark.parametrize(
"data_type, index_param", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS
)
def test_valid_vector_index_params(
self,
data_type,
index_param,
single_vector_schema_with_index_param,
collection_temp_dir,
collection_option,
):
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=single_vector_schema_with_index_param,
option=collection_option,
)
check_collection_info(
coll,
single_vector_schema_with_index_param,
collection_option,
collection_temp_dir,
)
check_collection_basic(coll, True)
@pytest.mark.parametrize(
"data_type, index_param", INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS
)
def test_invalid_vector_index_params(
self,
data_type,
index_param,
single_vector_schema_with_index_param,
collection_temp_dir,
collection_option,
):
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=single_vector_schema_with_index_param,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
def test_open_concurrent_same_path(self, tmp_path_factory, collection_option):
"""Test concurrent opening of the same collection path.
- Multi-threading concurrency: 5 threads simultaneously open the same collection
- Result verification: Verify that only one can open successfully, others must fail
"""
# Create a temporary directory and path for the collection
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "concurrent_open_test_collection"
# First, create a collection that we'll try to open concurrently
field_list = [
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name", DataType.STRING, nullable=False, index_param=InvertIndexParam()
),
]
vector_list = [
VectorSchema(
"dense_vector",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
]
collection_schema = zvec.CollectionSchema(
name="concurrent_open_test_collection",
fields=field_list,
vectors=vector_list,
)
# Create the collection first
coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
# Close the collection so we can test opening it
if hasattr(coll, "close") and coll is not None:
coll.close()
# Shared variables to collect results from threads
results = []
errors = []
# Lock for thread-safe operations
lock = threading.Lock()
# Clean up the created collection reference
del coll
# Function to be executed by each thread
def open_collection_thread(thread_id):
try:
reopened_coll = zvec.open(
path=str(collection_path), option=collection_option
)
with lock:
results.append((thread_id, reopened_coll))
# Clean up the collection if opened successfully
if hasattr(reopened_coll, "close") and reopened_coll is not None:
reopened_coll.close()
except Exception as e:
with lock:
errors.append((thread_id, str(e)))
# Create and start 5 threads
threads = []
for i in range(5):
thread = threading.Thread(target=open_collection_thread, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Verify results:
# 1. Only one open should succeed (exactly one collection in results)
# 2. Others should fail (4 errors in errors)
assert len(results) == 1, (
f"Expected exactly one successful open, but got {len(results)}"
)
assert len(errors) == 4, (
f"Expected exactly four failures, but got {len(errors)}"
)
# Additional verification: check that the successful open has a valid collection
successful_thread_id, successful_collection = results[0]
assert successful_collection is not None, (
"Successful open should return a valid collection"
)
assert successful_collection.path == str(collection_path), (
"Collection path mismatch"
)
@pytest.mark.parametrize("read_only,enable_mmap", valid_collection_options)
def test_valid_option(
self, collection_temp_dir, basic_schema, read_only, enable_mmap
):
option = CollectionOption(read_only=read_only, enable_mmap=enable_mmap)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=basic_schema,
option=option,
)
check_collection_info(coll, basic_schema, option, collection_temp_dir)
check_collection_basic(coll)
coll.destroy()
def test_valid_none_option(self, collection_temp_dir, basic_schema):
zvec.create_and_open(
path=collection_temp_dir,
schema=basic_schema,
option=None,
)
@pytest.mark.parametrize("read_only,enable_mmap", invalid_collection_options)
def test_invalid_option(
self, collection_temp_dir, basic_schema, read_only, enable_mmap
):
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=basic_schema,
option=CollectionOption(read_only=read_only, enable_mmap=enable_mmap),
)
assert CREATE_READ_ONLY_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize(
"field_name1,field_name2,vector_name1,vector_name2",
duplicate_names_test,
)
def test_duplicate_field_names(
self,
collection_temp_dir,
collection_option,
field_name1,
field_name2,
vector_name1,
vector_name2,
):
with pytest.raises(Exception) as exc_info:
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
field_name1,
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
field_name2,
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
vector_name1,
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
vector_name2,
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
],
)
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize("long_name", long_names)
def test_invalid_long_field_names(
self, collection_option, collection_temp_dir, long_name
):
collection_schema = zvec.CollectionSchema(
name=long_name,
fields=[
FieldSchema(
long_name + "_field",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
long_name + "_vector",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
def test_invalid_empty_fields_and_vectors(
self, collection_temp_dir, collection_option
):
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[], # Empty fields
vectors=[], # Empty vectors
)
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=collection_temp_dir,
schema=collection_schema,
option=collection_option,
)
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
@pytest.mark.parametrize("valid_path", valid_path_list)
def test_valid_path(self, basic_schema, collection_option, valid_path):
if os.path.exists(valid_path):
import shutil
shutil.rmtree(valid_path)
coll = zvec.create_and_open(
path=valid_path, schema=basic_schema, option=collection_option
)
check_collection_info(coll, basic_schema, collection_option, valid_path)
coll.destroy()
@pytest.mark.parametrize("invalid_path", invalid_path_list)
def test_invalid_path(self, basic_schema, collection_option, invalid_path):
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=invalid_path, schema=basic_schema, option=collection_option
)
assert INVALID_PATH_ERROR_MSG in str(exc_info.value), str(exc_info.value)
python/tests/detail/test_collection_ddl.py
+1720
View File
File diff suppressed because it is too large Load Diff
python/tests/detail/test_collection_dml.py
+1111
View File
File diff suppressed because it is too large Load Diff
python/tests/detail/test_collection_dql.py
+1434
View File
File diff suppressed because it is too large Load Diff
python/tests/detail/test_collection_exception.py
+328
View File
@@ -0,0 +1,328 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import pytest
import numpy as np
import zvec
from zvec import (
CollectionOption,
InvertIndexParam,
HnswIndexParam,
DataType,
Collection,
Doc,
FieldSchema,
Query,
VectorSchema,
)
class TestCollectionExceptionHandling:
@pytest.fixture(scope="function")
def test_collection(self, tmp_path_factory):
"""Fixture to create a test collection"""
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
FieldSchema("weight", DataType.FLOAT, nullable=True),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open collection"
yield coll
# Clean up
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_create_and_open_missing_path(self, tmp_path_factory):
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
schema=collection_schema, option=collection_option
)
assert exc_info.value is not None, (
"Expected exception for missing path parameter"
)
def test_create_and_open_missing_schema(self, tmp_path_factory):
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
collection_option = CollectionOption(read_only=False, enable_mmap=True)
with pytest.raises(Exception) as exc_info:
coll = zvec.create_and_open(
path=str(collection_path), option=collection_option
)
assert exc_info.value is not None, (
"Expected exception for missing schema parameter"
)
def test_open_missing_path(self):
collection_option = CollectionOption(read_only=False, enable_mmap=True)
with pytest.raises(Exception) as exc_info:
coll = zvec.open(option=collection_option)
assert exc_info.value is not None, (
"Expected exception for missing path parameter"
)
def test_insert_missing_docs(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.insert()
assert exc_info.value is not None, (
"Expected exception for missing docs parameter"
)
def test_update_missing_docs(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.update()
assert exc_info.value is not None, (
"Expected exception for missing docs parameter"
)
def test_upsert_missing_docs(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.upsert()
assert exc_info.value is not None, (
"Expected exception for missing docs parameter"
)
def test_delete_missing_ids(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.delete()
assert exc_info.value is not None, (
"Expected exception for missing ids parameter"
)
def test_fetch_missing_ids(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.fetch()
assert exc_info.value is not None, (
"Expected exception for missing ids parameter"
)
def test_query_missing_query_field_name(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
result = test_collection.query([Query()])
assert exc_info.value is not None, (
"Expected exception for missing Query field_name parameter"
)
def test_add_column_missing_field_schema(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
test_collection.add_column()
assert exc_info.value is not None, (
"Expected exception for missing field_schema parameter"
)
def test_alter_column_missing_old_name(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
test_collection.alter_column(new_name="new_name")
assert exc_info.value is not None, (
"Expected exception for missing old_name parameter"
)
def test_alter_column_missing_new_name(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
test_collection.alter_column(old_name="old_name")
assert exc_info.value is not None, (
"Expected exception for missing new_name parameter"
)
def test_drop_column_missing_field_name(self, test_collection: Collection):
with pytest.raises(Exception) as exc_info:
test_collection.drop_column()
assert exc_info.value is not None, (
"Expected exception for missing field_name parameter"
)
def test_invalid_parameter_types(self, test_collection: Collection):
# This test depends on specific implementation details
# Generally, we would expect TypeErrors or similar exceptions
pass
def test_missing_required_parameters(self, test_collection: Collection):
# This test depends on specific implementation details
# Generally, we would expect TypeErrors or similar exceptions
pass
def test_empty_collection_operations(self, tmp_path_factory):
collection_schema = zvec.CollectionSchema(
name="empty_test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "empty_test_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open collection"
# Test fetch on empty collection
result = coll.fetch(["1"])
assert len(result) >= 0 # May be empty or have special handling
# Test query on empty collection
result = coll.query()
assert len(result) == 0
# Test update on empty collection
doc = Doc(
id="1",
fields={"id": 1, "name": "test"},
vectors={"dense": np.random.random(128).tolist()},
)
result = coll.update(doc)
# Should handle gracefully, possibly with NOT_FOUND status
# Clean up
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_resource_management(self, test_collection: Collection):
doc = Doc(
id="1",
fields={"id": 1, "name": "test", "weight": 80.5},
vectors={
"dense": np.random.random(128).tolist(),
"sparse": {1: 1.0, 2: 2.0},
},
)
# Insert
result = test_collection.insert(doc)
assert result.ok()
# Fetch
result = test_collection.fetch(["1"])
assert len(result) == 1
# Query
result = test_collection.query()
assert len(result) >= 0
# Update
result = test_collection.update(doc)
assert result.ok()
# Delete
result = test_collection.delete("1")
assert result.ok()
def test_exception_resource_cleanup(self, test_collection: Collection):
# This test would need to simulate exception conditions
# which is difficult without specific failure injection points
pass
python/tests/detail/test_collection_open.py
+967
View File
@@ -0,0 +1,967 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import threading
import numpy as np
from fixture_helper import *
COLLECTION_OPTION_TEST_CASES_VALID = [
# (read_only, enable_mmap, description)
(False, True, "Read-write with mmap enabled"),
(False, False, "Read-write with mmap disabled"),
(True, True, "Read-only with mmap enabled"),
(True, False, "Read-only with mmap disabled"),
]
# Test data for invalid paths
INVALID_PATH_LIST = [
"/nonexistent/directory/test_collection",
"invalid:path",
"", # Empty path
]
@pytest.fixture(scope="session")
def collection_schema():
return zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name", DataType.STRING, nullable=False, index_param=InvertIndexParam()
),
FieldSchema(
"weight", DataType.FLOAT, nullable=False, index_param=InvertIndexParam()
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
@pytest.fixture
def single_doc():
id = 0
return Doc(
id=f"{id}",
fields={"id": id, "name": "test"},
vectors={
"dense": [id + 0.1] * 128,
},
)
@pytest.fixture(scope="function")
def test_collection(
tmp_path_factory, collection_schema, collection_option
) -> Generator[Any, Any, Collection]:
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
coll = zvec.create_and_open(
path=str(collection_path), schema=collection_schema, option=collection_option
)
assert coll is not None, "Failed to create and open collection"
assert coll.path == str(collection_path)
assert coll.schema.name == collection_schema.name
assert list(coll.schema.fields) == list(collection_schema.fields)
assert list(coll.schema.vectors) == list(collection_schema.vectors)
assert coll.option.read_only == collection_option.read_only
assert coll.option.enable_mmap == collection_option.enable_mmap
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
class TestCollectionOpen:
def test_open_basic_functionality(
self, tmp_path_factory, collection_schema, collection_option
):
import sys
import time
import os
# Create unique temp directory
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
# Ensure the path exists
collection_path_str = str(collection_path)
print(f"DEBUG: Collection path: {collection_path_str}")
print(f"DEBUG: Temp directory exists: {temp_dir.exists()}")
# Create and open collection first
created_coll = zvec.create_and_open(
path=collection_path_str, schema=collection_schema, option=collection_option
)
assert created_coll is not None, (
f"Failed to create collection, returned None instead of valid Collection object. Path: {collection_path_str}"
)
assert created_coll.path == collection_path_str, (
f"Collection path mismatch. Expected: {collection_path_str}, Actual: {created_coll.path}"
)
assert created_coll.schema.name == "test_collection", (
f"Collection schema name mismatch. Expected: test_collection, Actual: {created_coll.schema.name}"
)
# Insert multiple documents to verify persistence
docs = []
for i in range(3):
doc = Doc(
id=f"{i}",
fields={"id": i, "name": f"test_{i}", "weight": float(i * 10)},
vectors={
"dense": [float(j + i) for j in range(128)],
"sparse": {j: float(j + i) for j in range(5)},
},
)
docs.append(doc)
result = created_coll.insert(docs)
assert len(result) == 3, f"Expected 3 insertion results, but got {len(result)}"
for i, res in enumerate(result):
assert res.ok(), (
f"Insertion result {i} is not OK. Status code: {res.code()}, Message: {res.message()}"
)
# Verify documents were inserted using fetch interface
fetched_docs_after_insert = created_coll.fetch(["0", "1", "2"])
assert len(fetched_docs_after_insert) == 3, (
f"Expected 3 fetched documents after insertion, but got {len(fetched_docs_after_insert)}"
)
assert "0" in fetched_docs_after_insert, (
"Document with ID '0' not found in fetched results after insertion"
)
assert "1" in fetched_docs_after_insert, (
"Document with ID '1' not found in fetched results after insertion"
)
assert "2" in fetched_docs_after_insert, (
"Document with ID '2' not found in fetched results after insertion"
)
# Verify fetched document content after insertion
for i in range(3):
doc = fetched_docs_after_insert[f"{i}"]
assert doc is not None, (
f"Fetched document with ID '{i}' is None after insertion"
)
assert doc.id == f"{i}", (
f"Document ID mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.id}"
)
assert doc.field("id") == i, (
f"Document id field mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.field('id')}"
)
assert doc.field("name") == f"test_{i}", (
f"Document name field mismatch for document '{i}' after insertion. Expected: test_{i}, Actual: {doc.field('name')}"
)
assert doc.field("weight") == float(i * 10), (
f"Document weight field mismatch for document '{i}' after insertion. Expected: {float(i * 10)}, Actual: {doc.field('weight')}"
)
# Verify vector access after insertion
assert doc.vector("dense") is not None, (
f"Document {i} should have dense vector after insertion"
)
assert doc.vector("sparse") is not None, (
f"Document {i} should have sparse vector after insertion"
)
# Verify vector types after insertion
assert isinstance(doc.vector("dense"), list), (
f"Document {i} dense vector should be dict after insertion, got {type(doc.vector('dense'))}"
)
assert isinstance(doc.vector("sparse"), dict), (
f"Document {i} sparse vector should be dict after insertion, got {type(doc.vector('sparse'))}"
)
# Verify documents were inserted using stats
stats = created_coll.stats
assert stats is not None, "Collection stats should not be None"
assert stats.doc_count == 3, (
f"Document count mismatch after insertion. Expected: 3, Actual: {stats.doc_count}"
)
# Store the collection path before cleanup
collection_path = created_coll.path
# Clean up the created collection reference
del created_coll
# Wait and verify the path still exists
print(f"DEBUG: Collection path after destroy: {collection_path}")
print(f"DEBUG: Path exists after destroy: {os.path.exists(collection_path)}")
# Now open the existing collection
try:
print(f"DEBUG: Path exists before open: {os.path.exists(collection_path)}")
# List contents of parent directory for debugging
parent_dir = os.path.dirname(collection_path)
if os.path.exists(parent_dir):
print(f"DEBUG: Parent directory contents: {os.listdir(parent_dir)}")
opened_coll = zvec.open(path=collection_path, option=collection_option)
assert opened_coll is not None, (
f"Failed to open existing collection at path: {collection_path}. Returned None instead of valid Collection object"
)
assert opened_coll.path == collection_path, (
f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}"
)
assert opened_coll.schema.name == "test_collection", (
f"Opened collection schema name mismatch. Expected: test_collection, Actual: {opened_coll.schema.name}"
)
# Check reference count of opened collection
opened_ref_count = sys.getrefcount(opened_coll)
print(f"DEBUG: Reference count of opened collection: {opened_ref_count}")
# Verify data persistence
# Verify data persistence using fetch interface
fetched_docs = opened_coll.fetch(["0", "1", "2"])
assert len(fetched_docs) == 3, (
f"Expected 3 fetched documents after reopening, but got {len(fetched_docs)}"
)
assert "0" in fetched_docs, (
"Document with ID '0' not found in fetched results after reopening"
)
assert "1" in fetched_docs, (
"Document with ID '1' not found in fetched results after reopening"
)
assert "2" in fetched_docs, (
"Document with ID '2' not found in fetched results after reopening"
)
# Verify fetched document content after reopening collection
for i in range(3):
doc = fetched_docs[f"{i}"]
assert doc is not None, (
f"Fetched document with ID '{i}' is None after reopening collection"
)
assert doc.id == f"{i}", (
f"Document ID mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.id}"
)
assert doc.field("id") == i, (
f"Document id field mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.field('id')}"
)
assert doc.field("name") == f"test_{i}", (
f"Document name field mismatch for document '{i}' after reopening. Expected: test_{i}, Actual: {doc.field('name')}"
)
assert doc.field("weight") == float(i * 10), (
f"Document weight field mismatch for document '{i}' after reopening. Expected: {float(i * 10)}, Actual: {doc.field('weight')}"
)
# Verify vector access after reopening
assert doc.vector("dense") is not None, (
f"Document {i} should have dense vector after reopening"
)
assert doc.vector("sparse") is not None, (
f"Document {i} should have sparse vector after reopening"
)
# Verify vector types after reopening
assert isinstance(doc.vector("dense"), list), (
f"Document {i} dense vector should be dict after reopening, got {type(doc.vector('dense'))}"
)
assert isinstance(doc.vector("sparse"), dict), (
f"Document {i} sparse vector should be dict after reopening, got {type(doc.vector('sparse'))}"
)
# Verify score attribute exists
assert hasattr(doc, "score"), (
f"Document {i} should have a score attribute after reopening"
)
assert isinstance(doc.score, (int, float)), (
f"Document {i} score should be numeric after reopening, got {type(doc.score)}"
)
# For fetch operations, score is typically 0.0
assert doc.score == 0.0, (
f"Document {i} score should be 0.0 for fetch operation after reopening, but got {doc.score}"
)
# Test query functionality
query_result = opened_coll.query(include_vector=True)
assert len(query_result) == 3, (
f"Expected 3 query results, but got {len(query_result)}"
)
# Verify query results have proper structure and content with detailed validation
returned_doc_ids = set()
for doc in query_result:
# Verify basic document structure
assert doc.id is not None, f"Query result document should have an ID"
assert doc.id in ["0", "1", "2"], (
f"Query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
)
returned_doc_ids.add(doc.id)
# Verify field access
assert doc.field("id") is not None, (
f"Document {doc.id} should have id field"
)
assert doc.field("name") is not None, (
f"Document {doc.id} should have name field"
)
assert doc.field("weight") is not None, (
f"Document {doc.id} should have weight field"
)
# Verify field values
expected_id = int(doc.id)
assert doc.field("id") == expected_id, (
f"Document {doc.id} id field mismatch. Expected: {expected_id}, Actual: {doc.field('id')}"
)
assert doc.field("name") == f"test_{expected_id}", (
f"Document {doc.id} name field mismatch. Expected: test_{expected_id}, Actual: {doc.field('name')}"
)
assert doc.field("weight") == float(expected_id * 10), (
f"Document {doc.id} weight field mismatch. Expected: {float(expected_id * 10)}, Actual: {doc.field('weight')}"
)
# Verify vector access
assert doc.vector("dense") is not None, (
f"Document {doc.id} should have dense vector"
)
assert doc.vector("sparse") is not None, (
f"Document {doc.id} should have sparse vector"
)
# Verify vector types
assert isinstance(doc.vector("dense"), list), (
f"Document {doc.id} dense vector should be list, got {type(doc.vector('dense'))}"
)
assert isinstance(doc.vector("sparse"), dict), (
f"Document {doc.id} sparse vector should be dict, got {type(doc.vector('sparse'))}"
)
# Verify score attribute exists
assert hasattr(doc, "score"), (
f"Document {doc.id} should have a score attribute"
)
assert isinstance(doc.score, (int, float)), (
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
)
# Verify all expected documents are returned
expected_doc_ids = {"0", "1", "2"}
assert returned_doc_ids == expected_doc_ids, (
f"Query should return all expected documents. Expected: {expected_doc_ids}, Actual: {returned_doc_ids}"
)
# === Enhanced validation based on test_collection_dql_operations.py ===
# Verify vector field names accessibility for all documents
for doc in query_result:
vector_names = doc.vector_names()
expected_vector_names = {"dense", "sparse"}
assert set(vector_names) == expected_vector_names, (
f"Document {doc.id} vector names mismatch. Expected: {expected_vector_names}, Actual: {set(vector_names)}"
)
# Verify all vector fields can be accessed
for vector_name in expected_vector_names:
vector_data = doc.vector(vector_name)
assert vector_data is not None, (
f"Document {doc.id} should have accessible vector '{vector_name}'"
)
if vector_name == "dense":
assert isinstance(vector_data, list), (
f"Document {doc.id} vector '{vector_name}' should be list, got {type(vector_data)}"
)
else:
assert isinstance(vector_data, dict), (
f"Document {doc.id} vector '{vector_name}' should be dict, got {type(vector_data)}"
)
# Test query with filter
filtered_result = opened_coll.query(filter="id >= 1", include_vector=True)
assert len(filtered_result) == 2, (
f"Expected 2 filtered query results (id >= 1), but got {len(filtered_result)}"
)
# Verify filtered query results
filtered_doc_ids = set()
for doc in filtered_result:
assert doc.id is not None, (
f"Filtered query result document should have an ID"
)
assert doc.id in ["1", "2"], (
f"Filtered query result document ID should be one of ['1', '2'], but got {doc.id}"
)
filtered_doc_ids.add(doc.id)
# Verify filter condition is satisfied
doc_id = int(doc.id)
assert doc_id >= 1, (
f"Document {doc.id} should satisfy filter condition id >= 1"
)
# Verify document structure
assert doc.field("id") is not None, (
f"Document {doc.id} should have id field"
)
assert doc.field("name") is not None, (
f"Document {doc.id} should have name field"
)
assert doc.field("weight") is not None, (
f"Document {doc.id} should have weight field"
)
# Verify field values
assert doc.field("id") == doc_id, (
f"Document {doc.id} id field mismatch. Expected: {doc_id}, Actual: {doc.field('id')}"
)
assert doc.field("name") == f"test_{doc_id}", (
f"Document {doc.id} name field mismatch. Expected: test_{doc_id}, Actual: {doc.field('name')}"
)
assert doc.field("weight") == float(doc_id * 10), (
f"Document {doc.id} weight field mismatch. Expected: {float(doc_id * 10)}, Actual: {doc.field('weight')}"
)
# Verify vector access
assert doc.vector("dense") is not None, (
f"Document {doc.id} should have dense vector"
)
assert doc.vector("sparse") is not None, (
f"Document {doc.id} should have sparse vector"
)
# Verify score attribute exists
assert hasattr(doc, "score"), (
f"Document {doc.id} should have a score attribute"
)
assert isinstance(doc.score, (int, float)), (
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
)
# Verify filtered documents
expected_filtered_ids = {"1", "2"}
assert filtered_doc_ids == expected_filtered_ids, (
f"Filtered query should return expected documents. Expected: {expected_filtered_ids}, Actual: {filtered_doc_ids}"
)
# Test vector query functionality for dense vectors
query_vector_dense = [0.1] * 128
vector_query_result = opened_coll.query(
Query(field_name="dense", vector=query_vector_dense)
)
assert len(vector_query_result) > 0, (
f"Expected at least 1 vector query result, but got {len(vector_query_result)}"
)
# Verify vector query results structure
for doc in vector_query_result[:3]: # Check first 3 results
assert doc.id is not None, (
f"Vector query result document should have an ID"
)
assert doc.id in ["0", "1", "2"], (
f"Vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
)
# Verify document structure
assert doc.field("id") is not None, (
f"Document {doc.id} should have id field"
)
assert doc.field("name") is not None, (
f"Document {doc.id} should have name field"
)
assert doc.field("weight") is not None, (
f"Document {doc.id} should have weight field"
)
# Verify vector access
assert doc.vector("dense") is not None, (
f"Document {doc.id} should have dense vector"
)
assert doc.vector("sparse") is not None, (
f"Document {doc.id} should have sparse vector"
)
# Verify score attribute exists and is numeric
assert hasattr(doc, "score"), (
f"Document {doc.id} should have a score attribute"
)
assert isinstance(doc.score, (int, float)), (
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
)
# For dense vector queries, score should typically be non-negative (depending on metric)
# Note: This may vary based on the metric type used
assert doc.score >= 0 or doc.score < 0, (
f"Document {doc.id} score should be a valid number"
)
# Test vector query functionality for sparse vectors
query_vector_sparse = {1: 1.0, 2: 2.0, 3: 3.0}
sparse_vector_query_result = opened_coll.query(
Query(field_name="sparse", vector=query_vector_sparse)
)
assert len(sparse_vector_query_result) > 0, (
f"Expected at least 1 sparse vector query result, but got {len(sparse_vector_query_result)}"
)
# Verify sparse vector query results structure
for doc in sparse_vector_query_result[:3]: # Check first 3 results
assert doc.id is not None, (
f"Sparse vector query result document should have an ID"
)
assert doc.id in ["0", "1", "2"], (
f"Sparse vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
)
# Verify document structure
assert doc.field("id") is not None, (
f"Document {doc.id} should have id field"
)
assert doc.field("name") is not None, (
f"Document {doc.id} should have name field"
)
assert doc.field("weight") is not None, (
f"Document {doc.id} should have weight field"
)
# Verify vector access
assert doc.vector("dense") is not None, (
f"Document {doc.id} should have dense vector"
)
assert doc.vector("sparse") is not None, (
f"Document {doc.id} should have sparse vector"
)
# Verify score attribute exists and is numeric
assert hasattr(doc, "score"), (
f"Document {doc.id} should have a score attribute"
)
assert isinstance(doc.score, (int, float)), (
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
)
# Clean up
if hasattr(opened_coll, "destroy") and opened_coll is not None:
opened_coll.destroy()
print("DEBUG: Opened collection destroyed successfully")
except Exception as e:
logging.error("Exception occurred: [{}]".format(e))
raise e
@pytest.mark.parametrize(
"read_only,enable_mmap,description", COLLECTION_OPTION_TEST_CASES_VALID
)
@pytest.mark.parametrize("createAndopen_enable_mmap", [True, False])
def test_open_with_different_collection_options_valid(
self,
tmp_path_factory,
createAndopen_enable_mmap,
read_only,
enable_mmap,
description,
collection_schema,
):
# Create collection with initial option
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
initial_option = CollectionOption(
read_only=False, enable_mmap=createAndopen_enable_mmap
)
# Create and open collection first
created_coll = zvec.create_and_open(
path=str(collection_path), schema=collection_schema, option=initial_option
)
assert created_coll is not None, "Failed to create collection"
# Clean up the created collection reference
del created_coll
# Now open with different options
collection_option = CollectionOption(
read_only=read_only, enable_mmap=enable_mmap
)
try:
opened_coll = zvec.open(path=str(collection_path), option=collection_option)
assert opened_coll is not None, (
f"Failed to open collection with option: {description}. Returned None instead of valid Collection object. Path: {collection_path}"
)
assert opened_coll.path == str(collection_path), (
f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}"
)
assert opened_coll.schema.name == collection_schema.name, (
f"Opened collection schema name mismatch. Expected: {collection_schema.name}, Actual: {opened_coll.schema.name}"
)
assert opened_coll.option.read_only == read_only, (
f"Opened collection read_only option mismatch. Expected: {read_only}, Actual: {opened_coll.option.read_only}"
)
assert opened_coll.option.enable_mmap == createAndopen_enable_mmap, (
f"Opened collection mmap option mismatch. Expected: {createAndopen_enable_mmap}, Actual: {opened_coll.option.enable_mmap}"
)
# Clean up
if (
hasattr(opened_coll, "destroy")
and opened_coll is not None
and read_only == False
):
opened_coll.destroy()
except Exception as e:
logging.error("Exception occurred: [{}]".format(e))
pytest.fail(f"Failed to open collection with different options: {e}")
def test_open_with_none_option(self, tmp_path_factory, collection_schema):
# Create collection
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
initial_option = CollectionOption(read_only=False, enable_mmap=True)
# Create and open collection first
created_coll = zvec.create_and_open(
path=str(collection_path), schema=collection_schema, option=initial_option
)
assert created_coll is not None, (
f"Failed to create collection. Returned None instead of valid Collection object. Path: {collection_path}"
)
# Clean up the created collection reference
del created_coll
# Now open with None option
with pytest.raises(Exception) as exc_info:
zvec.open(path=str(collection_path), option=None)
assert "incompatible function arguments" in str(exc_info.value), (
f"Expected 'incompatible function arguments' error, but got: {exc_info.value}"
)
def test_reopen_collection(self, tmp_path_factory):
# Prepare schema
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
FieldSchema(
"description",
DataType.STRING,
nullable=True,
index_param=InvertIndexParam(),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
# Create collection
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
# Create and open collection
coll1 = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert coll1 is not None, "Failed to create and open collection"
# Insert some data
doc = Doc(
id="1",
fields={"id": 1, "name": "test", "description": "这是一个中文描述。"},
vectors={"dense": np.random.random(128).tolist()},
)
result = coll1.insert(doc)
assert result.ok()
# Close the first collection (delete reference)
del coll1
# Reopen the collection
coll2 = zvec.open(path=str(collection_path), option=collection_option)
assert coll2 is not None, "Failed to reopen collection"
assert coll2.path == str(collection_path)
assert coll2.schema.name == collection_schema.name
# Verify data is still there
fetched_docs = coll2.fetch(["1"])
assert "1" in fetched_docs
fetched_doc = fetched_docs["1"]
assert fetched_doc.id == "1"
assert fetched_doc.field("name") == "test"
assert fetched_doc.field("description") == "这是一个中文描述。"
# Clean up
if hasattr(coll2, "destroy") and coll2 is not None:
try:
coll2.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_open_concurrent_same_path(self, tmp_path_factory):
# First create a collection
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
# Create collection path
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
# First create the collection
created_coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert created_coll is not None, "Failed to create collection"
# Close the collection so we can test concurrent opening
if hasattr(created_coll, "close") and created_coll is not None:
created_coll.close()
# Shared variables to collect results from threads
results = []
errors = []
# Lock for thread-safe operations
lock = threading.Lock()
# Clean up the created collection reference
del created_coll
# Function to be executed by each thread
def open_collection_thread(thread_id):
try:
coll = zvec.open(path=str(collection_path), option=collection_option)
with lock:
results.append((thread_id, coll))
# Close the collection if opened successfully
if hasattr(coll, "close") and coll is not None:
coll.close()
except Exception as e:
with lock:
errors.append((thread_id, str(e)))
# Create 5 threads to call open concurrently
threads = []
for i in range(5):
thread = threading.Thread(target=open_collection_thread, args=(i,))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Verify concurrency safety: only one should succeed, others should fail
assert len(results) == 1, (
f"Expected exactly one successful open, but got {len(results)}"
)
assert len(errors) == 4, (
f"Expected exactly four failures, but got {len(errors)}"
)
# Additional verification: check that the successful open has a valid collection
successful_thread_id, successful_collection = results[0]
assert successful_collection is not None, (
"Successful open should return a valid collection"
)
assert successful_collection.path == str(collection_path), (
"Collection path mismatch"
)
# Clean up the successfully opened collection
if (
hasattr(successful_collection, "destroy")
and successful_collection is not None
):
try:
successful_collection.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_open_with_corrupted_files(self, tmp_path_factory):
# First create a collection
collection_schema = zvec.CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
)
],
)
collection_option = CollectionOption(read_only=False, enable_mmap=True)
# Create collection path
temp_dir = tmp_path_factory.mktemp("zvec")
collection_path = temp_dir / "test_collection"
# First create the collection
created_coll = zvec.create_and_open(
path=str(collection_path),
schema=collection_schema,
option=collection_option,
)
assert created_coll is not None, "Failed to create collection"
# Close the collection so we can manipulate its files
if hasattr(created_coll, "close") and created_coll is not None:
created_coll.close()
# Test case 1: Delete some files in the collection directory (simulate partial corruption)
import os
import shutil
import random
# Get the collection directory path
collection_dir = str(collection_path)
# List all files in the collection directory
files_in_dir = []
for root, dirs, files in os.walk(collection_dir):
for file in files:
files_in_dir.append(os.path.join(root, file))
# Randomly delete approximately half of the files to simulate partial corruption
if files_in_dir:
# Shuffle the list to randomly select files
random.shuffle(files_in_dir)
files_to_delete = files_in_dir[: len(files_in_dir) // 2]
for file_path in files_to_delete:
try:
os.remove(file_path)
except Exception as e:
pass # Ignore errors during deletion
# Try to open the collection with missing files - should raise an exception
with pytest.raises(Exception):
zvec.open(path=str(collection_path), option=collection_option)
# Test case 2: Delete all files in the collection directory (simulate complete corruption)
# Recreate the collection
recreated_coll = zvec.create_and_open(
path=str(collection_path) + "_all",
schema=collection_schema,
option=collection_option,
)
assert recreated_coll is not None, "Failed to recreate collection"
# Close the collection so we can manipulate its files
if hasattr(recreated_coll, "close") and recreated_coll is not None:
recreated_coll.close()
# Delete all files in the collection directory
try:
shutil.rmtree(collection_dir)
os.makedirs(collection_dir) # Recreate empty directory
except Exception as e:
pass # Ignore errors during deletion
# Try to open the collection with missing files - should raise an exception
with pytest.raises(Exception):
zvec.open(path=str(collection_path), option=collection_option)
python/tests/detail/test_collection_recall.py
+740
View File
@@ -0,0 +1,740 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
from zvec.model import Collection, Doc, Query
from zvec.model.param import (
CollectionOption,
InvertIndexParam,
HnswIndexParam,
FlatIndexParam,
IVFIndexParam,
DiskAnnIndexParam,
HnswQueryParam,
IVFQueryParam,
DiskAnnQueryParam,
)
from zvec.model.schema import FieldSchema, VectorSchema
from zvec.extension import RrfReRanker, WeightedReRanker, QwenReRanker
from distance_helper import *
from zvec import StatusCode
from distance_helper import *
from fixture_helper import *
from doc_helper import *
from params_helper import *
import time
# ==================== helper ====================
def batchdoc_and_check(collection: Collection, multiple_docs, operator="insert"):
if operator == "insert":
result = collection.insert(multiple_docs)
elif operator == "upsert":
result = collection.upsert(multiple_docs)
elif operator == "update":
result = collection.update(multiple_docs)
else:
logging.error("operator value is error!")
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok(), (
f"result={result},Insert operation failed with code {item.code()}"
)
stats = collection.stats
assert stats is not None, "Collection stats should not be None"
"""assert stats.doc_count == len(multiple_docs), (
f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}"
)"""
doc_ids = [doc.id for doc in multiple_docs]
fetched_docs = collection.fetch(doc_ids)
assert len(fetched_docs) == len(multiple_docs), (
f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}"
)
for original_doc in multiple_docs:
assert original_doc.id in fetched_docs, (
f"Expected document ID {original_doc.id} in fetched documents"
)
fetched_doc = fetched_docs[original_doc.id]
assert is_doc_equal(fetched_doc, original_doc, collection.schema)
assert hasattr(fetched_doc, "score"), "Document should have a score attribute"
assert fetched_doc.score == 0.0, (
"Fetch operation should return default score of 0.0"
)
def compute_exact_similarity_scores(
vectors_a,
vectors_b,
metric_type=MetricType.IP,
DataType=DataType.VECTOR_FP32,
QuantizeType=QuantizeType.UNDEFINED,
):
similarities = []
for i, vec_a in enumerate(vectors_a):
for j, vec_b in enumerate(vectors_b):
similarity = distance_recall(vec_a, vec_b, metric_type, DataType)
similarities.append((j, similarity))
# For L2,COSINE metric, smaller distances mean higher similarity, so sort in ascending order
if (
metric_type in [MetricType.L2]
and DataType
in [DataType.VECTOR_FP32, DataType.VECTOR_FP16, DataType.VECTOR_INT8]
) or (
metric_type in [MetricType.COSINE]
and DataType in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]
):
similarities.sort(key=lambda x: x[1], reverse=False) # Ascending order for L2
else:
similarities.sort(
key=lambda x: x[1], reverse=True
) # Descending order for others
# Special handling for COSINE in FP16 to address precision issues
if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16:
# Clamp values to valid cosine distance range [0, 2] and handle floating point errors
similarities = [(idx, max(0.0, min(2.0, score))) for idx, score in similarities]
return similarities
def get_ground_truth_for_vector_query(
collection,
query_vector,
field_name,
all_docs,
query_idx,
metric_type,
k,
use_exact_computation=False,
):
if use_exact_computation:
all_vectors = [doc.vectors[field_name] for doc in all_docs]
for d, f in DEFAULT_VECTOR_FIELD_NAME.items():
if field_name == f:
DataType = d
break
similarities = compute_exact_similarity_scores(
[query_vector],
all_vectors,
metric_type,
DataType=DataType,
QuantizeType=QuantizeType,
)
if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16:
# Filter out tiny non-zero values that may be caused by precision errors
similarities = [
(idx, max(0.0, min(2.0, score))) for idx, score in similarities
]
ground_truth_ids_scores = similarities[:k]
print("Get the most similar k document IDs k:,ground_truth_ids_scores")
print(k, ground_truth_ids_scores)
return ground_truth_ids_scores
else:
full_result = collection.query(
Query(field_name=field_name, vector=query_vector),
topk=min(len(all_docs), 1024),
include_vector=True,
)
ground_truth_ids_scores = [
(result.id, result.score) for result in full_result[:k]
]
if not ground_truth_ids_scores:
ground_truth_ids_scores = [(all_docs[query_idx].id, 0)]
return ground_truth_ids_scores
def get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type, k):
ground_truth_map = {}
for field_name, query_vectors in query_vectors_map.items():
ground_truth_map[field_name] = {}
# Support per-field metric type: metric_type can be a dict mapping
# field_name -> MetricType, or a single MetricType applied to all fields.
if isinstance(metric_type, dict):
field_metric = metric_type.get(field_name, MetricType.IP)
else:
field_metric = metric_type
for i, query_vector in enumerate(query_vectors):
# Get the ground truth for this query
relevant_doc_ids_scores = get_ground_truth_for_vector_query(
collection,
query_vector,
field_name,
test_docs,
i,
field_metric,
k,
True,
)
ground_truth_map[field_name][i] = relevant_doc_ids_scores
print("ground_truth_map:\n")
print(ground_truth_map)
return ground_truth_map
def calculate_recall_at_k(
collection: Collection,
test_docs,
query_vectors_map,
schema,
k=1,
expected_doc_ids_scores_map=None,
tolerance=0.01,
):
recall_stats = {}
for field_name, query_vectors in query_vectors_map.items():
recall_stats[field_name] = {
"relevant_retrieved_count": 0,
"total_relevant_count": 0,
"retrieved_count": 0,
"recall_at_k": 0.0,
}
for i, query_vector in enumerate(query_vectors):
print("Starting %dth query" % i)
query_result_list = collection.query(
Query(field_name=field_name, vector=query_vector),
topk=1024,
include_vector=True,
)
retrieved_count = len(query_result_list)
query_result_ids_scores = []
for word in query_result_list:
query_result_ids_scores.append((word.id, word.score))
recall_stats[field_name]["retrieved_count"] += retrieved_count
print("expected_doc_ids_scores_map:\n")
print(expected_doc_ids_scores_map)
if i in (expected_doc_ids_scores_map[field_name]):
expected_relevant_ids_scores = expected_doc_ids_scores_map[field_name][
i
]
print(
"field_name,i,expected_relevant_ids_scores, query_result_ids_scores:\n"
)
print(
field_name,
i,
"\n",
expected_relevant_ids_scores,
"\n",
len(query_result_ids_scores),
query_result_ids_scores,
)
# Update total relevant documents count
recall_stats[field_name]["total_relevant_count"] += len(
expected_relevant_ids_scores
)
relevant_found_count = 0
for ids_scores_except in expected_relevant_ids_scores:
for ids_scores_result in query_result_ids_scores[:k]:
if int(ids_scores_result[0]) == int(ids_scores_except[0]):
relevant_found_count += 1
break
elif (
int(ids_scores_result[0]) != int(ids_scores_except[0])
and abs(ids_scores_result[1] - ids_scores_except[1])
<= tolerance
):
print("IDs are not equal, but the error is small, tolerance")
print(
ids_scores_result[0],
ids_scores_except[0],
ids_scores_result[1],
ids_scores_except[1],
tolerance,
)
relevant_found_count += 1
break
else:
continue
recall_stats[field_name]["relevant_retrieved_count"] += relevant_found_count
# Calculate Recall@K
if recall_stats[field_name]["total_relevant_count"] > 0:
recall_stats[field_name]["recall_at_k"] = (
recall_stats[field_name]["relevant_retrieved_count"]
/ recall_stats[field_name]["total_relevant_count"]
)
return recall_stats
class TestRecall:
@pytest.mark.parametrize(
"full_schema_new",
[
(True, True, HnswIndexParam()),
(False, True, IVFIndexParam()),
(False, True, DiskAnnIndexParam()),
(False, True, FlatIndexParam()), # ——ok
(
True,
True,
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
),
),
(
True,
True,
HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
),
),
(
True,
True,
HnswIndexParam(
metric_type=MetricType.L2,
m=32,
ef_construction=200,
),
),
(
False,
True,
FlatIndexParam(
metric_type=MetricType.IP,
),
),
(
True,
True,
FlatIndexParam(
metric_type=MetricType.COSINE,
),
),
(
True,
True,
FlatIndexParam(
metric_type=MetricType.L2,
),
),
(
True,
True,
IVFIndexParam(
metric_type=MetricType.IP,
n_list=100,
n_iters=10,
use_soar=False,
),
),
(
True,
True,
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
),
(
True,
True,
IVFIndexParam(
metric_type=MetricType.COSINE,
n_list=150,
n_iters=15,
use_soar=False,
),
),
(
True,
True,
DiskAnnIndexParam(
metric_type=MetricType.IP,
max_degree=32,
),
),
(
True,
True,
DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
),
],
indirect=True,
)
@pytest.mark.parametrize("doc_num", [500])
@pytest.mark.parametrize("query_num", [10])
@pytest.mark.parametrize("top_k", [1])
def test_recall_with_single_vector_valid_500(
self,
full_collection_new: Collection,
doc_num,
query_num,
top_k,
full_schema_new,
request,
):
full_schema_params = request.getfixturevalue("full_schema_new")
# Build per-field metric type map so ground truth uses each field's
# actual index metric (fields may fall back to HnswIndexParam/IP).
field_metric_map = {}
for vector_para in full_schema_params.vectors:
if vector_para.index_param is not None:
field_metric_map[vector_para.name] = vector_para.index_param.metric_type
else:
field_metric_map[vector_para.name] = MetricType.IP
metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
multiple_docs = [
generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
]
print("len(multiple_docs):\n")
print(len(multiple_docs))
# print(multiple_docs)
for i in range(10):
if i != 0:
pass
# print(multiple_docs[i * 1000:1000 * (i + 1)])
batchdoc_and_check(
full_collection_new,
multiple_docs[i * 1000 : 1000 * (i + 1)],
operator="insert",
)
stats = full_collection_new.stats
assert stats.doc_count == len(multiple_docs)
doc_ids = ["0", "1"]
fetched_docs = full_collection_new.fetch(doc_ids)
print("fetched_docs,multiple_docs")
print(
fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"],
fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"],
fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"],
fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"],
"\n",
multiple_docs[0].vectors["sparse_vector_fp32_field"],
multiple_docs[0].vectors["sparse_vector_fp32_field"],
multiple_docs[1].vectors["sparse_vector_fp32_field"],
multiple_docs[1].vectors["sparse_vector_fp16_field"],
)
full_collection_new.optimize(option=OptimizeOption())
time.sleep(2)
query_vectors_map = {}
for field_name in DEFAULT_VECTOR_FIELD_NAME.values():
query_vectors_map[field_name] = [
multiple_docs[i].vectors[field_name] for i in range(query_num)
]
# Get ground truth mapping (pass per-field metric map)
ground_truth_map = get_ground_truth_map(
full_collection_new,
multiple_docs,
query_vectors_map,
field_metric_map,
top_k,
)
# Validate ground truth mapping structure
for field_name in DEFAULT_VECTOR_FIELD_NAME.values():
assert field_name in ground_truth_map
field_gt = ground_truth_map[field_name]
assert len(field_gt) == query_num
for query_idx in range(query_num):
assert query_idx in field_gt
relevant_ids = field_gt[query_idx]
assert isinstance(relevant_ids, list)
assert len(relevant_ids) <= top_k
# Print ground truth statistics
print(f"Ground Truth for Top-{top_k} Retrieval:")
for field_name, field_gt in ground_truth_map.items():
print(f" {field_name}:")
for query_idx, relevant_ids in field_gt.items():
print(
f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}"
)
# Calculate Recall@K using ground truth
recall_at_k_stats = calculate_recall_at_k(
full_collection_new,
multiple_docs,
query_vectors_map,
full_schema_new,
k=top_k,
expected_doc_ids_scores_map=ground_truth_map,
tolerance=0.01,
)
print("ground_truth_map:\n")
print(ground_truth_map)
print("(recall_at_k_stats:\n")
print(recall_at_k_stats)
print("field_metric_map:")
print(field_metric_map)
# Print Recall@K statistics
print(f"Recall@{top_k} using Ground Truth:")
for field_name, stats in recall_at_k_stats.items():
print(f" {field_name}:")
print(
f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}"
)
print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}")
for k, v in recall_at_k_stats.items():
assert v["recall_at_k"] == 1.0
@pytest.mark.parametrize(
"full_schema_new",
[
(True, True, HnswIndexParam()),
(False, True, IVFIndexParam()),
(False, True, FlatIndexParam()), # ——ok
(
True,
True,
HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
),
),
(
True,
True,
HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
),
),
# (True, True, HnswIndexParam(metric_type=MetricType.L2, m=32, ef_construction=200, )),
(
False,
True,
FlatIndexParam(
metric_type=MetricType.IP,
),
),
(
True,
True,
FlatIndexParam(
metric_type=MetricType.COSINE,
),
),
# (True, True, FlatIndexParam(metric_type=MetricType.L2, )),
(
True,
True,
IVFIndexParam(
metric_type=MetricType.IP,
n_list=100,
n_iters=10,
use_soar=False,
),
),
(
True,
True,
IVFIndexParam(
metric_type=MetricType.L2,
n_list=200,
n_iters=20,
use_soar=True,
),
),
(
True,
True,
DiskAnnIndexParam(metric_type=MetricType.IP, max_degree=32),
),
(
True,
True,
DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
),
(
True,
True,
DiskAnnIndexParam(metric_type=MetricType.COSINE, max_degree=32),
),
],
indirect=True,
)
@pytest.mark.parametrize("doc_num", [2000])
@pytest.mark.parametrize("query_num", [2])
@pytest.mark.parametrize("top_k", [1])
@pytest.mark.skip(reason="known bug")
def test_recall_with_single_vector_valid_2000(
self,
full_collection_new: Collection,
doc_num,
query_num,
top_k,
full_schema_new,
request,
):
full_schema_params = request.getfixturevalue("full_schema_new")
# Build per-field metric type map so ground truth uses each field's
# actual index metric (fields may fall back to HnswIndexParam/IP).
field_metric_map = {}
for vector_para in full_schema_params.vectors:
if vector_para.index_param is not None:
field_metric_map[vector_para.name] = vector_para.index_param.metric_type
else:
field_metric_map[vector_para.name] = MetricType.IP
metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
multiple_docs = [
generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
]
print("len(multiple_docs):\n")
print(len(multiple_docs))
# print(multiple_docs)
for i in range(10):
if i != 0:
pass
# print(multiple_docs[i * 1000:1000 * (i + 1)])
batchdoc_and_check(
full_collection_new,
multiple_docs[i * 1000 : 1000 * (i + 1)],
operator="insert",
)
stats = full_collection_new.stats
assert stats.doc_count == len(multiple_docs)
doc_ids = ["0", "1"]
fetched_docs = full_collection_new.fetch(doc_ids)
print("fetched_docs,multiple_docs")
print(
fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"],
fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"],
fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"],
fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"],
"\n",
multiple_docs[0].vectors["sparse_vector_fp32_field"],
multiple_docs[0].vectors["sparse_vector_fp32_field"],
multiple_docs[1].vectors["sparse_vector_fp32_field"],
multiple_docs[1].vectors["sparse_vector_fp16_field"],
)
full_collection_new.optimize(option=OptimizeOption())
time.sleep(2)
query_vectors_map = {}
for field_name in DEFAULT_VECTOR_FIELD_NAME.values():
query_vectors_map[field_name] = [
multiple_docs[i].vectors[field_name] for i in range(query_num)
]
# Get ground truth mapping (pass per-field metric map)
ground_truth_map = get_ground_truth_map(
full_collection_new,
multiple_docs,
query_vectors_map,
field_metric_map,
top_k,
)
# Validate ground truth mapping structure
for field_name in DEFAULT_VECTOR_FIELD_NAME.values():
assert field_name in ground_truth_map
field_gt = ground_truth_map[field_name]
assert len(field_gt) == query_num
for query_idx in range(query_num):
assert query_idx in field_gt
relevant_ids = field_gt[query_idx]
assert isinstance(relevant_ids, list)
assert len(relevant_ids) <= top_k
# Print ground truth statistics
print(f"Ground Truth for Top-{top_k} Retrieval:")
for field_name, field_gt in ground_truth_map.items():
print(f" {field_name}:")
for query_idx, relevant_ids in field_gt.items():
print(
f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}"
)
# Calculate Recall@K using ground truth
recall_at_k_stats = calculate_recall_at_k(
full_collection_new,
multiple_docs,
query_vectors_map,
full_schema_new,
k=top_k,
expected_doc_ids_scores_map=ground_truth_map,
tolerance=0.01,
)
print("ground_truth_map:\n")
print(ground_truth_map)
print("(recall_at_k_stats:\n")
print(recall_at_k_stats)
print("field_metric_map:")
print(field_metric_map)
# Print Recall@K statistics
print(f"Recall@{top_k} using Ground Truth:")
for field_name, stats in recall_at_k_stats.items():
print(f" {field_name}:")
print(
f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}"
)
print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}")
for k, v in recall_at_k_stats.items():
assert v["recall_at_k"] == 1.0
python/tests/detail/test_db_config.py
+307
View File
@@ -0,0 +1,307 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import pytest
import tempfile
import os
import sys
import subprocess
import zvec
import zvec
from zvec import LogType, LogLevel
# Error messages
INITIALIZATION_ERROR_MSG = "initialization failed"
RUNTIME_ERROR_MSG = "RuntimeError"
VALUE_ERROR_MSG = "ValueError"
TYPE_ERROR_MSG = "TypeError"
# ==================== helper ====================
def run_in_subprocess(func):
def wrapper(*args, **kwargs):
if os.getenv("RUNNING_IN_SUBPROCESS"):
return func(*args, **kwargs)
env = os.environ.copy()
env["RUNNING_IN_SUBPROCESS"] = "1"
env["PYTEST_CURRENT_TEST"] = func.__name__
import inspect
filepath = inspect.getfile(func)
qualname = func.__qualname__.replace(".", "::")
test_id = f"{filepath}::{qualname}"
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
env["PYTHONPATH"] = project_root + ":" + env.get("PYTHONPATH", "")
cmd = [sys.executable, "-m", "pytest", "-v", "-s", test_id]
result = subprocess.run(cmd, env=env, capture_output=True, text=True)
if result.returncode != 0:
pytest.fail(
f"Subprocess test {func.__name__} failed with code {result.returncode}\n"
f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
)
return wrapper
# ==================== Fixtures ====================
@pytest.fixture(scope="function")
def temp_log_dir(tmp_path_factory):
return tmp_path_factory.mktemp("logs")
# ==================== Tests ====================
class TestDbConfigInitialization:
@run_in_subprocess
def test_init_default(self):
# default config
# log_type: Optional[LogType] = LogType.CONSOLE,
# log_level: Optional[LogLevel] = LogLevel.WARN,
# log_dir: Optional[str] = "./logs",
# log_basename: Optional[str] = "zvec.log",
# log_file_size: Optional[int] = 2048,
# log_overdue_days: Optional[int] = 7,
zvec.init()
@run_in_subprocess
def test_init_file_logger(self):
from pathlib import Path
import shutil
zvec.init(
log_level=LogLevel.DEBUG,
log_type=LogType.FILE,
)
# assert logdir exist
log_dir = Path("./logs")
assert log_dir.exists()
# validate write log
col = zvec.create_and_open(
"/tmp/test/1",
zvec.CollectionSchema(
name="test",
vectors=zvec.VectorSchema(
dimension=4,
data_type=zvec.DataType.VECTOR_FP32,
name="image",
),
),
)
col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})])
assert any(log_dir.glob("zvec.log.*"))
# clear
col.destroy()
shutil.rmtree(log_dir, ignore_errors=True)
@run_in_subprocess
def test_init_with_mixed_config(self):
zvec.init(
memory_limit_mb=128,
log_type=LogType.FILE,
query_threads=1,
log_level=LogLevel.WARN,
)
@run_in_subprocess
def test_repeated_initialization(self):
# Calling init() repeatedly is allowed:
# it succeeds but becomes a no-op after the first successful init()
zvec.init()
class TestDbConfigMemoryLimitValidation:
@run_in_subprocess
def test_memory_limit_min_valid(self):
# MIN_MEMORY_LIMIT_BYTES is 100M
with pytest.raises(RuntimeError):
zvec.init(memory_limit_mb=99)
@run_in_subprocess
def test_memory_limit_invalid_value(self):
# memory_limit_mb must >= 0 and must be int and if None, set default value
with pytest.raises(ValueError):
zvec.init(memory_limit_mb=0)
with pytest.raises(ValueError):
zvec.init(memory_limit_mb=-1)
with pytest.raises(TypeError):
zvec.init(memory_limit_mb="512")
with pytest.raises(TypeError):
zvec.init(memory_limit_mb=512.5)
class TestDbConfigThreadValidation:
@run_in_subprocess
def test_query_threads(self):
zvec.init(query_threads=1)
@run_in_subprocess
def test_query_threads_invalid(self):
# query_threads must >= 0 and must be int and if None, set default value
with pytest.raises(ValueError):
zvec.init(query_threads=0)
with pytest.raises(ValueError):
zvec.init(query_threads=-1)
with pytest.raises(TypeError):
zvec.init(query_threads="value")
with pytest.raises(TypeError):
zvec.init(query_threads=512.5)
with pytest.raises(TypeError):
zvec.init(query_threads="512")
@run_in_subprocess
def test_optimize_threads(self):
zvec.init(optimize_threads=1)
@run_in_subprocess
def test_optimize_threads_invalid(self):
# optimize_threads must >= 0 and must be int and if None, set default value
with pytest.raises(ValueError):
zvec.init(optimize_threads=0)
with pytest.raises(ValueError):
zvec.init(optimize_threads=-1)
with pytest.raises(TypeError):
zvec.init(optimize_threads="value")
with pytest.raises(TypeError):
zvec.init(optimize_threads=512.5)
with pytest.raises(TypeError):
zvec.init(optimize_threads="512")
class TestDbConfigRatioValidation:
@run_in_subprocess
def test_init_invert_to_forward_scan_ratio(self):
# must be in [0,1]
zvec.init(invert_to_forward_scan_ratio=0.8)
@run_in_subprocess
def test_init_invert_to_forward_scan_ratio_invalid(self):
with pytest.raises(ValueError):
zvec.init(invert_to_forward_scan_ratio=1.1)
with pytest.raises(ValueError):
zvec.init(invert_to_forward_scan_ratio=-0.1)
with pytest.raises(TypeError):
zvec.init(invert_to_forward_scan_ratio="0.8")
@run_in_subprocess
def test_init_brute_force_by_keys_ratio(self):
zvec.init(brute_force_by_keys_ratio=0.8)
@run_in_subprocess
def test_init_brute_force_by_keys_ratio_invalid(self):
with pytest.raises(ValueError):
zvec.init(brute_force_by_keys_ratio=1.1)
with pytest.raises(ValueError):
zvec.init(brute_force_by_keys_ratio=-0.1)
with pytest.raises(TypeError):
zvec.init(brute_force_by_keys_ratio="0.8")
class TestDbConfigLogValidation:
@run_in_subprocess
def test_log_type_valid(self):
zvec.init(log_type=LogType.CONSOLE)
@run_in_subprocess
def test_log_type_invalid(self):
with pytest.raises(TypeError):
zvec.init(log_type="FILE")
with pytest.raises(TypeError):
zvec.init(log_type="")
with pytest.raises(TypeError):
zvec.init(log_type="invalid")
with pytest.raises(TypeError):
zvec.init(log_type=123)
@run_in_subprocess
def test_log_level_valid(self):
zvec.init(log_level=LogLevel.ERROR)
@run_in_subprocess
def test_log_level_invalid(self):
with pytest.raises(TypeError):
zvec.init(log_level="WARN")
with pytest.raises(TypeError):
zvec.init(log_level="")
with pytest.raises(TypeError):
zvec.init(log_level="invalid")
with pytest.raises(TypeError):
zvec.init(log_level=123)
@run_in_subprocess
def test_init_file_logger(self):
from pathlib import Path
import shutil
temp_dir = tempfile.mkdtemp(prefix="log_test_")
abs_temp_dir = os.path.abspath(temp_dir)
zvec.init(
log_level=LogLevel.DEBUG,
log_type=LogType.FILE,
log_dir=abs_temp_dir,
log_basename="test",
)
# assert logdir exist
log_dir = Path(abs_temp_dir)
assert log_dir.exists()
# validate write log
col = zvec.create_and_open(
"/tmp/test/1",
zvec.CollectionSchema(
name="test",
vectors=zvec.VectorSchema(
dimension=4,
data_type=zvec.DataType.VECTOR_FP32,
name="image",
),
),
)
col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})])
assert any(log_dir.glob("test.*"))
# clear
col.destroy()
shutil.rmtree(log_dir, ignore_errors=True)
@run_in_subprocess
def test_log_file_size_invalid(self):
with pytest.raises(TypeError):
zvec.init(log_type=LogType.FILE, log_file_size="df")
with pytest.raises(ValueError):
zvec.init(log_type=LogType.FILE, log_file_size=0)
with pytest.raises(ValueError):
zvec.init(log_type=LogType.FILE, log_file_size=-1)
@run_in_subprocess
def test_log_overdue_days_invalid(self):
with pytest.raises(TypeError):
zvec.init(log_type=LogType.FILE, log_overdue_days="df")
with pytest.raises(ValueError):
zvec.init(log_type=LogType.FILE, log_overdue_days=0)
with pytest.raises(ValueError):
zvec.init(log_type=LogType.FILE, log_overdue_days=-1)
python/tests/test_collection.py
+1218
View File
File diff suppressed because it is too large Load Diff
python/tests/test_collection_diskann.py
+600
View File
@@ -0,0 +1,600 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""End-to-end collection tests for the DiskAnn index.
Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin.
Two platform-level prerequisites are enforced at module import time:
1. DiskAnn is currently built only for Linux x86_64 — other platforms are
skipped wholesale.
2. The DiskAnn backend lives in a *runtime-loaded* plugin
(``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL |
RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory``
singleton is unified with the one inside ``_zvec.so``. After ``import
zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once.
If either prerequisite fails the whole module is skipped so the rest of the
test-suite is not affected.
"""
from __future__ import annotations
import math
import os
import platform
import sys
import pytest
# --------------------------------------------------------------------------- #
# Platform gating (must happen BEFORE we touch zvec).
# --------------------------------------------------------------------------- #
pytestmark = pytest.mark.skipif(
not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
reason="DiskAnn plugin is only supported on Linux x86_64",
)
# Promote all symbols in subsequently-loaded DSOs to the global namespace and
# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see
# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa.
# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement.
if sys.platform == "linux":
sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW)
import zvec # noqa: E402
from zvec import ( # noqa: E402
Collection,
CollectionOption,
DataType,
DiskAnnIndexParam,
DiskAnnQueryParam,
Doc,
FieldSchema,
MetricType,
Query,
VectorSchema,
)
from zvec.typing import QuantizeType # noqa: E402
@pytest.fixture(scope="session")
def diskann_collection_schema():
"""Create a collection schema with a DiskAnn index."""
return zvec.CollectionSchema(
name="test_diskann_collection",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
FieldSchema("name", DataType.STRING, nullable=False),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=128,
index_param=DiskAnnIndexParam(
metric_type=MetricType.L2,
max_degree=64,
list_size=100,
pq_chunk_num=0,
quantize_type=QuantizeType.UNDEFINED,
),
),
],
)
@pytest.fixture(scope="session")
def collection_option():
"""Create collection options."""
return CollectionOption(read_only=False, enable_mmap=True)
@pytest.fixture
def single_doc():
"""Create a single document for testing."""
return Doc(
id="0",
fields={"id": 0, "name": "test_doc_0"},
vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]},
)
@pytest.fixture
def multiple_docs():
"""Create multiple documents for testing."""
return [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_doc_{i}"},
vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]},
)
for i in range(1, 101)
]
@pytest.fixture(scope="function")
def diskann_collection(
tmp_path_factory, diskann_collection_schema, collection_option
) -> Collection:
"""
Function-scoped fixture: creates and opens a collection with DiskAnn index.
"""
temp_dir = tmp_path_factory.mktemp("zvec_diskann")
collection_path = temp_dir / "test_diskann_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=diskann_collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open DiskAnn collection"
assert coll.path == str(collection_path)
assert coll.schema.name == diskann_collection_schema.name
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
@pytest.fixture
def collection_with_single_doc(
diskann_collection: Collection, single_doc: Doc
) -> Collection:
"""Setup: insert single doc into collection."""
assert diskann_collection.stats.doc_count == 0
result = diskann_collection.insert(single_doc)
assert bool(result)
assert result.ok()
assert diskann_collection.stats.doc_count == 1
yield diskann_collection
# Teardown: delete single doc
diskann_collection.delete(single_doc.id)
assert diskann_collection.stats.doc_count == 0
@pytest.fixture
def collection_with_multiple_docs(
diskann_collection: Collection, multiple_docs: list[Doc]
) -> Collection:
"""Setup: insert multiple docs into collection."""
assert diskann_collection.stats.doc_count == 0
result = diskann_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert diskann_collection.stats.doc_count == len(multiple_docs)
yield diskann_collection
# Teardown: delete multiple docs
diskann_collection.delete([doc.id for doc in multiple_docs])
# ==================== Tests ====================
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionCreation:
"""Test DiskAnn collection creation and schema validation."""
def test_collection_creation(
self, diskann_collection: Collection, diskann_collection_schema
):
"""Test that collection is created with correct schema."""
assert diskann_collection is not None
assert diskann_collection.schema.name == diskann_collection_schema.name
assert len(diskann_collection.schema.fields) == len(
diskann_collection_schema.fields
)
assert len(diskann_collection.schema.vectors) == len(
diskann_collection_schema.vectors
)
def test_vector_schema_validation(self, diskann_collection: Collection):
"""Test that vector schema has correct DiskAnn configuration."""
vector_schema = diskann_collection.schema.vector("embedding")
assert vector_schema is not None
assert vector_schema.name == "embedding"
assert vector_schema.data_type == DataType.VECTOR_FP32
assert vector_schema.dimension == 128
index_param = vector_schema.index_param
assert index_param is not None
assert index_param.metric_type == MetricType.L2
assert index_param.max_degree == 64
assert index_param.list_size == 100
assert index_param.pq_chunk_num == 0
def test_collection_stats(self, diskann_collection: Collection):
"""Test initial collection statistics."""
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == 0
assert len(stats.index_completeness) == 1
assert stats.index_completeness["embedding"] == 1
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionInsert:
"""Test document insertion into DiskAnn collection."""
def test_insert_single_doc(self, diskann_collection: Collection, single_doc: Doc):
"""Test inserting a single document."""
result = diskann_collection.insert(single_doc)
assert bool(result)
assert result.ok()
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == 1
def test_insert_multiple_docs(
self, diskann_collection: Collection, multiple_docs: list[Doc]
):
"""Test inserting multiple documents."""
result = diskann_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == len(multiple_docs)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionFetch:
"""Test document fetching from DiskAnn collection."""
def test_fetch_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test fetching a single document by ID."""
result = collection_with_single_doc.fetch(ids=[single_doc.id])
assert bool(result)
assert single_doc.id in result.keys()
doc = result[single_doc.id]
assert doc is not None
assert doc.id == single_doc.id
assert doc.field("id") == single_doc.field("id")
assert doc.field("name") == single_doc.field("name")
def test_fetch_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test fetching multiple documents by IDs."""
ids = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.fetch(ids=ids)
assert bool(result)
assert len(result) == len(ids)
for doc_id in ids:
assert doc_id in result
doc = result[doc_id]
assert doc is not None
assert doc.id == doc_id
def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):
"""Test fetching a non-existent document."""
result = collection_with_single_doc.fetch(ids=["nonexistent_id"])
assert len(result) == 0
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionQuery:
"""Test vector search queries on DiskAnn collection."""
def test_query_by_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by vector with DiskAnn index."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
# First result should be the query document itself (or very close)
first_doc = result[0]
assert first_doc is not None
assert first_doc.id is not None
def test_query_by_id(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by document ID with DiskAnn index."""
query = Query(
field_name="embedding",
id=multiple_docs[0].id,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
def test_query_with_different_list_size(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different list_size parameter values."""
query_vector = multiple_docs[0].vector("embedding")
# Test with list_size=50
query_small = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=50),
)
result_small = collection_with_multiple_docs.query(queries=query_small, topk=10)
assert len(result_small) > 0
# Test with list_size=200
query_large = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=200),
)
result_large = collection_with_multiple_docs.query(queries=query_large, topk=10)
assert len(result_large) > 0
def test_query_with_topk(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different topk values."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
# Test topk=5
result_5 = collection_with_multiple_docs.query(queries=query, topk=5)
assert len(result_5) <= 5
# Test topk=20
result_20 = collection_with_multiple_docs.query(queries=query, topk=20)
assert len(result_20) <= 20
def test_query_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with filter conditions."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
# Query with id filter
result = collection_with_multiple_docs.query(
queries=query, topk=10, filter="id < 50"
)
assert len(result) > 0
for doc in result:
assert doc.field("id") < 50
def test_query_with_output_fields(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with specific output fields."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, output_fields=["id", "name"]
)
assert len(result) > 0
first_doc = result[0]
assert "id" in first_doc.field_names()
assert "name" in first_doc.field_names()
def test_query_with_include_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with vector data included in results."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, include_vector=True
)
assert len(result) > 0
first_doc = result[0]
assert first_doc.vector("embedding") is not None
assert len(first_doc.vector("embedding")) == 128
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionUpdate:
"""Test document update in DiskAnn collection."""
def test_update_doc_fields(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document fields."""
updated_doc = Doc(
id=single_doc.id,
fields={"id": single_doc.field("id"), "name": "updated_name"},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(ids=[single_doc.id])
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.field("name") == "updated_name"
def test_update_doc_vector(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document vector."""
new_vector = [0.5 + i * 0.01 for i in range(128)]
updated_doc = Doc(
id=single_doc.id,
vectors={"embedding": new_vector},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(
ids=[single_doc.id],
)
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.vector("embedding") is not None
embedding = doc.vector("embedding")
assert len(embedding) == 128
# Verify vector values are approximately equal (float comparison)
for i in range(128):
assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionDelete:
"""Test document deletion from DiskAnn collection."""
def test_delete_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test deleting a single document."""
result = collection_with_single_doc.delete(single_doc.id)
assert bool(result)
assert result.ok()
stats = collection_with_single_doc.stats
assert stats.doc_count == 0
def test_delete_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test deleting multiple documents."""
ids_to_delete = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.delete(ids_to_delete)
assert len(result) == len(ids_to_delete)
for item in result:
assert item.ok()
stats = collection_with_multiple_docs.stats
assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionOptimizeAndReopen:
"""Test collection optimize and reopen functionality."""
def test_optimize_close_reopen_and_query(
self,
tmp_path_factory,
diskann_collection_schema,
collection_option,
multiple_docs: list[Doc],
):
"""Test inserting 100 docs, optimize, close, reopen and query."""
# Create collection and insert 100 documents
temp_dir = tmp_path_factory.mktemp("zvec_diskann_optimize")
collection_path = temp_dir / "test_optimize_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=diskann_collection_schema,
option=collection_option,
)
assert coll is not None
assert coll.stats.doc_count == 0
# Insert 100 documents
result = coll.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert coll.stats.doc_count == len(multiple_docs)
# Call optimize
from zvec import OptimizeOption
coll.optimize(option=OptimizeOption())
# Verify data is still accessible after optimize
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result_before_close = coll.query(queries=query, topk=10)
assert len(result_before_close) > 0
# Close collection (destroy will close it)
collection_path_str = str(collection_path)
del coll
# Reopen collection
reopened_coll = zvec.open(path=collection_path_str, option=collection_option)
assert reopened_coll is not None
assert reopened_coll.stats.doc_count == len(multiple_docs)
# Execute query on reopened collection
query_after_reopen = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result_after_reopen = reopened_coll.query(queries=query_after_reopen, topk=10)
assert len(result_after_reopen) > 0
assert len(result_after_reopen) <= 10
# Verify query results are valid
first_doc = result_after_reopen[0]
assert first_doc is not None
assert first_doc.id is not None
assert first_doc.field("id") is not None
assert first_doc.field("name") is not None
# Cleanup
reopened_coll.destroy()
python/tests/test_collection_fts.py
+188
View File
@@ -0,0 +1,188 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""End-to-end tests for FTS-only collections (no vector field).
The schema validation rule "must have at least one vector field" has been
lifted; these tests pin the new behavior so insert / query / delete /
optimize all work on a vector-less collection.
"""
from __future__ import annotations
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
OptimizeOption,
)
from zvec.model.param.query import Fts, Query
# ==================== Fixtures ====================
@pytest.fixture(scope="function")
def fts_collection(tmp_path_factory) -> Collection:
"""FTS-only collection: a STRING field for forward + an FTS-indexed STRING."""
temp_dir = tmp_path_factory.mktemp("zvec_fts_only")
collection_path = temp_dir / "fts_collection"
schema = zvec.CollectionSchema(
name="fts_only",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="standard",
filters=["lowercase"],
),
),
],
# vectors omitted on purpose — schema validation must accept this.
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def _make_docs() -> list[Doc]:
"""5-doc corpus where 4 contain 'hello' and doc 4 is the only outlier."""
return [
Doc(id="pk_0", fields={"title": "intro", "content": "hello world"}),
Doc(id="pk_1", fields={"title": "guide", "content": "hello foo bar"}),
Doc(id="pk_2", fields={"title": "tips", "content": "hello baz"}),
Doc(id="pk_3", fields={"title": "more", "content": "hello hello"}),
Doc(id="pk_4", fields={"title": "other", "content": "nothing relevant"}),
]
def _fts_query(coll: Collection, term: str) -> list[Doc]:
"""Run a single-term FTS match query against the `content` field."""
return coll.query(
queries=Query(field_name="content", fts=Fts(match_string=term)),
topk=10,
)
# ==================== Tests ====================
class TestFtsOnlyCollectionSchema:
def test_create_and_open_without_vectors(self, fts_collection: Collection):
"""Schema with zero vector fields must be accepted by validate()."""
assert fts_collection.schema.name == "fts_only"
assert {f.name for f in fts_collection.schema.fields} == {"title", "content"}
# Empty vectors is the whole point of the test.
assert list(fts_collection.schema.vectors) == []
assert fts_collection.stats.doc_count == 0
def test_create_schema_omitting_vectors_kwarg(self):
"""Constructing CollectionSchema without `vectors=` argument is valid."""
schema = zvec.CollectionSchema(
name="bare_fts",
fields=[
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(),
),
],
)
assert list(schema.vectors) == []
assert {f.name for f in schema.fields} == {"content"}
class TestFtsOnlyCollectionLifecycle:
def test_insert_and_fts_query(self, fts_collection: Collection):
"""FTS-only collection supports insert + FTS query end-to-end."""
results = fts_collection.insert(_make_docs())
assert all(r.ok() for r in results)
assert fts_collection.stats.doc_count == 5
hits = _fts_query(fts_collection, "hello")
assert len(hits) == 4
assert {doc.id for doc in hits} == {"pk_0", "pk_1", "pk_2", "pk_3"}
# Term that nothing in the surviving corpus contains.
assert _fts_query(fts_collection, "missing_term_xyz") == []
def test_delete_then_query(self, fts_collection: Collection):
"""Tombstone filter must drop deleted docs from FTS results."""
fts_collection.insert(_make_docs())
statuses = fts_collection.delete(["pk_0", "pk_4"])
assert all(s.ok() for s in statuses)
assert fts_collection.stats.doc_count == 3
hits = _fts_query(fts_collection, "hello")
assert len(hits) == 3
assert {doc.id for doc in hits} == {"pk_1", "pk_2", "pk_3"}
# pk_4's unique term is filtered out post-delete.
assert _fts_query(fts_collection, "nothing") == []
def test_optimize_rebuilds_fts(self, fts_collection: Collection):
"""Optimize with >30% deletes triggers ReduceFts; recall unchanged."""
fts_collection.insert(_make_docs())
# 40% delete ratio — above COMPACT_DELETE_RATIO_THRESHOLD=0.3, so
# build_compact_task picks the rebuild path and ReduceFts runs.
fts_collection.delete(["pk_0", "pk_4"])
before = {doc.id for doc in _fts_query(fts_collection, "hello")}
assert before == {"pk_1", "pk_2", "pk_3"}
fts_collection.optimize(option=OptimizeOption())
assert fts_collection.stats.doc_count == 3
after = {doc.id for doc in _fts_query(fts_collection, "hello")}
assert after == before
assert _fts_query(fts_collection, "nothing") == []
class TestFtsOnlyCollectionQueryValidation:
def test_vector_query_rejected(self, fts_collection: Collection):
"""Vector query on a no-vector collection must raise."""
with pytest.raises(ValueError, match="No vector field found"):
fts_collection.query(
queries=Query(field_name="content", vector=[0.1, 0.2, 0.3]),
topk=5,
)
def test_id_query_rejected(self, fts_collection: Collection):
"""ID-based query on a no-vector collection must raise."""
fts_collection.insert(_make_docs()[:1])
with pytest.raises(ValueError, match="No vector field found"):
fts_collection.query(
queries=Query(field_name="content", id="pk_0"),
topk=5,
)
python/tests/test_collection_fts_vector_hybrid.py
+391
View File
@@ -0,0 +1,391 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for FTS + vector hybrid retrieval via multi-query with reranker."""
from __future__ import annotations
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
HnswIndexParam,
VectorSchema,
)
from zvec.extension.multi_vector_reranker import RrfReRanker, WeightedReRanker
from zvec.model.param.query import Fts, Query
DIM = 16
# ==================== Fixtures ====================
@pytest.fixture(scope="function")
def hybrid_collection(tmp_path_factory) -> Collection:
"""Collection with one vector field + one FTS field."""
temp_dir = tmp_path_factory.mktemp("zvec_hybrid")
collection_path = temp_dir / "hybrid_collection"
schema = zvec.CollectionSchema(
name="hybrid_test",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="standard",
filters=["lowercase"],
),
),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=DIM,
index_param=HnswIndexParam(),
),
],
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def _make_docs() -> list[Doc]:
"""Corpus with both text content and vectors.
Docs 0-2: AI/ML topic, vectors clustered in one region.
Docs 3-4: retrieval topic, vectors clustered in another region.
Doc 5: unrelated topic.
"""
# AI cluster vectors
ai_vec = [1.0] * 8 + [0.0] * 8
# Retrieval cluster vectors
ret_vec = [0.0] * 8 + [1.0] * 8
# Unrelated vector
other_vec = [0.5] * 16
return [
Doc(
id="pk_0",
fields={
"title": "ML Intro",
"content": "machine learning is a branch of artificial intelligence",
},
vectors={"embedding": ai_vec},
),
Doc(
id="pk_1",
fields={
"title": "Deep Learning",
"content": "deep learning uses neural networks for pattern recognition",
},
vectors={"embedding": [0.9] * 8 + [0.1] * 8},
),
Doc(
id="pk_2",
fields={
"title": "NLP",
"content": "natural language processing handles text with artificial intelligence",
},
vectors={"embedding": [0.8] * 8 + [0.2] * 8},
),
Doc(
id="pk_3",
fields={
"title": "Search Engine",
"content": "search engine uses inverted index for text retrieval",
},
vectors={"embedding": ret_vec},
),
Doc(
id="pk_4",
fields={
"title": "Vector DB",
"content": "vector database enables similarity retrieval and search",
},
vectors={"embedding": [0.1] * 8 + [0.9] * 8},
),
Doc(
id="pk_5",
fields={
"title": "Cooking",
"content": "baking bread requires flour water yeast and salt",
},
vectors={"embedding": other_vec},
),
]
@pytest.fixture(scope="function")
def hybrid_collection_with_docs(hybrid_collection: Collection) -> Collection:
"""Hybrid collection pre-populated with test documents."""
results = hybrid_collection.insert(_make_docs())
assert all(r.ok() for r in results)
return hybrid_collection
# ==================== Tests ====================
class TestFtsVectorHybridQuery:
"""Test FTS + vector hybrid retrieval using multi-query with RRF reranker."""
def test_hybrid_fts_and_vector_basic(self, hybrid_collection_with_docs: Collection):
"""FTS + vector multi-query with RRF reranker returns results."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
# Results should have scores
for doc in result:
assert doc.score > 0
def test_hybrid_fts_and_vector_ranking(
self, hybrid_collection_with_docs: Collection
):
"""Docs relevant in both FTS and vector should rank higher."""
reranker = RrfReRanker(rank_constant=60)
# FTS: "retrieval search" matches pk_3, pk_4
# Vector: ret_vec cluster matches pk_3, pk_4
# Both signals agree: pk_3 and pk_4 should rank top
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval search")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
top_ids = {doc.id for doc in result[:3]}
assert "pk_3" in top_ids or "pk_4" in top_ids
def test_hybrid_scores_descending(self, hybrid_collection_with_docs: Collection):
"""Hybrid query results must be sorted by score descending."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="intelligence")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=6,
reranker=reranker,
)
assert len(result) >= 2
scores = [doc.score for doc in result]
assert scores == sorted(scores, reverse=True)
def test_hybrid_with_filter(self, hybrid_collection_with_docs: Collection):
"""Hybrid query respects SQL filter."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=10,
reranker=reranker,
filter="title like '%Learning%'",
)
for doc in result:
assert "Learning" in doc.fields["title"]
def test_hybrid_fts_no_match_still_returns_vector_results(
self, hybrid_collection_with_docs: Collection
):
"""When FTS matches nothing, vector results still appear."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(
field_name="content",
fts=Fts(match_string="nonexistent_term_xyz"),
),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker,
)
# Vector query alone should still produce results
assert len(result) > 0
def test_hybrid_query_string_syntax(self, hybrid_collection_with_docs: Collection):
"""Hybrid query works with FTS query_string (advanced syntax)."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(
field_name="content",
fts=Fts(query_string="artificial AND intelligence"),
),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
# pk_0 and pk_2 contain "artificial intelligence"
hit_ids = {doc.id for doc in result}
assert "pk_0" in hit_ids or "pk_2" in hit_ids
class TestFtsVectorHybridValidation:
"""Test validation rules for FTS + vector hybrid queries."""
def test_hybrid_requires_reranker(self, hybrid_collection_with_docs: Collection):
"""Multi-query with FTS + vector without reranker should raise."""
with pytest.raises(ValueError, match="[Rr]eranker"):
hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="embedding", vector=[1.0] * DIM),
],
topk=5,
)
def test_duplicate_field_name_allowed(
self, hybrid_collection_with_docs: Collection
):
"""Multi-query with duplicate field names is allowed and returns results."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="content", fts=Fts(match_string="intelligence")),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
def test_multiple_vectors_allowed(self, hybrid_collection_with_docs: Collection):
"""Two vector queries on the same field are allowed with a reranker."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="embedding", vector=[1.0] * DIM),
Query(field_name="embedding", vector=[0.5] * DIM),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
class TestFtsVectorHybridWeightedReranker:
"""Test FTS + vector hybrid retrieval using WeightedReranker."""
def test_weighted_reranker_fts_and_vector(
self, hybrid_collection_with_docs: Collection
):
"""WeightedReranker correctly normalizes FTS scores alongside vector scores."""
weights = [0.5, 0.5]
reranker = WeightedReRanker(weights=weights)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval search")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
for doc in result:
assert doc.score > 0
def test_weighted_reranker_scores_descending(
self, hybrid_collection_with_docs: Collection
):
"""WeightedReranker hybrid results are sorted by score descending."""
weights = [0.4, 0.6]
reranker = WeightedReRanker(weights=weights)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="intelligence")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=6,
reranker=reranker,
)
assert len(result) >= 2
scores = [doc.score for doc in result]
assert scores == sorted(scores, reverse=True)
def test_weighted_reranker_fts_weight_influence(
self, hybrid_collection_with_docs: Collection
):
"""Higher FTS weight should boost FTS-relevant docs in ranking."""
# High FTS weight: FTS signal dominates
weights_fts_heavy = [0.9, 0.1]
reranker_fts = WeightedReRanker(weights=weights_fts_heavy)
result_fts = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker_fts,
)
# High vector weight: vector signal dominates
weights_vec_heavy = [0.1, 0.9]
reranker_vec = WeightedReRanker(weights=weights_vec_heavy)
result_vec = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker_vec,
)
# Both should return results
assert len(result_fts) > 0
assert len(result_vec) > 0
# With FTS-heavy weight, FTS-relevant docs (pk_3, pk_4) should rank higher
fts_top = [doc.id for doc in result_fts[:2]]
vec_top = [doc.id for doc in result_vec[:2]]
# The rankings should differ due to weight difference
assert fts_top != vec_top or len(result_fts) == len(result_vec) == 1
python/tests/test_collection_hnsw_rabitq.py
+574
View File
@@ -0,0 +1,574 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import platform
import sys
import pytest
import math
import zvec
pytestmark = pytest.mark.skipif(
not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
reason="HNSW RaBitQ only supported on Linux x86_64",
)
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
MetricType,
VectorSchema,
Query,
)
# ==================== Fixtures ====================
@pytest.fixture(scope="session")
def hnsw_rabitq_collection_schema():
"""Create a collection schema with HNSW RaBitQ index."""
return zvec.CollectionSchema(
name="test_hnsw_rabitq_collection",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
FieldSchema("name", DataType.STRING, nullable=False),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswRabitqIndexParam(
metric_type=MetricType.L2,
m=16,
ef_construction=200,
total_bits=7,
num_clusters=64,
),
),
],
)
@pytest.fixture(scope="session")
def collection_option():
"""Create collection options."""
return CollectionOption(read_only=False, enable_mmap=True)
@pytest.fixture
def single_doc():
"""Create a single document for testing."""
return Doc(
id="0",
fields={"id": 0, "name": "test_doc_0"},
vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]},
)
@pytest.fixture
def multiple_docs():
"""Create multiple documents for testing."""
return [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_doc_{i}"},
vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]},
)
for i in range(1, 101)
]
@pytest.fixture(scope="function")
def hnsw_rabitq_collection(
tmp_path_factory, hnsw_rabitq_collection_schema, collection_option
) -> Collection:
"""
Function-scoped fixture: creates and opens a collection with HNSW RaBitQ index.
"""
temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq")
collection_path = temp_dir / "test_hnsw_rabitq_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=hnsw_rabitq_collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open HNSW RaBitQ collection"
assert coll.path == str(collection_path)
assert coll.schema.name == hnsw_rabitq_collection_schema.name
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
@pytest.fixture
def collection_with_single_doc(
hnsw_rabitq_collection: Collection, single_doc: Doc
) -> Collection:
"""Setup: insert single doc into collection."""
assert hnsw_rabitq_collection.stats.doc_count == 0
result = hnsw_rabitq_collection.insert(single_doc)
assert bool(result)
assert result.ok()
assert hnsw_rabitq_collection.stats.doc_count == 1
yield hnsw_rabitq_collection
# Teardown: delete single doc
hnsw_rabitq_collection.delete(single_doc.id)
assert hnsw_rabitq_collection.stats.doc_count == 0
@pytest.fixture
def collection_with_multiple_docs(
hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]
) -> Collection:
"""Setup: insert multiple docs into collection."""
assert hnsw_rabitq_collection.stats.doc_count == 0
result = hnsw_rabitq_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert hnsw_rabitq_collection.stats.doc_count == len(multiple_docs)
yield hnsw_rabitq_collection
# Teardown: delete multiple docs
hnsw_rabitq_collection.delete([doc.id for doc in multiple_docs])
# ==================== Tests ====================
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionCreation:
"""Test HNSW RaBitQ collection creation and schema validation."""
def test_collection_creation(
self, hnsw_rabitq_collection: Collection, hnsw_rabitq_collection_schema
):
"""Test that collection is created with correct schema."""
assert hnsw_rabitq_collection is not None
assert hnsw_rabitq_collection.schema.name == hnsw_rabitq_collection_schema.name
assert len(hnsw_rabitq_collection.schema.fields) == len(
hnsw_rabitq_collection_schema.fields
)
assert len(hnsw_rabitq_collection.schema.vectors) == len(
hnsw_rabitq_collection_schema.vectors
)
def test_vector_schema_validation(self, hnsw_rabitq_collection: Collection):
"""Test that vector schema has correct HNSW RaBitQ configuration."""
vector_schema = hnsw_rabitq_collection.schema.vector("embedding")
assert vector_schema is not None
assert vector_schema.name == "embedding"
assert vector_schema.data_type == DataType.VECTOR_FP32
assert vector_schema.dimension == 128
index_param = vector_schema.index_param
assert index_param is not None
assert index_param.metric_type == MetricType.L2
assert index_param.m == 16
assert index_param.ef_construction == 200
assert index_param.total_bits == 7
assert index_param.num_clusters == 64
def test_collection_stats(self, hnsw_rabitq_collection: Collection):
"""Test initial collection statistics."""
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == 0
assert len(stats.index_completeness) == 1
assert stats.index_completeness["embedding"] == 1
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionInsert:
"""Test document insertion into HNSW RaBitQ collection."""
def test_insert_single_doc(
self, hnsw_rabitq_collection: Collection, single_doc: Doc
):
"""Test inserting a single document."""
result = hnsw_rabitq_collection.insert(single_doc)
assert bool(result)
assert result.ok()
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == 1
def test_insert_multiple_docs(
self, hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]
):
"""Test inserting multiple documents."""
result = hnsw_rabitq_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == len(multiple_docs)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionFetch:
"""Test document fetching from HNSW RaBitQ collection."""
def test_fetch_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test fetching a single document by ID."""
result = collection_with_single_doc.fetch(ids=[single_doc.id])
assert bool(result)
assert single_doc.id in result.keys()
doc = result[single_doc.id]
assert doc is not None
assert doc.id == single_doc.id
assert doc.field("id") == single_doc.field("id")
assert doc.field("name") == single_doc.field("name")
def test_fetch_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test fetching multiple documents by IDs."""
ids = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.fetch(ids=ids)
assert bool(result)
assert len(result) == len(ids)
for doc_id in ids:
assert doc_id in result
doc = result[doc_id]
assert doc is not None
assert doc.id == doc_id
def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):
"""Test fetching a non-existent document."""
result = collection_with_single_doc.fetch(ids=["nonexistent_id"])
assert len(result) == 0
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionQuery:
"""Test vector search queries on HNSW RaBitQ collection."""
def test_query_by_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by vector with HNSW RaBitQ index."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
# First result should be the query document itself (or very close)
first_doc = result[0]
assert first_doc is not None
assert first_doc.id is not None
def test_query_by_id(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by document ID with HNSW RaBitQ index."""
query = Query(
field_name="embedding",
id=multiple_docs[0].id,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
def test_query_with_different_ef_values(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different ef parameter values."""
query_vector = multiple_docs[0].vector("embedding")
# Test with ef=100
query_100 = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=100),
)
result_100 = collection_with_multiple_docs.query(queries=query_100, topk=10)
assert len(result_100) > 0
# Test with ef=500
query_500 = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=500),
)
result_500 = collection_with_multiple_docs.query(queries=query_500, topk=10)
assert len(result_500) > 0
def test_query_with_topk(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different topk values."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
# Test topk=5
result_5 = collection_with_multiple_docs.query(queries=query, topk=5)
assert len(result_5) <= 5
# Test topk=20
result_20 = collection_with_multiple_docs.query(queries=query, topk=20)
assert len(result_20) <= 20
def test_query_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with filter conditions."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
# Query with id filter
result = collection_with_multiple_docs.query(
queries=query, topk=10, filter="id < 50"
)
assert len(result) > 0
for doc in result:
assert doc.field("id") < 50
def test_query_with_output_fields(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with specific output fields."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, output_fields=["id", "name"]
)
assert len(result) > 0
first_doc = result[0]
assert "id" in first_doc.field_names()
assert "name" in first_doc.field_names()
def test_query_with_include_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with vector data included in results."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, include_vector=True
)
assert len(result) > 0
first_doc = result[0]
assert first_doc.vector("embedding") is not None
assert len(first_doc.vector("embedding")) == 128
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionUpdate:
"""Test document update in HNSW RaBitQ collection."""
def test_update_doc_fields(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document fields."""
updated_doc = Doc(
id=single_doc.id,
fields={"id": single_doc.field("id"), "name": "updated_name"},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(ids=[single_doc.id])
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.field("name") == "updated_name"
def test_update_doc_vector(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document vector."""
new_vector = [0.5 + i * 0.01 for i in range(128)]
updated_doc = Doc(
id=single_doc.id,
vectors={"embedding": new_vector},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(
ids=[single_doc.id],
)
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.vector("embedding") is not None
embedding = doc.vector("embedding")
assert len(embedding) == 128
# Verify vector values are approximately equal (float comparison)
for i in range(128):
assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionDelete:
"""Test document deletion from HNSW RaBitQ collection."""
def test_delete_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test deleting a single document."""
result = collection_with_single_doc.delete(single_doc.id)
assert bool(result)
assert result.ok()
stats = collection_with_single_doc.stats
assert stats.doc_count == 0
def test_delete_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test deleting multiple documents."""
ids_to_delete = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.delete(ids_to_delete)
assert len(result) == len(ids_to_delete)
for item in result:
assert item.ok()
stats = collection_with_multiple_docs.stats
assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionOptimizeAndReopen:
"""Test collection optimize and reopen functionality."""
def test_optimize_close_reopen_and_query(
self,
tmp_path_factory,
hnsw_rabitq_collection_schema,
collection_option,
multiple_docs: list[Doc],
):
"""Test inserting 100 docs, optimize, close, reopen and query."""
# Create collection and insert 100 documents
temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq_optimize")
collection_path = temp_dir / "test_optimize_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=hnsw_rabitq_collection_schema,
option=collection_option,
)
assert coll is not None
assert coll.stats.doc_count == 0
# Insert 100 documents
result = coll.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert coll.stats.doc_count == len(multiple_docs)
# Call optimize
from zvec import OptimizeOption
coll.optimize(option=OptimizeOption())
# Verify data is still accessible after optimize
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result_before_close = coll.query(query, topk=10)
assert len(result_before_close) > 0
# Close collection (destroy will close it)
collection_path_str = str(collection_path)
del coll
# Reopen collection
reopened_coll = zvec.open(path=collection_path_str, option=collection_option)
assert reopened_coll is not None
assert reopened_coll.stats.doc_count == len(multiple_docs)
# Execute query on reopened collection
query_after_reopen = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result_after_reopen = reopened_coll.query(query_after_reopen, topk=10)
assert len(result_after_reopen) > 0
assert len(result_after_reopen) <= 10
# Verify query results are valid
first_doc = result_after_reopen[0]
assert first_doc is not None
assert first_doc.id is not None
assert first_doc.field("id") is not None
assert first_doc.field("name") is not None
# Cleanup
reopened_coll.destroy()
python/tests/test_convert.py
+584
View File
@@ -0,0 +1,584 @@
from __future__ import annotations
import math
import pytest
from zvec._zvec import _Doc
from zvec.model.convert import convert_to_py_doc, convert_to_cpp_doc
from zvec import Doc, CollectionSchema, DataType, FieldSchema, VectorSchema
# ----------------------------
# Convert Cpp Doc Test Case
# ----------------------------
class TestConvertCppDoc:
def test_default(self):
doc = Doc(id="1")
schema = CollectionSchema(
name="test_collection",
fields=FieldSchema("name", DataType.STRING),
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
def test_with_field_notin_schema(self):
doc = Doc(id="1", fields={"name": "Tom"})
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
],
)
with pytest.raises(ValueError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_scalar_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = Doc(
id="1",
fields={
"id": 1,
"salary": 1000,
"age": 18,
"create_at": 1640995200,
"bmi": 80.0 / 200.0,
"author": "Tom",
"weight": 80.0,
"is_male": True,
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
assert cpp_doc.get_any("id", DataType.UINT64) == 1
assert cpp_doc.get_any("salary", DataType.UINT32) == 1000
assert cpp_doc.get_any("age", DataType.INT32) == 18
assert cpp_doc.get_any("create_at", DataType.INT64) == 1640995200
assert cpp_doc.get_any("author", DataType.STRING) == "Tom"
assert math.isclose(
cpp_doc.get_any("weight", DataType.FLOAT), 80.0, rel_tol=1e-6
)
assert math.isclose(
cpp_doc.get_any("bmi", DataType.DOUBLE), 80.0 / 200.0, rel_tol=1e-6
)
assert cpp_doc.get_any("is_male", DataType.BOOL) == True
def test_with_array_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = Doc(
id="1",
fields={
"tags": ["tag1", "tag2", "tag3"],
"ids": [111111111111, 222222222222, 333333333333],
"marks": [100, 200, 300],
"x": [1, 2, 3],
"y": [100, 200, 300],
"scores": [1.1, 2.2, 3.3],
"ratios": [0.1, 0.2, 0.3],
"results": [True, False, True],
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
assert cpp_doc.get_any("tags", DataType.ARRAY_STRING) == doc.field("tags")
assert cpp_doc.get_any("ids", DataType.ARRAY_UINT64) == doc.field("ids")
assert cpp_doc.get_any("marks", DataType.ARRAY_UINT32) == doc.field("marks")
assert cpp_doc.get_any("x", DataType.ARRAY_INT32) == doc.field("x")
assert cpp_doc.get_any("y", DataType.ARRAY_INT64) == doc.field("y")
scores = cpp_doc.get_any("scores", DataType.ARRAY_FLOAT)
for i in range(len(doc.field("scores"))):
assert math.isclose(scores[i], doc.field("scores")[i], rel_tol=1e-1)
ratios = cpp_doc.get_any("ratios", DataType.ARRAY_DOUBLE)
for i in range(len(doc.field("ratios"))):
assert math.isclose(ratios[i], doc.field("ratios")[i], rel_tol=1e-1)
results = cpp_doc.get_any("results", DataType.ARRAY_BOOL)
for i in range(len(doc.field("results"))):
assert results[i] == doc.field("results")[i]
def test_with_dense_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = Doc(
id="1",
vectors={
"embedding": [1.1] * 4,
"image": [2.2] * 8,
"text": [4] * 32,
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
embedding_vector = cpp_doc.get_any("embedding", DataType.VECTOR_FP16)
assert len(embedding_vector) == 4
for i in range(4):
assert math.isclose(
embedding_vector[i], doc.vector("embedding")[i], rel_tol=1e-1
)
image_vector = cpp_doc.get_any("image", DataType.VECTOR_FP32)
assert len(image_vector) == 8
for i in range(8):
assert math.isclose(image_vector[i], doc.vector("image")[i], rel_tol=1e-1)
text_vector = cpp_doc.get_any("text", DataType.VECTOR_INT8)
assert len(text_vector) == 32
for i in range(32):
assert text_vector[i] == doc.vectors["text"][i]
def test_with_sparse_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = Doc(
id="1",
vectors={
"author": {1: 1.1, 2: 2.2, 3: 3.3},
"content": {4: 4.4, 5: 5.5, 6: 6.6},
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
author_vector = cpp_doc.get_any("author", DataType.SPARSE_VECTOR_FP32)
assert isinstance(author_vector, dict)
for key, value in doc.vector("author").items():
assert math.isclose(author_vector[key], value, rel_tol=1e-1)
content_vector = cpp_doc.get_any("content", DataType.SPARSE_VECTOR_FP16)
assert isinstance(content_vector, dict)
for key, value in doc.vector("content").items():
assert math.isclose(content_vector[key], value, rel_tol=1e-1)
def test_with_scalar_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = Doc(
id="1",
fields={
"id": "1",
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"salary": "1000"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"age": "18"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"create_at": "2021-01-01"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"author": 1})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"weight": "80.5"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"bmi": "25.0"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"is_male": "true"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_array_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = Doc(id="1", fields={"tags": [1, 2, 3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"ids": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"marks": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"x": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"y": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"scores": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"ratios": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"results": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_vector_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = Doc(id="1", vectors={"image": ["1.1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", vectors={"text": ["1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", vectors={"embedding": ["1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_sparse_vector_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = Doc(
id="1",
vectors={
"author": {"1": 1.1, "2": 2.2, "3": 3.3},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(
id="1",
vectors={
"content": {"1": 1.1, "2": 2.2, "3": 3.3},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(
id="1",
vectors={
"author": {1: "1", 2: "2", 3: "3"},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
# ----------------------------
# Convert Py Doc Test Case
# ----------------------------
class TestConvertPyDoc:
def test_default(self):
doc = _Doc()
doc.set_pk("1")
doc.set_score(1.0)
schema = CollectionSchema(
name="test_collection",
fields=FieldSchema("name", DataType.STRING),
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
assert py_doc.score == 1.0
def test_with_scalar_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any("id", schema.field("id")._get_object(), 1111111111111111)
doc.set_any("salary", schema.field("salary")._get_object(), 1000)
doc.set_any("age", schema.field("age")._get_object(), 18)
doc.set_any("create_at", schema.field("create_at")._get_object(), 1640995200)
doc.set_any("author", schema.field("author")._get_object(), "Tom")
doc.set_any("weight", schema.field("weight")._get_object(), 80.0)
doc.set_any("bmi", schema.field("bmi")._get_object(), 80.0 / 200.0)
doc.set_any("is_male", schema.field("is_male")._get_object(), True)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
assert py_doc.field("id") == 1111111111111111
assert py_doc.field("salary") == 1000
assert py_doc.field("age") == 18
assert py_doc.field("create_at") == 1640995200
assert py_doc.field("author") == "Tom"
assert py_doc.field("weight") == 80.0
assert py_doc.field("bmi") == 80.0 / 200.0
assert py_doc.field("is_male") == True
def test_with_array_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any(
"tags", schema.field("tags")._get_object(), ["tag1", "tag2", "tag3"]
)
doc.set_any(
"ids",
schema.field("ids")._get_object(),
[111111111111, 222222222222, 3333333333333],
)
doc.set_any("marks", schema.field("marks")._get_object(), [1000, 2000, 3000])
doc.set_any("x", schema.field("x")._get_object(), [1, 2, 3])
doc.set_any("y", schema.field("y")._get_object(), [100, 200, 300])
doc.set_any("scores", schema.field("scores")._get_object(), [0.1, 0.2, 0.3])
doc.set_any("ratios", schema.field("ratios")._get_object(), [0.1, 0.2, 0.3])
doc.set_any(
"results", schema.field("results")._get_object(), [True, False, True]
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.field("tags") == ["tag1", "tag2", "tag3"]
assert py_doc.field("ids") == [111111111111, 222222222222, 3333333333333]
assert py_doc.field("marks") == [1000, 2000, 3000]
assert py_doc.field("x") == [1, 2, 3]
assert py_doc.field("y") == [100, 200, 300]
scores = doc.get_any("scores", DataType.ARRAY_FLOAT)
for i in range(len(scores)):
assert math.isclose(scores[i], py_doc.field("scores")[i], rel_tol=1e-1)
ratios = doc.get_any("ratios", DataType.ARRAY_DOUBLE)
for i in range(len(ratios)):
assert math.isclose(ratios[i], py_doc.field("ratios")[i], rel_tol=1e-1)
results = doc.get_any("results", DataType.ARRAY_BOOL)
for i in range(len(results)):
assert results[i] == py_doc.field("results")[i]
def test_with_dense_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any("embedding", schema.vector("embedding")._get_object(), [1.1] * 4)
doc.set_any("image", schema.vector("image")._get_object(), [2.2] * 8)
doc.set_any("text", schema.vector("text")._get_object(), [4] * 32)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
embedding_vector = py_doc.vector("embedding")
assert len(embedding_vector) == 4
for i in range(4):
assert math.isclose(
py_doc.vector("embedding")[i], embedding_vector[i], rel_tol=1e-1
)
image_vector = py_doc.vector("image")
assert len(image_vector) == 8
for i in range(8):
assert math.isclose(
py_doc.vector("image")[i], image_vector[i], rel_tol=1e-1
)
text_vector = py_doc.vector("text")
assert len(text_vector) == 32
for i in range(32):
assert py_doc.vector("text")[i] == text_vector[i]
def test_with_sparse_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any(
"author", schema.vector("author")._get_object(), {1: 1.1, 2: 2.2, 3: 3.3}
)
doc.set_any(
"content", schema.vector("content")._get_object(), {4: 4.4, 5: 5.5, 6: 6.6}
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
author_vector = py_doc.vector("author")
assert isinstance(author_vector, dict)
for key, value in doc.get_any("author", DataType.SPARSE_VECTOR_FP32).items():
assert math.isclose(author_vector[key], value, rel_tol=1e-1)
content_vector = py_doc.vector("content")
assert isinstance(content_vector, dict)
for key, value in doc.get_any("content", DataType.SPARSE_VECTOR_FP16).items():
assert math.isclose(content_vector[key], value, rel_tol=1e-1)
python/tests/test_doc.py
+269
View File
@@ -0,0 +1,269 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import math
import pytest
from zvec._zvec import _Doc
from zvec import FieldSchema, VectorSchema, Doc, DataType
# ----------------------------
# PyDoc Test Case
# ----------------------------
class TestPyDoc:
def test_default(self):
Doc(id="1")
def test_with_single_vector(self):
doc = Doc(id="1", vectors={"dense": [1, 2, 3]})
assert doc is not None
assert doc.id == "1"
assert doc.vector("dense") == [1, 2, 3]
def test_with_hybrid_vectors(self):
doc = Doc(
id="1", vectors={"dense": [1, 2, 3], "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("dense") == [1, 2, 3]
assert doc.vector("sparse") == {1: 1.0, 2: 2.0, 3: 3.0}
def test_with_multi_vectors(self):
doc = Doc(
id="1",
vectors={
"image": [1, 2, 3],
"description": [4, 5, 6],
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
},
fields={"author": "Tom", "age": 19, "is_male": True, "weight": 60.5},
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("image") == [1, 2, 3]
assert doc.vector("description") == [4, 5, 6]
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
assert doc.field("author") == "Tom"
assert doc.field("age") == 19
assert doc.field("is_male") == True
assert doc.field("weight") == 60.5
def test_with_numpy_array(self):
import numpy as np
doc = Doc._from_tuple(
(
"1",
0.0,
None,
{
"image": np.array([1, 2, 3]),
"description": np.random.random(512),
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
},
)
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("image") == [1, 2, 3]
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
# ----------------------------
# CppDoc Test Case
# ----------------------------
class TestCppDoc:
def test_default(self):
doc = _Doc()
assert doc is not None
def test_doc_set_pk(self):
doc = _Doc()
doc.set_pk("1")
assert doc.pk() == "1"
def test_doc_set_score(self):
doc = _Doc()
doc.set_score(0.9)
assert math.isclose(doc.score(), 0.9, rel_tol=1e-6)
def test_doc_get_null_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING, nullable=True)
doc.set_any("author", schema._get_object(), None)
assert doc.has_field("author")
assert doc.get_any("author", schema.data_type) is None
def test_doc_get_set_has_null_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING, nullable=False)
with pytest.raises(ValueError):
doc.set_any("author", schema._get_object(), None)
def test_doc_get_set_has_string_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING)
doc.set_any("author", schema._get_object(), "Tom")
assert doc.has_field("author")
assert doc.get_any("author", DataType.STRING) == "Tom"
def test_doc_get_set_has_bool_field(self):
doc = _Doc()
schema = FieldSchema("is_male", DataType.BOOL)
doc.set_any("is_male", schema._get_object(), True)
assert doc.has_field("is_male")
assert doc.get_any("is_male", DataType.BOOL) == True
def test_doc_get_set_has_int32_field(self):
doc = _Doc()
schema = FieldSchema("age", DataType.INT32)
doc.set_any("age", schema._get_object(), 19)
assert doc.has_field("age")
assert doc.get_any("age", DataType.INT32) == 19
def test_doc_get_set_has_int64_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.INT64)
doc.set_any("id", schema._get_object(), 1111111111111111111)
assert doc.has_field("id")
assert doc.get_any("id", DataType.INT64) == 1111111111111111111
def test_doc_get_set_has_float_field(self):
doc = _Doc()
schema = FieldSchema("weight", DataType.FLOAT)
doc.set_any("weight", schema._get_object(), 60.5)
assert doc.has_field("weight")
assert math.isclose(doc.get_any("weight", DataType.FLOAT), 60.5, rel_tol=1e-6)
def test_doc_get_set_has_double_field(self):
doc = _Doc()
schema = FieldSchema("height", DataType.DOUBLE)
doc.set_any("height", schema._get_object(), 1.77777777777)
assert doc.has_field("height")
assert math.isclose(
doc.get_any("height", DataType.DOUBLE), 1.7777777777, rel_tol=1e-9
)
def test_doc_get_set_has_uint32_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.UINT32)
doc.set_any("id", schema._get_object(), 4294967295)
assert doc.has_field("id")
assert doc.get_any("id", DataType.UINT32) == 4294967295
def test_doc_get_set_has_uint64_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.UINT64)
doc.set_any("id", schema._get_object(), 18446744073709551615)
assert doc.has_field("id")
assert doc.get_any("id", DataType.UINT64) == 18446744073709551615
def test_doc_get_set_has_array_string_field(self):
doc = _Doc()
schema = FieldSchema("tags", DataType.ARRAY_STRING)
doc.set_any("tags", schema._get_object(), ["tag1", "tag2", "tag3"])
assert doc.has_field("tags")
assert doc.get_any("tags", DataType.ARRAY_STRING) == ["tag1", "tag2", "tag3"]
def test_doc_get_set_has_array_int32_field(self):
doc = _Doc()
schema = FieldSchema("ids", DataType.ARRAY_INT32)
doc.set_any("ids", schema._get_object(), [1, 2, 3])
assert doc.has_field("ids")
assert doc.get_any("ids", DataType.ARRAY_INT32) == [1, 2, 3]
def test_doc_get_set_has_array_int64_field(self):
doc = _Doc()
schema = FieldSchema("ids", DataType.ARRAY_INT64)
doc.set_any("ids", schema._get_object(), [1, 2, 3])
assert doc.has_field("ids")
assert doc.get_any("ids", DataType.ARRAY_INT64) == [1, 2, 3]
def test_doc_get_set_has_array_float_field(self):
doc = _Doc()
schema = FieldSchema("weights", DataType.ARRAY_FLOAT)
doc.set_any("weights", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("weights")
assert doc.get_any("weights", DataType.ARRAY_FLOAT) == [1.0, 2.0, 3.0]
def test_doc_get_set_has_array_double_field(self):
doc = _Doc()
schema = FieldSchema("heights", DataType.ARRAY_DOUBLE)
doc.set_any("heights", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("heights")
assert doc.get_any("heights", DataType.ARRAY_DOUBLE) == [1.0, 2.0, 3.0]
def test_doc_get_set_has_array_bool_field(self):
doc = _Doc()
schema = FieldSchema("bools", DataType.ARRAY_BOOL)
doc.set_any("bools", schema._get_object(), [True, False, True])
assert doc.has_field("bools")
assert doc.get_any("bools", DataType.ARRAY_BOOL) == [True, False, True]
def test_doc_get_set_has_vector_fp16(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_FP16)
doc.set_any("image", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("image")
image_vector = doc.get_any("image", DataType.VECTOR_FP16)
assert image_vector is not None
for i in range(len(image_vector)):
assert math.isclose(image_vector[i], [1.0, 2.0, 3.0][i], rel_tol=1e-6)
def test_doc_get_set_has_vector_fp32(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_FP32)
doc.set_any("image", schema._get_object(), [1.111111, 2.222222, 3.333333])
assert doc.has_field("image")
vector = doc.get_any("image", DataType.VECTOR_FP32)
assert vector is not None
for i in range(len(vector)):
assert math.isclose(
vector[i], [1.111111, 2.222222, 3.333333][i], rel_tol=1e-6
)
def test_doc_get_set_has_vector_int8(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_INT8)
doc.set_any("image", schema._get_object(), [1, 2, 3])
assert doc.has_field("image")
assert doc.get_any("image", DataType.VECTOR_INT8) == [1, 2, 3]
def test_doc_get_set_has_sparse_vector_fp32(self):
doc = _Doc()
sparse = {1: 1.111111, 2: 2.222222, 3: 3.333333}
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP32)
doc.set_any("key", schema._get_object(), sparse)
assert doc.has_field("key")
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP32)
assert vector is not None
assert isinstance(vector, dict)
for key, value in sparse.items():
assert math.isclose(vector[key], value, rel_tol=1e-6)
def test_doc_get_set_has_sparse_vector_fp16(self):
doc = _Doc()
sparse = {1: 1.1, 2: 2.2, 3: 3.3}
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP16)
doc.set_any("key", schema._get_object(), sparse)
assert doc.has_field("key")
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP16)
assert vector is not None
assert isinstance(vector, dict)
for key, value in sparse.items():
assert math.isclose(vector[key], value, rel_tol=1e-1)
python/tests/test_embedding.py
+2099
View File
File diff suppressed because it is too large Load Diff
python/tests/test_fts_query.py
+158
View File
@@ -0,0 +1,158 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for FTS (Full-Text Search) query support in the Python SDK."""
import pickle
import pytest
from zvec.model.param.query import Fts, Query
class TestFtsQueryValidation:
"""Test FTS parameter validation in Query dataclass."""
def test_fts_query_string_only(self):
"""Query with only query_string in Fts should be valid."""
q = Query(
field_name="content", fts=Fts(query_string='+hello -world "exact phrase"')
)
q._validate()
assert q.fts.query_string == '+hello -world "exact phrase"'
assert q.fts.match_string is None
assert q.has_fts() is True
def test_fts_match_string_only(self):
"""Query with only match_string in Fts should be valid."""
q = Query(field_name="content", fts=Fts(match_string="machine learning"))
q._validate()
assert q.fts.match_string == "machine learning"
assert q.fts.query_string is None
assert q.has_fts() is True
def test_fts_query_string_and_match_string_mutually_exclusive(self):
"""Cannot provide both query_string and match_string in Fts."""
q = Query(
field_name="content",
fts=Fts(query_string="+hello", match_string="hello world"),
)
with pytest.raises(ValueError, match="mutually exclusive"):
q._validate()
def test_no_fts(self):
"""Query without FTS fields should have has_fts() == False."""
q = Query(field_name="embedding", vector=[0.1, 0.2, 0.3])
assert q.has_fts() is False
def test_vector_and_fts_mutually_exclusive(self):
"""Cannot combine vector search with FTS in a single Query."""
q = Query(
field_name="embedding",
vector=[0.1, 0.2, 0.3],
fts=Fts(match_string="deep learning"),
)
with pytest.raises(ValueError, match="Cannot combine fts with vector search"):
q._validate()
def test_fts_without_vector_or_id(self):
"""Query with only FTS (no vector, no id) should be valid."""
q = Query(field_name="content", fts=Fts(query_string="hello"))
q._validate()
assert q.has_vector() is False
assert q.has_id() is False
assert q.has_fts() is True
class TestFtsQueryBinding:
"""Test FTS binding layer (_Fts)."""
def test_import_fts_query(self):
"""_Fts should be importable from _zvec.param."""
from zvec._zvec.param import _Fts
fts = _Fts()
assert fts.query_string == ""
assert fts.match_string == ""
def test_fts_query_set_fields(self):
"""Setting fields on _Fts should work."""
from zvec._zvec.param import _Fts
fts = _Fts()
fts.query_string = "+hello -world"
assert fts.query_string == "+hello -world"
fts2 = _Fts()
fts2.match_string = "machine learning"
assert fts2.match_string == "machine learning"
def test_fts_query_pickle(self):
"""_Fts should support pickling."""
from zvec._zvec.param import _Fts
fts = _Fts()
fts.query_string = "+vector search"
fts.match_string = ""
data = pickle.dumps(fts)
restored = pickle.loads(data)
assert restored.query_string == "+vector search"
assert restored.match_string == ""
def test_search_query_fts_field(self):
"""_SearchQuery should have fts field."""
from zvec._zvec.param import _Fts, _SearchQuery
vq = _SearchQuery()
# fts should be None by default (optional)
assert vq.fts is None
# set fts
fts = _Fts()
fts.query_string = "hello"
vq.fts = fts
assert vq.fts is not None
assert vq.fts.query_string == "hello"
def test_search_query_pickle_with_fts(self):
"""_SearchQuery with fts should survive pickling."""
from zvec._zvec.param import _Fts, _SearchQuery
vq = _SearchQuery()
vq.topk = 10
vq.field_name = "embedding"
fts = _Fts()
fts.match_string = "test query"
vq.fts = fts
data = pickle.dumps(vq)
restored = pickle.loads(data)
assert restored.topk == 10
assert restored.field_name == "embedding"
assert restored.fts is not None
assert restored.fts.match_string == "test query"
def test_search_query_pickle_without_fts(self):
"""_SearchQuery without fts should survive pickling."""
from zvec._zvec.param import _SearchQuery
vq = _SearchQuery()
vq.topk = 5
vq.field_name = "vec"
data = pickle.dumps(vq)
restored = pickle.loads(data)
assert restored.topk == 5
assert restored.field_name == "vec"
assert restored.fts is None
python/tests/test_gil_release.py
+273
View File
@@ -0,0 +1,273 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests to verify that the GIL is released during native C++ query calls,
enabling true thread-level concurrency for multi-threaded Python applications."""
from __future__ import annotations
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
HnswIndexParam,
Query,
VectorSchema,
)
@pytest.fixture(scope="module")
def gil_test_collection(tmp_path_factory) -> Collection:
"""Create a collection with enough data to make queries take measurable time."""
schema = zvec.CollectionSchema(
name="gil_test",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
],
vectors=[
VectorSchema(
"vec",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
],
)
option = CollectionOption(read_only=False, enable_mmap=True)
temp_dir = tmp_path_factory.mktemp("zvec_gil_test")
collection_path = temp_dir / "gil_test_collection"
coll = zvec.create_and_open(path=str(collection_path), schema=schema, option=option)
# Insert enough docs to make queries non-trivial
docs = [
Doc(
id=str(i),
fields={"id": i},
vectors={"vec": [float(i % 100) + 0.1 * j for j in range(128)]},
)
for i in range(500)
]
result = coll.insert(docs)
for r in result:
assert r.ok()
yield coll
try:
coll.destroy()
except Exception:
pass
class TestGILRelease:
"""Verify that C++ query calls release the GIL, allowing true thread concurrency."""
def test_gil_released_during_query(self, gil_test_collection: Collection):
"""Prove the GIL is explicitly released during C++ Query calls.
Strategy:
- Calibrate per-query latency on the current platform (slow archs like
RISC-V can be 10x slower than x86), then dynamically pick a query count
whose total runtime fits comfortably inside switch_interval.
- Set switch_interval well above the projected total query time so that
CPython's involuntary GIL switching will NOT trigger during the run.
- A background thread (using time.sleep(0) to avoid deadlock) counts how
many times it got to run.
- Since total query time < switch_interval, the bg thread can ONLY run if
the C++ code explicitly releases the GIL.
- Reset counter just before queries; check counter > 0 after queries.
"""
query_vec = [1.0] * 128
def run_query():
gil_test_collection.query(
Query(field_name="vec", vector=query_vec),
topk=100,
)
# --- Calibrate: estimate per-query latency on this platform ---
# Warm up to avoid first-call overhead skewing the measurement.
for _ in range(3):
run_query()
calib_iters = 10
calib_start = time.monotonic()
for _ in range(calib_iters):
run_query()
per_query = max((time.monotonic() - calib_start) / calib_iters, 1e-6)
# Target total query window ~200ms, capped to a sane range so the test
# remains meaningful on both fast and slow archs.
target_total = 0.2
num_iters = max(1, min(500, int(target_total / per_query)))
projected_total = per_query * num_iters
# Pick switch_interval with a large safety margin (>=10x, >=2s) to absorb
# GC pauses, CPU throttling, and noisy-neighbor effects on CI / shared VMs.
switch_interval = max(2.0, projected_total * 10.0)
old_interval = sys.getswitchinterval()
sys.setswitchinterval(switch_interval)
try:
counter = {"value": 0}
stop_event = threading.Event()
def background_counter():
while not stop_event.is_set():
counter["value"] += 1
time.sleep(0) # Yield GIL to prevent deadlock
bg_thread = threading.Thread(target=background_counter, daemon=True)
bg_thread.start()
# Let bg thread start (sleep releases GIL)
time.sleep(0.05)
# --- Critical section: reset counter, run queries, capture counter ---
counter["value"] = 0
start = time.monotonic()
for _ in range(num_iters):
run_query()
elapsed = time.monotonic() - start
count_during_queries = counter["value"]
# --- End critical section ---
stop_event.set()
time.sleep(0.01)
bg_thread.join(timeout=5)
print(
f"\nPer-query: {per_query * 1000:.2f}ms, iters: {num_iters}, "
f"elapsed: {elapsed:.4f}s, switch_interval: {switch_interval:.2f}s"
)
print(f"Counter during queries: {count_during_queries}")
# Verify queries completed within the switch_interval window.
# If they did NOT, the run was contaminated by external jitter (GC,
# throttling, noisy neighbor) rather than a real GIL-release defect,
# so skip instead of failing to avoid flaky CI noise.
if elapsed >= switch_interval:
pytest.skip(
f"Queries took {elapsed:.3f}s >= switch_interval "
f"({switch_interval:.3f}s); calibration was outpaced by "
"runtime jitter, result is inconclusive."
)
# If elapsed < switch_interval, the ONLY way bg thread could run is
# via explicit GIL release.
assert count_during_queries > 0, (
"Background thread could not run during C++ execution despite "
"query time < switch_interval. GIL was NOT released."
)
finally:
sys.setswitchinterval(old_interval)
def test_parallel_queries_correctness(self, gil_test_collection: Collection):
"""Verify parallel queries return correct results and print timing info.
NOTE: The definitive proof of GIL release is test_gil_released_during_query
(counter + setswitchinterval). This test focuses on parallel correctness and
logs timing for manual inspection, since CI timing is too noisy for assertions.
"""
num_queries = 1000
query_vec = [1.0] * 128
def do_query():
return gil_test_collection.query(
Query(field_name="vec", vector=query_vec),
topk=100,
)
# Serial execution (baseline)
start_serial = time.monotonic()
for _ in range(num_queries):
do_query()
serial_time = time.monotonic() - start_serial
# Parallel execution
num_workers = os.cpu_count() or 2
start_parallel = time.monotonic()
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = [executor.submit(do_query) for _ in range(num_queries)]
for future in as_completed(futures):
result = future.result()
assert len(result) > 0
parallel_time = time.monotonic() - start_parallel
print(f"\nSerial time: {serial_time:.4f}s, Parallel time: {parallel_time:.4f}s")
print(
f"Speedup ratio: {serial_time / parallel_time:.2f}x (workers={num_workers})"
)
def test_thread_safety_concurrent_queries(self, gil_test_collection: Collection):
"""Verify no crashes or data corruption under concurrent query load."""
num_threads = 8
queries_per_thread = 10
errors = []
def worker(thread_id):
try:
for i in range(queries_per_thread):
vec = [float(thread_id + i) + 0.1 * j for j in range(128)]
result = gil_test_collection.query(
Query(field_name="vec", vector=vec),
topk=10,
)
assert len(result) > 0
except Exception as e:
errors.append((thread_id, e))
threads = [
threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads)
]
for t in threads:
t.start()
for t in threads:
t.join(timeout=60)
assert len(errors) == 0, f"Errors in threads: {errors}"
def test_concurrent_fetch_release_gil(self, gil_test_collection: Collection):
"""Verify Fetch operations also release the GIL correctly."""
num_threads = 4
errors = []
def worker(thread_id):
try:
ids = [str(i) for i in range(thread_id * 10, thread_id * 10 + 10)]
result = gil_test_collection.fetch(ids)
assert len(result) > 0
except Exception as e:
errors.append((thread_id, e))
threads = [
threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads)
]
for t in threads:
t.start()
for t in threads:
t.join(timeout=30)
assert len(errors) == 0, f"Errors in threads: {errors}"
python/tests/test_hnsw_contiguous_memory.py
+415
View File
@@ -0,0 +1,415 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Tests for the ``use_contiguous_memory`` optimization on ``HnswIndexParam``.
The HNSW streamer supports two allocation strategies for graph nodes:
* ``use_contiguous_memory=False`` (default): each node allocates its own
linked buffer. Lower peak memory usage, worse cache locality.
* ``use_contiguous_memory=True``: a single contiguous arena holds every
node. Higher peak memory usage, better cache locality and search
throughput.
These tests exercise the Python surface end-to-end and make sure that
when a collection is created / reopened with ``use_contiguous_memory=True``
the underlying HNSW streamer entity is constructed correctly and serves
search traffic.
"""
from __future__ import annotations
import pickle
import sys
import numpy as np
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
CollectionSchema,
Doc,
FieldSchema,
HnswIndexParam,
HnswQueryParam,
InvertIndexParam,
Query,
VectorSchema,
)
from zvec.typing import DataType, IndexType, MetricType, QuantizeType
DIMENSION = 32
NUM_DOCS = 128
TOPK = 5
# ---------------------------------------------------------------------------
def _debug_hnsw_storage_mode(coll: Collection, column: str = "dense") -> str:
"""Return the internal HNSW entity storage mode for ``column``.
Exposes the debug-only introspection hook on the pybind11 ``_Collection``.
Only meaningful after ``optimize()`` has built a persisted HNSW index; on
a pure writing segment it will raise ``KeyError``.
"""
underlying = coll._obj # type: ignore[attr-defined]
return underlying._debug_hnsw_storage_mode(column)
def _build_schema(name: str, *, use_contiguous_memory: bool) -> CollectionSchema:
"""Create a simple schema with a single FP32 HNSW vector column."""
return CollectionSchema(
name=name,
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=DIMENSION,
index_param=HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
use_contiguous_memory=use_contiguous_memory,
),
),
],
)
def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]:
"""Produce deterministic documents for insertion."""
docs: list[Doc] = []
for i in range(num):
vec = rng.standard_normal(DIMENSION).astype(np.float32)
docs.append(
Doc(
id=str(i),
fields={"id": i},
vectors={"dense": vec.tolist()},
)
)
return docs
def _assert_query_matches(coll: Collection, query_vec: list[float]) -> list[str]:
"""Run a top-k vector query and return the returned ids in order."""
vector_query = Query(
field_name="dense",
vector=query_vec,
param=HnswQueryParam(ef=128),
)
hits = coll.query(vector_query, topk=TOPK)
# Expect a single result group for the single vector query.
assert hits is not None, "query returned None"
assert len(hits) >= 1, f"expected at least one hit, got {hits!r}"
return [doc.id for doc in hits]
# ---------------------------------------------------------------------------
# 1) Pure Python surface: construction / property / to_dict / repr / pickle
# ---------------------------------------------------------------------------
class TestHnswIndexParamContiguousMemorySurface:
"""Verify the Python binding exposes ``use_contiguous_memory`` correctly."""
def test_default_is_false(self):
param = HnswIndexParam()
assert param.use_contiguous_memory is False
def test_custom_true(self):
param = HnswIndexParam(use_contiguous_memory=True)
assert param.use_contiguous_memory is True
assert param.type == IndexType.HNSW
# other fields keep their default values
assert param.m == 50
assert param.ef_construction == 500
def test_to_dict_includes_use_contiguous_memory(self):
param = HnswIndexParam(
metric_type=MetricType.L2,
m=16,
ef_construction=100,
quantize_type=QuantizeType.FP16,
use_contiguous_memory=True,
)
data = param.to_dict()
assert data["use_contiguous_memory"] is True
# Make sure existing fields are still present.
assert data["metric_type"] == "L2"
assert data["m"] == 16
assert data["ef_construction"] == 100
assert data["quantize_type"] == "FP16"
def test_repr_contains_flag(self):
on = repr(HnswIndexParam(use_contiguous_memory=True))
off = repr(HnswIndexParam(use_contiguous_memory=False))
assert "use_contiguous_memory" in on
assert "use_contiguous_memory" in off
assert "true" in on
assert "false" in off
def test_readonly_property(self):
param = HnswIndexParam(use_contiguous_memory=True)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.use_contiguous_memory = False # type: ignore[misc]
def test_pickle_roundtrip(self):
original = HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
quantize_type=QuantizeType.INT8,
use_contiguous_memory=True,
)
restored = pickle.loads(pickle.dumps(original))
assert restored.use_contiguous_memory is True
assert restored.metric_type == MetricType.COSINE
assert restored.m == 24
assert restored.ef_construction == 150
assert restored.quantize_type == QuantizeType.INT8
# ---------------------------------------------------------------------------
# 2) End-to-end: create collection, insert, query with contiguous memory on
# ---------------------------------------------------------------------------
@pytest.fixture
def rng() -> np.random.Generator:
return np.random.default_rng(seed=42)
# NOTE: the ``enable_mmap=False`` (BufferPool) variant is intentionally
# omitted from this fixture. Building a persisted HNSW index via
# ``optimize()`` / ``create_vector_index`` / ``drop_vector_index``
# currently requires mmap-backed storage, because the BufferPool backend
# has not implemented the ``create_new`` semantics yet and the guard in
# ``SegmentImpl::merge_vector_indexer`` rejects that combination. Once
# BufferPool gains write support, re-add ``False`` to ``params`` (and
# drop the guard in segment.cc) so these end-to-end tests cover both
# storage modes again.
@pytest.fixture(params=[True], ids=["mmap_on"])
def collection_option(request) -> CollectionOption:
return CollectionOption(read_only=False, enable_mmap=request.param)
# Building a new persisted HNSW index currently requires mmap-backed storage
# because the BufferPool backend has not implemented `create_new` semantics
# yet. Collections opened with ``enable_mmap=False`` therefore cannot run
# optimize()/create_vector_index/drop_vector_index. Tests use this fixture
# to know which behaviour to assert, and once BufferPool gains write support
# the guard in segment.cc (and these branches) can be removed together.
@pytest.fixture
def build_index_supported(collection_option: CollectionOption) -> bool:
return bool(collection_option.enable_mmap)
# Error message fragments emitted by the NotSupported guard in
# SegmentImpl::merge_vector_indexer / drop_vector_index. If the C++ message
# changes, update these together.
_BUILD_NOT_SUPPORTED_FRAGMENTS = ("not yet supported", "enable_mmap=false")
class TestHnswContiguousMemoryEndToEnd:
"""End-to-end: schema -> create_and_open -> insert -> query works."""
def test_create_with_contiguous_memory_and_query(
self,
tmp_path_factory,
collection_option,
rng,
):
"""With the flag on, the schema round-trips and search works end-to-end.
After ``optimize()`` the writing segment is compacted into a persisted
segment backed by the configured HNSW entity. We assert both the
user-observable behaviour (schema + search) and, via the debug hook,
that the entity type actually honours ``use_contiguous_memory``.
"""
schema = _build_schema("hnsw_contig_create", use_contiguous_memory=True)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_create"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
# Schema round-trips with the flag set.
vec_schema = coll.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is True
docs = _generate_docs(rng)
insert_result = coll.insert(docs=docs)
for r in insert_result:
assert r.ok(), f"insert failed: code={r.code()}"
assert coll.stats.doc_count == NUM_DOCS
# Build persisted HNSW index; this is where the contiguous entity
# is actually instantiated.
coll.optimize()
assert _debug_hnsw_storage_mode(coll) == "contiguous", (
"use_contiguous_memory=True should produce a contiguous entity"
)
# Pick an existing vector as the query; top-1 must be itself.
query_vec = docs[0].vector("dense")
ids = _assert_query_matches(coll, query_vec)
assert ids[0] == "0", f"expected self-recall, got top-1 id={ids[0]}"
finally:
coll.destroy()
def test_create_without_contiguous_memory_uses_mmap_entity(
self,
tmp_path_factory,
collection_option,
rng,
):
"""Baseline: when the flag is omitted the default (mmap) entity is used."""
schema = _build_schema("hnsw_contig_default", use_contiguous_memory=False)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_default"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
vec_schema = coll.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is False
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
assert coll.stats.doc_count == NUM_DOCS
coll.optimize()
# With the flag off and mmap on, the persisted entity must be the
# default mmap layout — specifically, not the contiguous arena.
assert _debug_hnsw_storage_mode(coll) == "mmap", (
"use_contiguous_memory=False + enable_mmap=True should "
"produce the mmap entity"
)
# Search still functions with the default entity backing.
query_vec = docs[0].vector("dense")
ids = _assert_query_matches(coll, query_vec)
assert ids[0] == "0"
finally:
coll.destroy()
def test_close_and_reopen_with_contiguous_memory(
self,
tmp_path_factory,
collection_option,
rng,
):
"""Reopening a collection must preserve the ``use_contiguous_memory`` flag.
The core property: the flag survives the schema persist/reload
round-trip so the HNSW streamer entity — constructed lazily on first
persisted-segment build — honours the user's choice. We run
``optimize()`` after reopen and confirm the contiguous entity was
materialized.
"""
schema = _build_schema("hnsw_contig_reopen", use_contiguous_memory=True)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_reopen"
path_str = str(path)
created = zvec.create_and_open(
path=path_str, schema=schema, option=collection_option
)
docs = _generate_docs(rng)
for r in created.insert(docs=docs):
assert r.ok()
assert created.stats.doc_count == NUM_DOCS
# Persist pending writes so that reopen reconstructs state from disk.
created.flush()
del created # close the handle
reopened = zvec.open(path=path_str, option=collection_option)
try:
assert reopened is not None
assert reopened.stats.doc_count == NUM_DOCS
# Schema persisted the flag across the reopen boundary.
vec_schema = reopened.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is True
reopened.optimize()
assert _debug_hnsw_storage_mode(reopened) == "contiguous"
# Entity actually works: exact self-recall + fetch parity.
query_vec = docs[7].vector("dense")
ids = _assert_query_matches(reopened, query_vec)
assert ids[0] == "7"
fetched = reopened.fetch([d.id for d in docs[:10]])
assert len(fetched) == 10
finally:
reopened.destroy()
def test_result_parity_with_and_without_contiguous_memory(
self,
tmp_path_factory,
rng,
):
"""
Two collections built from the same documents must return the same
top-k neighbors regardless of whether contiguous memory is enabled:
the flag is a memory-layout optimization and must not alter recall
for identical graph construction parameters on the same data.
"""
docs = _generate_docs(rng)
query_vec = docs[3].vector("dense")
def _build_and_query(tag: str, flag: bool) -> list[str]:
schema = _build_schema(f"hnsw_parity_{tag}", use_contiguous_memory=flag)
option = CollectionOption(read_only=False, enable_mmap=True)
path = tmp_path_factory.mktemp("zvec") / f"hnsw_parity_{tag}"
coll = zvec.create_and_open(path=str(path), schema=schema, option=option)
try:
for r in coll.insert(docs=docs):
assert r.ok()
coll.optimize()
expected_mode = "contiguous" if flag else "mmap"
assert _debug_hnsw_storage_mode(coll) == expected_mode, (
f"{tag}: unexpected entity type"
)
return _assert_query_matches(coll, query_vec)
finally:
coll.destroy()
ids_off = _build_and_query("off", flag=False)
ids_on = _build_and_query("on", flag=True)
# The graph is built with the same (m, ef_construction, data, order),
# so top-k results must match exactly.
assert ids_on == ids_off, (
f"top-{TOPK} results diverged between use_contiguous_memory modes: "
f"on={ids_on}, off={ids_off}"
)
# Sanity: self-recall is still perfect.
assert ids_on[0] == "3"
python/tests/test_jieba_default_dict.py
+201
View File
@@ -0,0 +1,201 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""End-to-end: jieba FTS works without any user configuration.
`import zvec` is supposed to register the wheel-bundled jieba dict
directory via `set_default_jieba_dict_dir`. With that in place a user can
declare an FTS field with `tokenizer_name="jieba"`, leave `extra_params`
empty, and Chinese full-text search just works.
Falls back to GTEST_SKIP-equivalent when running against a build that did
not bundle the dict (e.g., source-tree dev install without the install
step). In that case CI will rely on the C++ unit tests instead.
"""
from __future__ import annotations
import os
import sys
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
)
from zvec.model.param.query import Fts, Query
def _bundled_dict_dir() -> str:
"""Path zvec.__init__ would have registered; empty when not bundled."""
return zvec.get_default_jieba_dict_dir()
def _bundled_dict_files_exist() -> bool:
"""Whether the registered default actually contains the dict files.
`importlib.resources` happily returns a path even when the data dir was
not installed (e.g. source-tree dev runs); only an installed wheel has
the files on disk.
"""
import os
base = _bundled_dict_dir()
if not base:
return False
return os.path.isfile(os.path.join(base, "jieba.dict.utf8")) and os.path.isfile(
os.path.join(base, "hmm_model.utf8")
)
@pytest.fixture(scope="module", autouse=True)
def _require_bundled_dict():
if not _bundled_dict_files_exist():
pytest.skip(
"Bundled jieba dict not found at zvec/data/jieba_dict/ — "
"this test requires an installed wheel (not a source-tree dev "
"build without the install step).",
)
@pytest.fixture(scope="function")
def jieba_collection(tmp_path_factory) -> Collection:
"""FTS-only collection using jieba tokenizer and no explicit dict path."""
# env-var shadows GlobalConfig in the priority chain.
if os.environ.get("ZVEC_JIEBA_DICT_DIR"):
pytest.skip("ZVEC_JIEBA_DICT_DIR shadows the bundled default")
temp_dir = tmp_path_factory.mktemp("zvec_jieba_default")
collection_path = temp_dir / "fts_jieba"
schema = zvec.CollectionSchema(
name="fts_jieba_default",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
# Deliberately omit extra_params — the bundled default must
# be picked up via GlobalConfig.jieba_dict_dir.
index_param=FtsIndexParam(
tokenizer_name="jieba",
filters=["lowercase"],
),
),
],
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_jieba_works_without_explicit_dict_path(jieba_collection: Collection):
"""User opens collection, inserts CJK doc, searches — no init() / no
extra_params / no env var / no manual setter call. Just `import zvec`."""
docs = [
Doc(id="pk_1", fields={"title": "t1", "content": "中华人民共和国成立"}),
Doc(id="pk_2", fields={"title": "t2", "content": "无关文档"}),
]
insert_results = jieba_collection.insert(docs)
assert all(r.ok() for r in insert_results)
hits = jieba_collection.query(
queries=Query(field_name="content", fts=Fts(match_string="中华")),
topk=10,
)
ids = {doc.id for doc in hits}
assert "pk_1" in ids
assert "pk_2" not in ids
def test_default_dict_dir_is_registered_on_import():
"""Sanity check: zvec.__init__ registered a non-empty default."""
assert _bundled_dict_dir() != ""
def test_user_can_override_default_at_runtime():
"""zvec.set_default_jieba_dict_dir can be called any time to override."""
saved = zvec.get_default_jieba_dict_dir()
try:
zvec.set_default_jieba_dict_dir("/tmp/zvec/jieba-override")
assert zvec.get_default_jieba_dict_dir() == "/tmp/zvec/jieba-override"
finally:
zvec.set_default_jieba_dict_dir(saved)
@pytest.mark.skipif(
sys.platform == "win32",
reason="os.environ writes may not propagate across CRT to zvec.pyd",
)
def test_env_var_overrides_global_config(monkeypatch, tmp_path_factory):
"""ZVEC_JIEBA_DICT_DIR beats GlobalConfig in jieba's resolution chain."""
bundled = _bundled_dict_dir()
monkeypatch.setenv("ZVEC_JIEBA_DICT_DIR", bundled)
saved_global = zvec.get_default_jieba_dict_dir()
try:
zvec.set_default_jieba_dict_dir("/zvec/intentionally/missing/global")
temp_dir = tmp_path_factory.mktemp("zvec_jieba_env")
schema = zvec.CollectionSchema(
name="fts_jieba_env",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="jieba",
filters=["lowercase"],
),
),
],
)
coll = zvec.create_and_open(
path=str(temp_dir / "fts_jieba_env"),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
results = coll.insert(
[
Doc(id="pk_1", fields={"title": "t", "content": "搜索引擎技术"}),
]
)
assert all(r.ok() for r in results)
hits = coll.query(
queries=Query(field_name="content", fts=Fts(match_string="搜索")),
topk=10,
)
assert {d.id for d in hits} == {"pk_1"}
finally:
coll.destroy()
finally:
zvec.set_default_jieba_dict_dir(saved_global)
python/tests/test_params.py
+540
View File
@@ -0,0 +1,540 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import sys
import time
import numpy as np
import pytest
from zvec import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
FlatIndexParam,
HnswIndexParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
OptimizeOption,
HnswQueryParam,
IVFQueryParam,
Query,
VectorQuery,
IndexType,
MetricType,
QuantizeType,
QuantizerParam,
DataType,
VectorSchema,
)
from zvec._zvec.param import _SearchQuery
# ----------------------------
# Invert Index Param Test Case
# ----------------------------
class TestInvertIndexParam:
def test_default(self):
param = InvertIndexParam()
assert param.enable_range_optimization is False
assert param.enable_extended_wildcard is False
assert param.type == IndexType.INVERT
def test_custom(self):
param = InvertIndexParam(
enable_range_optimization=True, enable_extended_wildcard=True
)
assert param.enable_range_optimization is True
assert param.enable_extended_wildcard is True
def test_readonly(self):
param = InvertIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.enable_range_optimization = False
param.enable_extended_wildcard = False
# ----------------------------
# Hnsw Index Param Test Case
# ----------------------------
class TestHnswIndexParam:
def test_default(self):
param = HnswIndexParam()
assert param.metric_type == MetricType.IP
assert param.m == 50
assert param.ef_construction == 500
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.type == IndexType.HNSW
def test_custom(self):
param = HnswIndexParam(
metric_type=MetricType.L2,
m=10,
ef_construction=1000,
quantize_type=QuantizeType.FP16,
)
assert param.metric_type == MetricType.L2
assert param.m == 10
assert param.ef_construction == 1000
assert param.quantize_type == QuantizeType.FP16
@pytest.mark.parametrize(
"attr", ["metric_type", "m", "ef_construction", "quantize_type"]
)
def test_readonly_attributes(self, attr):
param = HnswIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# Flat Index Param Test Case
# ----------------------------
class TestFlatIndexParam:
def test_default(self):
param = FlatIndexParam()
assert param.type == IndexType.FLAT
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.metric_type == MetricType.IP
def test_custom(self):
param = FlatIndexParam(
metric_type=MetricType.L2, quantize_type=QuantizeType.INT8
)
assert param.metric_type == MetricType.L2
assert param.quantize_type == QuantizeType.INT8
@pytest.mark.parametrize("attr", ["metric_type", "quantize_type"])
def test_readonly_attributes(self, attr):
param = FlatIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# Ivf Index Param Test Case
# ----------------------------
class TestIVFIndexParam:
def test_default(self):
param = IVFIndexParam()
assert param.metric_type == MetricType.IP
assert param.n_list == 10
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.type == IndexType.IVF
def test_custom(self):
param = IVFIndexParam(
metric_type=MetricType.L2, n_list=1000, quantize_type=QuantizeType.FP16
)
assert param.metric_type == MetricType.L2
assert param.n_list == 1000
assert param.quantize_type == QuantizeType.FP16
assert param.type == IndexType.IVF
@pytest.mark.parametrize("attr", ["metric_type", "n_list", "quantize_type"])
def test_readonly_attributes(self, attr):
param = IVFIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# CollectionOption Test Case
# ----------------------------
class TestCollectionOption:
def test_default(self):
option = CollectionOption()
assert option is not None
assert option.read_only == False
assert option.enable_mmap == True
def test_custom(self):
option = CollectionOption(read_only=True, enable_mmap=False)
assert option.read_only == True
assert option.enable_mmap == False
option = CollectionOption(read_only=False, enable_mmap=True)
assert option.read_only == False
assert option.enable_mmap == True
@pytest.mark.parametrize("attr", ["read_only", "enable_mmap"])
def test_readonly_attributes(self, attr):
param = CollectionOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# IndexOption Test Case
# ----------------------------
class TestIndexOption:
def test_default(self):
option = IndexOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = IndexOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = IndexOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# AddColumnOption Test Case
# ----------------------------
class TestAddColumnOption:
def test_default(self):
option = AddColumnOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = AddColumnOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = AddColumnOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# AlterColumnOption Test Case
# ----------------------------
class TestAlterColumnOption:
def test_default(self):
option = AlterColumnOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = AlterColumnOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = AlterColumnOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# OptimizeOption Test Case
# ----------------------------
class TestOptimizeOption:
def test_default(self):
option = OptimizeOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = OptimizeOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = OptimizeOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# HnswQueryParam Test Case
# ----------------------------
class TestHnswQueryParam:
def test_default(self):
param = HnswQueryParam()
assert param is not None
assert param.ef == 300
assert param.is_using_refiner == False
assert param.radius == 0
assert param.is_linear == False
assert param.prefetch_offset == 8
assert param.prefetch_lines == 0
def test_custom(self):
param = HnswQueryParam(
ef=10,
is_using_refiner=True,
radius=30,
is_linear=True,
extra_params={
"prefetch_offset": 16,
"prefetch_lines": 4,
},
)
assert param.ef == 10
assert param.is_using_refiner == True
assert param.radius == 30
assert param.is_linear == True
assert param.prefetch_offset == 16
assert param.prefetch_lines == 4
def test_readonly_attributes(self):
param = HnswQueryParam()
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.ef = 10
param.is_using_refiner = True
param.radius = 30
param.is_linear = True
# # ----------------------------
# # IVFQueryParam Test Case
# # ----------------------------
# class TestIVFQueryParam:
# def test_default(self):
# param = IVFQueryParam()
# assert param is not None
# assert param.nprobe == 10
# assert param.is_using_refiner == False
# assert param.radius == 0
# assert param.is_linear == False
# assert param.scale_factor == 10
#
# def test_custom(self):
# param = IVFQueryParam(
# nprobe=20,
# is_using_refiner=True,
# radius=30,
# is_linear=True,
# scale_factor=40
# )
# assert param.nprobe == 20
# assert param.is_using_refiner == True
# assert param.radius == 30
# assert param.is_linear == True
# assert param.scale_factor == 40
class TestQuery:
def test_init_with_valid_id(self):
vq = Query(field_name="embedding", id="doc123")
assert vq.field_name == "embedding"
assert vq.id == "doc123"
assert vq.vector is None
assert vq.param is None
def test_init_with_valid_vector(self):
vec = [0.1, 0.2, 0.3]
param = HnswQueryParam(ef=300)
vq = Query(field_name="embedding", vector=vec, param=param)
assert vq.field_name == "embedding"
assert vq.vector == vec
assert vq.param == param
def test_init_both_id_and_vector_raises_error(self):
with pytest.raises(ValueError):
Query(field_name="embedding", id="doc123", vector=[0.1])._validate()
def test_init_without_field_name_raises_error(self):
with pytest.raises(ValueError):
Query(field_name=None)._validate()
def test_has_id_returns_true_when_id_set(self):
vq = Query(field_name="embedding", id="doc123")
assert vq.has_id()
def test_has_id_returns_false_when_no_id(self):
vq = Query(field_name="embedding", vector=[0.1])
assert not vq.has_id()
def test_has_vector_returns_true_with_non_empty_vector(self):
vq = Query(field_name="embedding", vector=[0.1])
assert vq.has_vector()
def test_validate_fails_on_both_id_and_vector(self):
vq = Query(field_name="test", id="doc123", vector=[0.1])
with pytest.raises(ValueError):
vq._validate()
def test_validate_fails_on_both_id_and_numpy_vector(self):
vq = Query(field_name="test", id="doc123", vector=np.array([0.1]))
with pytest.raises(ValueError, match="Cannot provide both id and vector"):
vq._validate()
class TestVectorQueryDeprecated:
def test_deprecation_warning(self):
import warnings
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
vq = VectorQuery(field_name="embedding", id="doc123")
assert len(w) == 1
assert issubclass(w[0].category, DeprecationWarning)
assert "Query" in str(w[0].message)
def test_isinstance_compatibility(self):
import warnings
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
vq = VectorQuery(field_name="embedding", id="doc123")
assert isinstance(vq, Query)
# ----------------------------
# QuantizerParam Test Case
# ----------------------------
class TestQuantizerParam:
def test_default(self):
qp = QuantizerParam()
assert qp.enable_rotate is False
def test_enable_rotate_true(self):
qp = QuantizerParam(enable_rotate=True)
assert qp.enable_rotate is True
def test_enable_rotate_false(self):
qp = QuantizerParam(enable_rotate=False)
assert qp.enable_rotate is False
def test_equality(self):
qp1 = QuantizerParam(enable_rotate=True)
qp2 = QuantizerParam(enable_rotate=True)
qp3 = QuantizerParam(enable_rotate=False)
assert qp1 == qp2
assert qp1 != qp3
def test_to_dict(self):
qp = QuantizerParam(enable_rotate=True)
d = qp.to_dict()
assert isinstance(d, dict)
assert d.get("enable_rotate") is True
def test_repr(self):
qp = QuantizerParam(enable_rotate=True)
r = repr(qp)
assert "enable_rotate" in r or "QuantizerParam" in r
def test_pickle_roundtrip(self):
import pickle
qp = QuantizerParam(enable_rotate=True)
data = pickle.dumps(qp)
qp2 = pickle.loads(data)
assert qp2.enable_rotate is True
assert qp == qp2
# ----------------------------
# HnswIndexParam with QuantizerParam
# ----------------------------
class TestHnswIndexParamQuantizer:
def test_default_quantizer_param(self):
param = HnswIndexParam()
assert param.quantizer_param is not None
assert param.quantizer_param.enable_rotate is False
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = HnswIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8
# ----------------------------
# FlatIndexParam with QuantizerParam
# ----------------------------
class TestFlatIndexParamQuantizer:
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = FlatIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8
python/tests/test_query_executor.py
+329
View File
@@ -0,0 +1,329 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Dict, Union
from unittest.mock import MagicMock, patch
import numpy as np
import math
from zvec._zvec.param import _SearchQuery
import pytest
from zvec.executor.query_executor import (
QueryContext,
QueryExecutor,
)
from zvec import (
RrfReRanker,
WeightedReRanker,
HnswQueryParam,
CollectionSchema,
VectorSchema,
DataType,
MetricType,
Query,
VectorQuery,
)
from zvec.extension.multi_vector_reranker import CallbackReRanker
# ----------------------------
# Mock Collection Schema
# ----------------------------
class MockCollectionSchema(CollectionSchema):
def __init__(self, vectors=Union[VectorSchema, Dict[str, VectorSchema]]):
self._vectors = (
[vectors] if not isinstance(vectors, Dict) else list(vectors.values())
)
@property
def vectors(self):
return self._vectors
# ----------------------------
# VectorQuery Test Case
# ----------------------------
class TestQuery:
def test_init(self):
query = Query(field_name="test_field")
assert query.field_name == "test_field"
assert query.id is None
assert query.vector is None
assert query.param is None
param = HnswQueryParam()
query = Query(
field_name="test_field", id="test_id", vector=[1, 2, 3], param=param
)
assert query.field_name == "test_field"
assert query.id == "test_id"
assert query.vector == [1, 2, 3]
assert query.param == param
def test_has_id(self):
query = Query(field_name="test_field")
assert not query.has_id()
query = Query(field_name="test_field", id="test_id")
assert query.has_id()
def test_has_vector(self):
query = Query(field_name="test_field")
assert not query.has_vector()
query = Query(field_name="test_field", vector=[])
assert not query.has_vector()
query = Query(field_name="test_field", vector=[1, 2, 3])
assert query.has_vector()
def test_validate_dense_fp16_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP16)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float16)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_fp32_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float32)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_fp64_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP64)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float64)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_int8_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_INT8)
vec = np.array([1, 2, 3], dtype=np.int8)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_sparse_fp32_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP32)
vec = {1: 1.1, 2: 2.2, 3: 3.3}
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
for k in vec.keys():
assert math.isclose(vec[k], ret[k], abs_tol=1e-6)
def test_validate_sparse_fp16_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP16)
vec = {1: 1.1, 2: 2.2, 3: 3.3}
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
for k in vec.keys():
assert math.isclose(np.float16(vec[k]), ret[k], abs_tol=1e-6)
class TestVectorQueryDeprecated:
def test_deprecation_warning(self):
import warnings
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
vq = VectorQuery(field_name="test_field")
assert len(w) == 1
assert issubclass(w[0].category, DeprecationWarning)
assert "Query" in str(w[0].message)
def test_isinstance_compatibility(self):
import warnings
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
vq = VectorQuery(field_name="test_field")
assert isinstance(vq, Query)
class TestQueryContext:
def test_init(self):
ctx = QueryContext(topk=10)
assert ctx.topk == 10
assert ctx.queries == []
assert ctx.filter is None
assert ctx.reranker is None
assert ctx.output_fields is None
assert ctx.include_vector is False
def test_properties(self):
queries = [Query(field_name="test")]
reranker = RrfReRanker()
output_fields = ["field1", "field2"]
ctx = QueryContext(
topk=5,
filter="test_filter",
include_vector=True,
queries=queries,
output_fields=output_fields,
reranker=reranker,
)
assert ctx.topk == 5
assert ctx.queries == queries
assert ctx.filter == "test_filter"
assert ctx.reranker == reranker
assert ctx.output_fields == output_fields
assert ctx.include_vector is True
def test_properties_with_weighted_reranker(self):
queries = [Query(field_name="test")]
reranker = WeightedReRanker(
weights=[1.0],
)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
assert ctx.reranker.weights == [1.0]
def test_properties_with_callback_reranker(self):
queries = [Query(field_name="test")]
cb = lambda query_results, topn: []
reranker = CallbackReRanker(callback=cb)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
class TestQueryExecutor:
def test_init(self):
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
assert isinstance(executor, QueryExecutor)
def test_do_build_without_queries(self):
# When no queries are given, build a single vector-less query.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5, filter="test_filter")
result = executor._build_queries(ctx, MagicMock())
assert len(result) == 1
assert result[0].topk == 5
assert result[0].filter == "test_filter"
def test_do_build_query_wo_vector(self):
# Vector-less core query should carry the context query params.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=7, filter="f", include_vector=True)
core_vector = executor._build_base_search_query(ctx)
assert core_vector.topk == 7
assert core_vector.filter == "f"
assert core_vector.include_vector is True
def test_do_merge_rerank_results_single_without_reranker(self):
# A single result list without a reranker is returned as-is.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
docs_list = [["doc1", "doc2"]]
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["doc1", "doc2"]
def test_do_merge_rerank_results_empty(self):
# Empty results should raise an error.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
with pytest.raises(ValueError, match="Query results is empty"):
executor._merge_and_rerank(ctx, [])
def test_do_merge_rerank_results_with_reranker(self):
# Multiple result lists are merged through the reranker.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
reranker = MagicMock()
reranker.rerank.return_value = ["merged"]
ctx = QueryContext(
topk=5,
queries=[Query(field_name="test1"), Query(field_name="test2")],
reranker=reranker,
)
docs_list = [["d1"], ["d2"]]
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["merged"]
reranker.rerank.assert_called_once_with(docs_list, ctx.topk)
def test_execute_python_pipeline(self):
# Each query is executed serially and converted into a result list.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
collection = MagicMock()
collection.Query.side_effect = [["raw1"], ["raw2"]]
vectors = [MagicMock(), MagicMock()]
with patch(
"zvec.executor.query_executor.convert_to_py_doc",
side_effect=lambda doc, schema: doc,
):
results = executor._execute_python_pipeline(vectors, collection)
assert results == [["raw1"], ["raw2"]]
assert collection.Query.call_count == 2
def test_build_search_query_by_missing_id_raises_value_error(self):
vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
schema = CollectionSchema(name="test_collection", vectors=[vector_schema])
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
collection = MagicMock()
collection.Fetch.return_value = {}
with pytest.raises(ValueError, match="Document with id 'missing' not found"):
executor._build_search_query(
ctx, Query(field_name="test", id="missing"), collection
)
def test_build_search_query_validates_query(self):
vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
schema = CollectionSchema(name="test_collection", vectors=[vector_schema])
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
collection = MagicMock()
with pytest.raises(ValueError, match="Cannot provide both id and vector"):
executor._build_search_query(
ctx,
Query(field_name="test", id="doc1", vector=np.array([0.1])),
collection,
)
python/tests/test_reranker.py
+948
View File
@@ -0,0 +1,948 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from unittest.mock import patch, MagicMock
import pytest
import os
from zvec import Doc, MetricType, VectorSchema, DataType, FlatIndexParam
from zvec.extension.multi_vector_reranker import (
CallbackReRanker,
RrfReRanker,
WeightedReRanker,
)
from zvec.extension.sentence_transformer_rerank_function import (
DefaultLocalReRanker,
)
from zvec.extension.qwen_rerank_function import QwenReRanker
# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests
RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1"
# ----------------------------
# RrfReRanker Test Case
# ----------------------------
class TestRrfReRanker:
def test_init(self):
reranker = RrfReRanker(rank_constant=100)
assert reranker.rank_constant == 100
def test_default_rank_constant(self):
reranker = RrfReRanker()
assert reranker.rank_constant == 60
def test_rerank(self):
reranker = RrfReRanker(rank_constant=60)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
doc3 = Doc(id="3", score=0.9)
doc4 = Doc(id="4", score=0.6)
query_results = [[doc1, doc2, doc3], [doc3, doc1, doc4]]
results = reranker.rerank(query_results, topn=3)
assert len(results) <= 3
for doc in results:
assert hasattr(doc, "score")
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True)
# ----------------------------
# WeightedReRanker Test Case
# ----------------------------
class TestWeightedReRanker:
@staticmethod
def _make_fields(metrics):
return [
VectorSchema(
name=f"vector{i}",
data_type=DataType.VECTOR_FP32,
dimension=4,
index_param=FlatIndexParam(metric_type=metric),
)
for i, metric in enumerate(metrics)
]
def test_init(self):
reranker = WeightedReRanker([0.7, 0.3])
assert reranker.weights == [0.7, 0.3]
def test_rerank(self):
reranker = WeightedReRanker([0.7, 0.3])
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
doc3 = Doc(id="3", score=0.9)
query_results = [[doc1, doc2], [doc2, doc3]]
fields = self._make_fields([MetricType.L2, MetricType.L2])
results = reranker.rerank(query_results, topn=3, fields=fields)
assert len(results) <= 3
for doc in results:
assert hasattr(doc, "score")
# ----------------------------
# CallbackReRanker Test Case
# ----------------------------
class TestCallbackReRanker:
def test_rerank(self):
def my_callback(query_results, fields, topn):
all_docs = []
for docs in query_results:
all_docs.extend(docs)
all_docs.sort(key=lambda d: d.score, reverse=True)
return all_docs[:topn]
reranker = CallbackReRanker(my_callback)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.9)
doc3 = Doc(id="3", score=0.7)
doc4 = Doc(id="4", score=0.6)
query_results = [[doc1, doc2], [doc3, doc4]]
results = reranker.rerank(query_results, topn=3)
assert len(results) == 3
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True)
def test_callback_with_topn(self):
received_topn = []
def my_callback(query_results, fields, topn):
received_topn.append(topn)
return []
reranker = CallbackReRanker(my_callback)
reranker.rerank([[Doc(id="1", score=0.5)]], topn=7)
assert received_topn == [7]
# ----------------------------
# QwenReRanker Test Case
# ----------------------------
class TestQwenReRanker:
def test_init_without_query(self):
with pytest.raises(ValueError, match="Query is required for QwenReRanker"):
QwenReRanker(api_key="test_key")
def test_init_without_api_key(self):
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="DashScope API key is required"):
QwenReRanker(query="test")
@patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"})
def test_init_with_env_api_key(self):
reranker = QwenReRanker(query="test", rerank_field="content")
assert reranker.query == "test"
assert reranker._api_key == "test_key"
assert reranker.rerank_field == "content"
def test_init_with_explicit_api_key(self):
reranker = QwenReRanker(
query="test", api_key="explicit_key", rerank_field="content"
)
assert reranker.query == "test"
assert reranker._api_key == "explicit_key"
def test_model_property(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
assert reranker.model == "gte-rerank-v2"
reranker = QwenReRanker(
query="test",
model="custom-model",
api_key="test_key",
rerank_field="content",
)
assert reranker.model == "custom-model"
def test_query_property(self):
reranker = QwenReRanker(
query="test query", api_key="test_key", rerank_field="content"
)
assert reranker.query == "test query"
def test_rerank_field_property(self):
reranker = QwenReRanker(query="test", api_key="test_key", rerank_field="title")
assert reranker.rerank_field == "title"
def test_rerank_empty_results(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_success(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.95},
{"index": 1, "relevance_score": 0.85},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test query", api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Document 1"}),
Doc(id="2", fields={"content": "Document 2"}),
]
}
results = reranker.rerank(query_results, topn=2)
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
# Verify API call
mock_dashscope.TextReRank.call.assert_called_once_with(
model="gte-rerank-v2",
query="test query",
documents=["Document 1", "Document 2"],
top_n=2,
return_documents=False,
)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_deduplicate_documents(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.9},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
query_results = {"vector1": [doc1], "vector2": [doc1]}
results = reranker.rerank(query_results, topn=5)
# Should only call API with document once
call_args = mock_dashscope.TextReRank.call.call_args
assert len(call_args[1]["documents"]) == 1
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_api_error(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API error response
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.message = "Invalid request"
mock_response.code = "InvalidParameter"
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(ValueError, match="DashScope API error"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_runtime_error(self, mock_require_module):
# Mock dashscope module that raises exception
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
mock_dashscope.TextReRank.call.side_effect = Exception("Network error")
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to call DashScope API"):
reranker.rerank(query_results)
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_qwen_rerank(self):
"""Integration test with real DashScope TextReRank API.
To run this test, set environment variables:
export ZVEC_RUN_INTEGRATION_TESTS=1
export DASHSCOPE_API_KEY=your-api-key
"""
# Create reranker with real API
reranker = QwenReRanker(
query="What is machine learning?",
rerank_field="content",
model="gte-rerank-v2",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real API
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
assert doc.score > 0, "Score should be positive"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nReranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")
# ----------------------------
# DefaultLocalReRanker Test Case
# ----------------------------
class TestDefaultLocalReRanker:
"""Test cases for DefaultLocalReRanker."""
def test_init_without_query(self):
"""Test initialization fails without query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(rerank_field="content")
def test_init_with_empty_query(self):
"""Test initialization fails with empty query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(query="", rerank_field="content")
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_success(self, mock_require_module):
"""Test successful initialization with mocked model."""
# Mock sentence_transformers module
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock() # Cross-encoder has predict method
mock_model.device = "cpu"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="test query",
rerank_field="content",
model_name="cross-encoder/ms-marco-MiniLM-L6-v2",
)
assert reranker.query == "test query"
assert reranker.rerank_field == "content"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L6-v2"
assert reranker.model_source == "huggingface"
assert reranker.batch_size == 32
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_with_custom_params(self, mock_require_module):
"""Test initialization with custom parameters."""
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_model.device = "cuda"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="custom query",
rerank_field="title",
model_name="cross-encoder/ms-marco-MiniLM-L12-v2",
model_source="modelscope",
device="cuda",
batch_size=64,
)
assert reranker.query == "custom query"
assert reranker.rerank_field == "title"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L12-v2"
assert reranker.model_source == "modelscope"
assert reranker.batch_size == 64
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_invalid_model(self, mock_require_module):
"""Test initialization fails with non-cross-encoder model."""
# Mock a model without predict method (not a cross-encoder)
mock_st = MagicMock()
mock_model = MagicMock(spec=[]) # No predict method
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
with pytest.raises(ValueError, match="does not appear to be a cross-encoder"):
DefaultLocalReRanker(query="test", rerank_field="content")
def test_query_property(self):
"""Test query property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
assert reranker.query == "test query"
def test_rerank_field_property(self):
"""Test rerank_field property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="title")
assert reranker.rerank_field == "title"
def test_batch_size_property(self):
"""Test batch_size property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=128
)
assert reranker.batch_size == 128
def test_rerank_empty_results(self):
"""Test rerank with empty query_results."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
"""Test rerank with documents missing rerank_field."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
"""Test rerank skips documents with empty content."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_success(self):
"""Test successful rerank with mocked model."""
# Mock standard cross-encoder model
mock_model = MagicMock()
# Mock predict method to return scores
import numpy as np
mock_scores = np.array([0.95, 0.85, 0.75])
mock_model.predict.return_value = mock_scores
mock_model.device = "cpu"
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", score=0.8, fields={"content": "Document 1"}),
Doc(id="2", score=0.7, fields={"content": "Document 2"}),
Doc(id="3", score=0.6, fields={"content": "Document 3"}),
]
}
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) == 3
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
assert results[2].id == "3"
assert results[2].score == 0.75
# Verify model.predict was called correctly
assert mock_model.predict.called
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 3
assert pairs[0] == ["test query", "Document 1"]
assert pairs[1] == ["test query", "Document 2"]
assert pairs[2] == ["test query", "Document 3"]
assert call_args[1]["batch_size"] == 32
assert call_args[1]["show_progress_bar"] is False
def test_rerank_with_topn_limit(self):
"""Test rerank respects topn limit."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
Doc(id="4", fields={"content": "Doc 4"}),
Doc(id="5", fields={"content": "Doc 5"}),
]
}
results = reranker.rerank(query_results, topn=2)
# Should only return top 2
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.9
assert results[1].id == "2"
assert results[1].score == 0.8
def test_rerank_deduplicate_documents(self):
"""Test rerank deduplicates documents across multiple vectors."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.95, 0.85])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
doc2 = Doc(id="2", fields={"content": "Document 2"})
query_results = {
"vector1": [doc1, doc2],
"vector2": [doc1], # doc1 appears in both
}
results = reranker.rerank(query_results, topn=5)
# Should only process each document once
assert len(results) == 2
assert mock_model.predict.call_count == 1
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 2 # Only 2 unique documents
def test_rerank_sorting(self):
"""Test rerank sorts documents by score in descending order."""
mock_model = MagicMock()
import numpy as np
# Return scores in non-sorted order
mock_scores = np.array([0.6, 0.9, 0.7])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
]
}
results = reranker.rerank(query_results, topn=3)
# Should be sorted by score (descending)
assert len(results) == 3
assert results[0].id == "2" # score 0.9
assert results[0].score == 0.9
assert results[1].id == "3" # score 0.7
assert results[1].score == 0.7
assert results[2].id == "1" # score 0.6
assert results[2].score == 0.6
def test_rerank_model_error(self):
"""Test rerank handles model prediction errors."""
mock_model = MagicMock()
# Mock predict to raise exception
mock_model.predict.side_effect = Exception("Model inference error")
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to compute rerank scores"):
reranker.rerank(query_results)
def test_rerank_with_custom_batch_size(self):
"""Test rerank uses custom batch_size."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=64
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
]
}
reranker.rerank(query_results)
# Verify batch_size is passed to predict
call_args = mock_model.predict.call_args
assert call_args[1]["batch_size"] == 64
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_sentence_transformer_rerank(self):
"""Integration test with real SentenceTransformer cross-encoder model.
To run this test, set environment variable:
export ZVEC_RUN_INTEGRATION_TESTS=1
Note: This test requires sentence-transformers package and will
download the MS MARCO MiniLM model (~80MB) on first run.
"""
# Create reranker with real model (using default lightweight model)
reranker = DefaultLocalReRanker(
query="What is machine learning?",
rerank_field="content",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real model
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Documents 1, 3, and 4 are about machine learning, should rank higher
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nSentenceTransformer Reranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")
python/tests/test_schema.py
+246
View File
@@ -0,0 +1,246 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import pytest
from zvec import (
CollectionSchema,
CollectionStats,
FieldSchema,
VectorSchema,
HnswIndexParam,
InvertIndexParam,
DataType,
IndexType,
MetricType,
)
# ----------------------------
# FieldSchema Test Case
# ----------------------------
class TestFieldSchema:
def test_default(self):
field = FieldSchema("field", data_type=DataType.FLOAT)
assert field.name == "field"
assert field.data_type == DataType.FLOAT
assert field.nullable is False
assert field.index_param is None
def test_custom(self):
field_1 = FieldSchema(
name="float",
data_type=DataType.FLOAT,
nullable=True,
index_param=InvertIndexParam(),
)
assert field_1.name == "float"
assert field_1.data_type == DataType.FLOAT
assert field_1.nullable is True
assert field_1.index_param.enable_range_optimization is False
field_2 = FieldSchema(
name="str",
data_type=DataType.STRING,
nullable=True,
index_param=InvertIndexParam(enable_range_optimization=True),
)
assert field_2.name == "str"
assert field_2.data_type == DataType.STRING
assert field_2.nullable is True
assert field_2.index_param.enable_range_optimization is True
def test_readonly(self):
field = FieldSchema(
name="float",
data_type=DataType.FLOAT,
nullable=True,
index_param=InvertIndexParam(),
)
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
field.index_param = InvertIndexParam(enable_range_optimization=True)
# ----------------------------
# VectorSchema Test Case
# ----------------------------
class TestVectorSchema:
def test_default(self):
field = VectorSchema("vector", data_type=DataType.VECTOR_FP32, dimension=128)
assert field.name == "vector"
assert field.data_type == DataType.VECTOR_FP32
assert field.dimension == 128
assert field.index_param is not None
assert field.index_param.type == IndexType.FLAT
assert field.index_param.metric_type == MetricType.IP
def test_custom(self):
field = VectorSchema(
name="vector",
data_type=DataType.VECTOR_INT8,
dimension=512,
index_param=HnswIndexParam(
metric_type=MetricType.COSINE, m=15, ef_construction=300
),
)
assert field.name == "vector"
assert field.data_type == DataType.VECTOR_INT8
assert field.index_param.metric_type == MetricType.COSINE
assert field.index_param.m == 15
assert field.index_param.ef_construction == 300
def test_readonly(self):
field = VectorSchema(
name="vector",
dimension=128,
data_type=DataType.VECTOR_INT8,
)
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
field.dimension = 4
# ----------------------------
# CollectionSchema Test Case
# ----------------------------
class TestCollectionSchema:
def test_collection_schema_with_single_field(self):
collection_schema = CollectionSchema(
name="test_collection",
fields=FieldSchema(
name="id",
data_type=DataType.INT64,
index_param=InvertIndexParam(),
nullable=False,
),
vectors=VectorSchema(
name="vector",
data_type=DataType.VECTOR_INT8,
dimension=128,
index_param=HnswIndexParam(),
),
)
assert collection_schema is not None
assert collection_schema.name == "test_collection"
assert len(collection_schema.fields) == 1
assert len(collection_schema.vectors) == 1
field = collection_schema.field("id")
assert field is not None
assert field.name == "id"
assert field.data_type == DataType.INT64
assert not field.nullable
assert field.index_param.type == IndexType.INVERT
assert not field.index_param.enable_range_optimization
vector = collection_schema.vector("vector")
assert vector is not None
assert vector.name == "vector"
assert vector.data_type == DataType.VECTOR_INT8
assert vector.dimension == 128
assert vector.index_param.type == IndexType.HNSW
assert vector.index_param.m == 50
assert vector.index_param.ef_construction == 500
assert vector.index_param.metric_type == MetricType.IP
def test_collection_schema_with_multi_fields(self):
collection_schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
FieldSchema(
"name",
DataType.STRING,
nullable=False,
index_param=InvertIndexParam(),
),
FieldSchema(
"weight",
DataType.INT32,
nullable=True,
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
assert collection_schema is not None
assert collection_schema.name == "test_collection"
assert len(collection_schema.fields) == 3
assert len(collection_schema.vectors) == 2
field_id = collection_schema.field("id")
assert field_id is not None
assert field_id.name == "id"
assert field_id.data_type == DataType.INT64
assert not field_id.nullable
assert field_id.index_param.type == IndexType.INVERT
dense = collection_schema.vector("dense")
assert dense is not None
assert dense.name == "dense"
assert dense.data_type == DataType.VECTOR_FP32
assert dense.dimension == 128
assert dense.index_param.type == IndexType.HNSW
sparse = collection_schema.vector("sparse")
assert sparse is not None
assert sparse.name == "sparse"
assert sparse.data_type == DataType.SPARSE_VECTOR_FP32
assert sparse.dimension == 0
assert sparse.index_param.type == IndexType.HNSW
assert str(collection_schema) is not None
# ----------------------------
# CollectionStats Test Case
# ----------------------------
class TestCollectionStats:
"""
The constructor of CollectionStats is not provided.
It can only be obtained through collection.stats()
"""
def test_collection_stats(self):
stats = CollectionStats()
assert stats is not None
python/tests/test_typing.py
+141
View File
@@ -0,0 +1,141 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import pytest
from zvec import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
# ----------------------------
# Enum Test Case
# ----------------------------
@pytest.mark.parametrize(
"member, name",
[
(DataType.FLOAT, "FLOAT"),
(IndexType.HNSW, "HNSW"),
(MetricType.COSINE, "COSINE"),
(QuantizeType.INT8, "INT8"),
(StatusCode.OK, "OK"),
],
)
def test_enum_names(member, name):
assert member.name == name
@pytest.mark.parametrize(
"member, value",
[
(DataType.FLOAT, 8),
(IndexType.HNSW, 1),
(MetricType.COSINE, 3),
(QuantizeType.INT8, 2),
(StatusCode.OK, 0),
],
)
def test_enum_values(member, value):
assert member.value == value
@pytest.mark.parametrize("member", ["L2", "IP", "COSINE"])
def test_metric_type_has_member(member):
assert member in MetricType.__members__
@pytest.mark.parametrize(
"member",
[
"STRING",
"BOOL",
"INT32",
"INT64",
"FLOAT",
"DOUBLE",
"UINT32",
"UINT64",
"VECTOR_FP16",
"VECTOR_FP32",
"VECTOR_FP64",
"VECTOR_INT8",
"SPARSE_VECTOR_FP32",
"SPARSE_VECTOR_FP16",
"ARRAY_STRING",
"ARRAY_INT32",
"ARRAY_INT64",
"ARRAY_FLOAT",
"ARRAY_DOUBLE",
"ARRAY_BOOL",
"ARRAY_UINT32",
"ARRAY_UINT64",
],
)
def test_data_type_has_member(member):
assert member in DataType.__members__
@pytest.mark.parametrize("member", ["HNSW", "IVF", "FLAT", "INVERT"])
def test_index_type_has_member(member):
assert member in IndexType.__members__
@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"])
def test_quantize_type_has_member(member):
assert member in QuantizeType.__members__
@pytest.mark.parametrize(
"member",
[
"OK",
"UNKNOWN",
"NOT_FOUND",
"ALREADY_EXISTS",
"INVALID_ARGUMENT",
"PERMISSION_DENIED",
"FAILED_PRECONDITION",
"RESOURCE_EXHAUSTED",
"UNAVAILABLE",
"INTERNAL_ERROR",
"NOT_SUPPORTED",
],
)
def test_status_code_has_member(member):
assert member in StatusCode.__members__
# ----------------------------
# Status Test Case
# ----------------------------
class TestStatus:
def test_status_code(self):
status = Status(StatusCode.OK)
assert status.code() == StatusCode.OK
def test_status_message(self):
status = Status(StatusCode.OK, "OK")
assert status.message() == "OK"
status = Status(StatusCode.NOT_FOUND, "Not Found")
assert status.message() == "Not Found"
def test_status_ok(self):
status = Status(StatusCode.OK)
assert status.ok()
python/tests/test_util.py
+89
View File
@@ -0,0 +1,89 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from zvec import require_module
# ----------------------------
# require_module func Test Case
# ----------------------------
def test_require_module_success():
module = require_module("os")
assert module is not None
assert hasattr(module, "path")
def test_require_module_with_submodule_success():
module = require_module("os.path")
assert module is not None
assert hasattr(module, "join")
def test_require_module_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("nonexistent_module")
exception_msg = str(exc_info.value)
assert "Required package 'nonexistent_module' is not installed." in exception_msg
def test_require_module_with_mitigation_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("nonexistent_module.submodule", mitigation="custom_package")
exception_msg = str(exc_info.value)
assert "Required package 'custom_package' is not installed." in exception_msg
assert (
"Module 'nonexistent_module.submodule' is part of 'nonexistent_module'"
in exception_msg
)
assert "please pip install 'custom_package'." in exception_msg
def test_require_module_submodule_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("os.nonexistent_submodule")
exception_msg = str(exc_info.value)
assert (
"Required package 'os.nonexistent_submodule' is not installed." in exception_msg
)
assert "Module 'os.nonexistent_submodule' is part of 'os'" in exception_msg
assert "please pip install 'os'." in exception_msg
@patch("importlib.import_module")
def test_require_module_wraps_original_exception(mock_import_module):
original_exception = ImportError("Original error")
mock_import_module.side_effect = original_exception
with pytest.raises(ImportError) as exc_info:
require_module("some_module")
assert exc_info.value.__cause__ is original_exception
@patch("importlib.import_module")
def test_require_module_calls_importlib(mock_import_module):
mock_module = MagicMock()
mock_import_module.return_value = mock_module
result = require_module("test_module")
mock_import_module.assert_called_once_with("test_module")
assert result is mock_module
python/tests/test_vamana.py
+501
View File
@@ -0,0 +1,501 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Tests for the Python entry point of the Vamana (DiskANN) dense vector index.
Mirrors the structure of ``test_hnsw_contiguous_memory.py`` (the closest
hnsw dense reference), and is split into two parts:
1. **Surface tests** — verify that ``VamanaIndexParam`` / ``VamanaQueryParam``
are correctly bound: construction defaults, readonly properties,
``to_dict``, ``__repr__``, pickle round-trip, and that they appear in the
public ``zvec`` namespace with the expected ``IndexType.VAMANA`` value.
2. **End-to-end tests** — build a collection that uses Vamana on a dense
FP32 column, insert deterministic documents, then run a top-k query
through ``VamanaQueryParam`` on both the writer segment and the
persisted (post-``optimize()``) segment.
"""
from __future__ import annotations
import pickle
import sys
import numpy as np
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
CollectionSchema,
Doc,
FieldSchema,
InvertIndexParam,
VamanaIndexParam,
VamanaQueryParam,
Query,
VectorSchema,
)
from zvec.typing import DataType, IndexType, MetricType, QuantizeType
DIMENSION = 32
NUM_DOCS = 128
TOPK = 5
# Defaults pulled from src/include/zvec/core/interface/constants.h. Keep
# in sync with kDefaultVamana* if the engine defaults ever change.
DEFAULT_MAX_DEGREE = 64
DEFAULT_SEARCH_LIST_SIZE = 100
DEFAULT_ALPHA = 1.2
DEFAULT_EF_SEARCH = 200
DEFAULT_SATURATE_GRAPH = False
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _build_schema(
name: str,
*,
metric_type: MetricType = MetricType.IP,
max_degree: int = 32,
search_list_size: int = 64,
alpha: float = 1.2,
use_contiguous_memory: bool = False,
) -> CollectionSchema:
"""Create a simple schema with a single FP32 Vamana vector column."""
return CollectionSchema(
name=name,
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=DIMENSION,
index_param=VamanaIndexParam(
metric_type=metric_type,
max_degree=max_degree,
search_list_size=search_list_size,
alpha=alpha,
use_contiguous_memory=use_contiguous_memory,
),
),
],
)
def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]:
"""Produce deterministic documents for insertion."""
docs: list[Doc] = []
for i in range(num):
vec = rng.standard_normal(DIMENSION).astype(np.float32)
docs.append(
Doc(
id=str(i),
fields={"id": i},
vectors={"dense": vec.tolist()},
)
)
return docs
def _query_topk(
coll: Collection, query_vec: list[float], *, ef_search: int = 64
) -> list[str]:
"""Run a top-k vector query and return the returned ids in order."""
vector_query = Query(
field_name="dense",
vector=query_vec,
param=VamanaQueryParam(ef_search=ef_search),
)
hits = coll.query(vector_query, topk=TOPK)
assert hits is not None, "query returned None"
assert len(hits) >= 1, f"expected at least one hit, got {hits!r}"
return [doc.id for doc in hits]
# ---------------------------------------------------------------------------
# 1) Surface: construction / property / to_dict / repr / pickle / namespace
# ---------------------------------------------------------------------------
class TestVamanaIndexParamSurface:
"""Verify the Python binding for ``VamanaIndexParam``."""
def test_defaults(self):
param = VamanaIndexParam()
assert param.type == IndexType.VAMANA
assert param.metric_type == MetricType.IP
assert param.max_degree == DEFAULT_MAX_DEGREE
assert param.search_list_size == DEFAULT_SEARCH_LIST_SIZE
assert param.alpha == pytest.approx(DEFAULT_ALPHA)
assert param.saturate_graph is DEFAULT_SATURATE_GRAPH
assert param.use_contiguous_memory is False
assert param.use_id_map is False
assert param.quantize_type == QuantizeType.UNDEFINED
def test_custom_construction(self):
param = VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=48,
search_list_size=128,
alpha=1.5,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.INT8,
)
assert param.type == IndexType.VAMANA
assert param.metric_type == MetricType.COSINE
assert param.max_degree == 48
assert param.search_list_size == 128
assert param.alpha == pytest.approx(1.5)
assert param.saturate_graph is True
assert param.use_contiguous_memory is True
assert param.use_id_map is False
assert param.quantize_type == QuantizeType.INT8
def test_to_dict_includes_all_fields(self):
param = VamanaIndexParam(
metric_type=MetricType.L2,
max_degree=32,
search_list_size=80,
alpha=1.3,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.FP16,
)
data = param.to_dict()
assert data["type"] == "VAMANA"
assert data["metric_type"] == "L2"
assert data["max_degree"] == 32
assert data["search_list_size"] == 80
assert data["alpha"] == pytest.approx(1.3)
assert data["saturate_graph"] is True
assert data["use_contiguous_memory"] is True
assert data["use_id_map"] is False
assert data["quantize_type"] == "FP16"
def test_repr_contains_key_fields(self):
text = repr(
VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=24,
search_list_size=72,
alpha=1.4,
saturate_graph=True,
use_contiguous_memory=True,
)
)
# Spot-check the most diagnostic fields are rendered.
assert "VAMANA" in text
assert "COSINE" in text
assert "max_degree" in text and "24" in text
assert "search_list_size" in text and "72" in text
assert "alpha" in text
assert "saturate_graph" in text and "true" in text
assert "use_contiguous_memory" in text and "true" in text
@pytest.mark.parametrize(
"field, kwargs",
[
("max_degree", dict(max_degree=99)),
("search_list_size", dict(search_list_size=99)),
("alpha", dict(alpha=1.7)),
("saturate_graph", dict(saturate_graph=True)),
("use_contiguous_memory", dict(use_contiguous_memory=True)),
("use_id_map", dict(use_id_map=True)),
],
)
def test_readonly_properties(self, field, kwargs):
param = VamanaIndexParam(**kwargs)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, field, getattr(param, field))
def test_pickle_roundtrip(self):
original = VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=48,
search_list_size=120,
alpha=1.4,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.INT8,
)
restored = pickle.loads(pickle.dumps(original))
assert restored.type == IndexType.VAMANA
assert restored.metric_type == MetricType.COSINE
assert restored.max_degree == 48
assert restored.search_list_size == 120
assert restored.alpha == pytest.approx(1.4)
assert restored.saturate_graph is True
assert restored.use_contiguous_memory is True
assert restored.use_id_map is False
assert restored.quantize_type == QuantizeType.INT8
# to_dict equality is the strongest end-to-end equivalence we have.
assert restored.to_dict() == original.to_dict()
class TestVamanaQueryParamSurface:
"""Verify the Python binding for ``VamanaQueryParam``."""
def test_defaults(self):
q = VamanaQueryParam()
assert q.type == IndexType.VAMANA
assert q.ef_search == DEFAULT_EF_SEARCH
assert q.radius == pytest.approx(0.0)
assert q.is_linear is False
assert q.is_using_refiner is False
assert q.prefetch_offset == 8
assert q.prefetch_lines == 0
def test_custom_construction(self):
q = VamanaQueryParam(
ef_search=300,
radius=0.5,
is_linear=True,
is_using_refiner=True,
extra_params={
"prefetch_offset": 8,
"prefetch_lines": 2,
},
)
assert q.type == IndexType.VAMANA
assert q.ef_search == 300
assert q.radius == pytest.approx(0.5)
assert q.is_linear is True
assert q.is_using_refiner is True
assert q.prefetch_offset == 8
assert q.prefetch_lines == 2
def test_repr_contains_key_fields(self):
text = repr(VamanaQueryParam(ef_search=128, radius=0.25))
assert "VAMANA" in text
assert "ef_search" in text and "128" in text
assert "radius" in text
def test_readonly_ef_search(self):
q = VamanaQueryParam(ef_search=100)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
q.ef_search = 200 # type: ignore[misc]
def test_pickle_roundtrip(self):
original = VamanaQueryParam(
ef_search=256,
radius=0.3,
is_linear=False,
is_using_refiner=True,
extra_params={
"prefetch_offset": 4,
"prefetch_lines": 3,
},
)
restored = pickle.loads(pickle.dumps(original))
assert restored.type == IndexType.VAMANA
assert restored.ef_search == 256
assert restored.radius == pytest.approx(0.3)
assert restored.is_linear is False
assert restored.is_using_refiner is True
assert restored.prefetch_offset == 4
assert restored.prefetch_lines == 3
class TestVamanaPublicNamespace:
"""The Vamana entry points must be importable from the top-level ``zvec``."""
def test_top_level_exports(self):
assert zvec.VamanaIndexParam is VamanaIndexParam
assert zvec.VamanaQueryParam is VamanaQueryParam
assert "VamanaIndexParam" in zvec.__all__
assert "VamanaQueryParam" in zvec.__all__
def test_index_type_enum_member(self):
# Sanity: the IndexType enum exposes VAMANA and it is what the
# bound params advertise.
assert IndexType.VAMANA is not None
assert VamanaIndexParam().type == IndexType.VAMANA
assert VamanaQueryParam().type == IndexType.VAMANA
# ---------------------------------------------------------------------------
# 2) End-to-end: create collection, insert, query through the writer segment
# ---------------------------------------------------------------------------
@pytest.fixture
def rng() -> np.random.Generator:
return np.random.default_rng(seed=42)
# Mirror the hnsw dense test fixture: only the mmap-backed variant is
# currently usable for vector index construction. BufferPool (enable_mmap=
# False) is intentionally omitted because the same write-path guard in
# ``SegmentImpl::merge_vector_indexer`` rejects that combination.
@pytest.fixture(params=[True], ids=["mmap_on"])
def collection_option(request) -> CollectionOption:
return CollectionOption(read_only=False, enable_mmap=request.param)
class TestVamanaEndToEnd:
"""End-to-end: schema -> create_and_open -> insert -> query works."""
def test_schema_round_trip(self, tmp_path_factory, collection_option):
"""The Vamana index params survive the schema persist path."""
schema = _build_schema(
"vamana_schema_rt",
metric_type=MetricType.COSINE,
max_degree=32,
search_list_size=80,
alpha=1.3,
use_contiguous_memory=True,
)
path = tmp_path_factory.mktemp("zvec") / "vamana_schema_rt"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
vec_schema = coll.schema.vectors[0]
ip = vec_schema.index_param
assert ip.type == IndexType.VAMANA
assert ip.metric_type == MetricType.COSINE
assert ip.max_degree == 32
assert ip.search_list_size == 80
assert ip.alpha == pytest.approx(1.3)
assert ip.use_contiguous_memory is True
finally:
coll.destroy()
def test_insert_and_query_self_recall(
self, tmp_path_factory, collection_option, rng
):
"""Top-1 of a query equal to an inserted vector must be that vector.
Exercises the writer-segment Vamana streamer end-to-end through the
Python entry point: ``VamanaIndexParam`` for build and
``VamanaQueryParam`` for search.
"""
schema = _build_schema("vamana_e2e_recall")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_recall"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok(), f"insert failed: code={r.code()}"
assert coll.stats.doc_count == NUM_DOCS
# Self-recall: query with the i-th inserted vector, expect id i
# to be the top result.
for probe in (0, 7, 42, NUM_DOCS - 1):
query_vec = docs[probe].vector("dense")
ids = _query_topk(coll, query_vec)
assert ids[0] == str(probe), (
f"expected self-recall at probe={probe}, got top-1 id={ids[0]} "
f"(top-{TOPK}={ids})"
)
finally:
coll.destroy()
def test_query_param_ef_search_affects_only_quality(
self, tmp_path_factory, collection_option, rng
):
"""``ef_search`` is a search-time knob and must not crash for any
sensible value. Larger ``ef_search`` should be at least as good as
smaller for self-recall."""
schema = _build_schema("vamana_e2e_ef")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_ef"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
query_vec = docs[3].vector("dense")
ids_small = _query_topk(coll, query_vec, ef_search=16)
ids_large = _query_topk(coll, query_vec, ef_search=256)
# Both should self-recall the probe vector at top-1.
assert ids_small[0] == "3"
assert ids_large[0] == "3"
assert len(ids_small) == TOPK
assert len(ids_large) == TOPK
finally:
coll.destroy()
def test_optimize_then_query(self, tmp_path_factory, collection_option, rng):
"""The persisted Vamana segment built by ``optimize()`` must serve
queries correctly.
Until the cmake fix to force-load ``core_knn_vamana_static`` into the
``_zvec`` pybind module, this path failed at ``VamanaStreamer``
creation because the global factory registration in
``vamana_streamer.cc`` was never linked in. This test pins down the
regression.
"""
schema = _build_schema("vamana_e2e_optimize")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_optimize"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
assert coll.stats.doc_count == NUM_DOCS
# Snapshot the writer-segment top-k for a probe vector.
query_vec = docs[5].vector("dense")
ids_pre = _query_topk(coll, query_vec)
assert ids_pre[0] == "5"
# Trigger persisted segment build. Pre-fix this raised
# RuntimeError("Failed to create index").
coll.optimize()
# Persisted segment must still serve queries with the same
# top-1 self-recall guarantee. We do not assert full top-k
# equality with the writer segment because the persisted
# streamer may visit nodes in a different order; top-1 self-
# recall is the strong invariant.
ids_post = _query_topk(coll, query_vec)
assert ids_post[0] == "5", (
f"post-optimize top-1 should still be probe id, got {ids_post}"
)
assert len(ids_post) == TOPK
finally:
coll.destroy()
python/zvec/__init__.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import sys
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from importlib.metadata import PackageNotFoundError
# zvec ships a native C++ extension that is only built and tested for 64-bit
# CPython. A 32-bit interpreter would fail to load the extension with an
# obscure error, so fail fast here with an actionable message.
if sys.maxsize <= 2**32:
raise ImportError(
"zvec requires a 64-bit Python interpreter; "
"the current interpreter is 32-bit and is not supported."
)
# Register the wheel-bundled jieba dict dir so `import zvec` alone makes
# the jieba FTS tokenizer usable. Users can still override via
# zvec.init(jieba_dict_dir=...), zvec.set_default_jieba_dict_dir(...),
# ZVEC_JIEBA_DICT_DIR, or per-field FtsIndexParam.extra_params.
try:
from importlib.resources import files as _resource_files
from zvec._zvec import (
get_default_jieba_dict_dir,
set_default_jieba_dict_dir,
)
set_default_jieba_dict_dir(str(_resource_files("zvec").joinpath("data/jieba_dict")))
except Exception:
# Custom builds without bundled dict; users must configure explicitly.
pass
# ==============================
# Public API — grouped by category
# ==============================
# —— DiskAnn runtime plugin ——
# Re-export the plugin management entry points defined by the C++ extension.
# DiskAnn normally auto-loads on first use; these APIs let tests and
# diagnostic tools preload the plugin and get a clear error if libaio is
# missing or the plugin shared object cannot be located.
from zvec._zvec import (
DISKANN_PLUGIN_DLOPEN_FAILED,
DISKANN_PLUGIN_LIBAIO_MISSING,
DISKANN_PLUGIN_OK,
DISKANN_PLUGIN_UNSUPPORTED_PLATFORM,
is_diskann_plugin_loaded,
is_libaio_available,
load_diskann_plugin,
)
from . import model as model
# —— Extensions ——
from .extension import (
BM25EmbeddingFunction,
DefaultLocalDenseEmbedding,
DefaultLocalReRanker,
DefaultLocalSparseEmbedding,
DenseEmbeddingFunction,
OpenAIDenseEmbedding,
OpenAIFunctionBase,
QwenDenseEmbedding,
QwenFunctionBase,
QwenReRanker,
QwenSparseEmbedding,
ReRanker,
RrfReRanker,
SentenceTransformerFunctionBase,
SparseEmbeddingFunction,
WeightedReRanker,
)
# —— Typing ——
from .model import param as param
from .model import schema as schema
# —— Core data structures ——
from .model.collection import Collection
from .model.doc import Doc, DocList
# —— Query & index parameters ——
# —— FTS params (C++ binding) ——
from .model.param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
DiskAnnIndexParam,
DiskAnnQueryParam,
FlatIndexParam,
FtsIndexParam,
FtsQueryParam,
HnswIndexParam,
HnswQueryParam,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
from .model.param.query import Fts, Query, VectorQuery
# —— Schema & field definitions ——
from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema
# —— tools ——
from .tool import require_module
from .typing import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
from .typing.enum import LogLevel, LogType
# —— lifecycle ——
from .zvec import create_and_open, init, open
# ==============================
# Public interface declaration
# ==============================
__all__ = [
# Zvec functions
"create_and_open",
"init",
"open",
"set_default_jieba_dict_dir",
"get_default_jieba_dict_dir",
# Core classes
"Collection",
"Doc",
"DocList",
# Schema
"CollectionSchema",
"FieldSchema",
"VectorSchema",
"CollectionStats",
# Parameters
"Query",
"VectorQuery",
"Fts",
"FtsIndexParam",
"FtsQueryParam",
"InvertIndexParam",
"HnswIndexParam",
"HnswRabitqIndexParam",
"FlatIndexParam",
"IVFIndexParam",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"CollectionOption",
"IndexOption",
"OptimizeOption",
"AddColumnOption",
"AlterColumnOption",
"HnswQueryParam",
"HnswRabitqQueryParam",
"IVFQueryParam",
"QuantizerParam",
"VamanaIndexParam",
"VamanaQueryParam",
# Extensions
"DenseEmbeddingFunction",
"SparseEmbeddingFunction",
"QwenFunctionBase",
"OpenAIFunctionBase",
"SentenceTransformerFunctionBase",
"ReRanker",
"DefaultLocalDenseEmbedding",
"DefaultLocalSparseEmbedding",
"BM25EmbeddingFunction",
"OpenAIDenseEmbedding",
"QwenDenseEmbedding",
"QwenSparseEmbedding",
"RrfReRanker",
"WeightedReRanker",
"DefaultLocalReRanker",
"QwenReRanker",
# Typing
"DataType",
"MetricType",
"QuantizeType",
"IndexType",
"LogLevel",
"LogType",
"Status",
"StatusCode",
# Tools
"require_module",
# DiskAnn plugin
"load_diskann_plugin",
"is_diskann_plugin_loaded",
"is_libaio_available",
"DISKANN_PLUGIN_OK",
"DISKANN_PLUGIN_UNSUPPORTED_PLATFORM",
"DISKANN_PLUGIN_LIBAIO_MISSING",
"DISKANN_PLUGIN_DLOPEN_FAILED",
]
# ==============================
# Version handling
# ==============================
__version__: str
try:
from importlib.metadata import version
except ImportError:
from importlib_metadata import version # Python < 3.8
try:
__version__ = version("zvec")
except Exception:
__version__ = "unknown"
python/zvec/__init__.pyi
+202
View File
@@ -0,0 +1,202 @@
"""
Zvec core module
"""
from __future__ import annotations
import collections
from . import typing
from .extension import ReRanker, RrfReRanker, WeightedReRanker
from .extension.embedding import DenseEmbeddingFunction
from .model import param, schema
from .model.collection import Collection
from .model.doc import Doc, DocList
from .model.param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
DiskAnnIndexParam,
DiskAnnQueryParam,
FlatIndexParam,
FtsIndexParam,
FtsQueryParam,
HnswIndexParam,
HnswQueryParam,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
from .model.param.query import Fts, Query, VectorQuery
from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema
from .tool import require_module
from .typing import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
from .typing.enum import LogLevel, LogType
from .zvec import create_and_open, init, open
__all__: list = [
"AddColumnOption",
"AlterColumnOption",
"Collection",
"CollectionOption",
"CollectionSchema",
"CollectionStats",
"DataType",
"DenseEmbeddingFunction",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"Doc",
"DocList",
"FieldSchema",
"FlatIndexParam",
"Fts",
"FtsIndexParam",
"FtsQueryParam",
"HnswIndexParam",
"HnswQueryParam",
"HnswRabitqIndexParam",
"HnswRabitqQueryParam",
"IVFIndexParam",
"IVFQueryParam",
"IndexOption",
"IndexType",
"InvertIndexParam",
"LogLevel",
"LogType",
"MetricType",
"OptimizeOption",
"QuantizeType",
"QuantizerParam",
"Query",
"ReRanker",
"RrfReRanker",
"Status",
"StatusCode",
"VamanaIndexParam",
"VamanaQueryParam",
"VectorQuery",
"VectorSchema",
"WeightedReRanker",
"create_and_open",
"init",
"open",
"require_module",
]
class _Collection:
@staticmethod
def CreateAndOpen(
arg0: str, arg1: schema._CollectionSchema, arg2: param.CollectionOption
) -> _Collection: ...
@staticmethod
def Open(arg0: str, arg1: param.CollectionOption) -> _Collection: ...
def AddColumn(
self,
arg0: schema._FieldSchema,
arg1: str,
arg2: param.AddColumnOption,
) -> None: ...
def AlterColumn(
self,
arg0: str,
arg1: str,
arg2: schema._FieldSchema,
arg3: param.AlterColumnOption,
) -> None: ...
def CreateIndex(
self, arg0: str, arg1: param.IndexParam, arg2: param.IndexOption
) -> None: ...
def Delete(self, arg0: collections.abc.Sequence[str]) -> list[typing.Status]: ...
def DeleteByFilter(self, arg0: str) -> None: ...
def Destroy(self) -> None: ...
def DropColumn(self, arg0: str) -> None: ...
def DropIndex(self, arg0: str) -> None: ...
def Fetch(
self,
pks: collections.abc.Sequence[str],
output_fields: list[str] | None = None,
include_vector: bool = True,
) -> dict[str, _Doc]: ...
def Flush(self) -> None: ...
def GroupByQuery(self, arg0: ...) -> list[...]: ...
def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def Optimize(self, arg0: param.OptimizeOption) -> None: ...
def Options(self) -> param.CollectionOption: ...
def Path(self) -> str: ...
def Query(self, arg0: param._SearchQuery) -> list[_Doc]: ...
def Schema(self) -> schema._CollectionSchema: ...
def Stats(self) -> schema.CollectionStats: ...
def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def Upsert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def _debug_hnsw_storage_mode(self, column_name: str) -> str:
"""Debug-only: returns the storage mode of the HNSW entity on the
given vector column. One of 'mmap', 'buffer_pool', 'contiguous'.
Raises KeyError if no HNSW index exists on the column, or
ValueError if the column's index is not an HNSW index. Intended
for introspection and testing only; not part of the stable API."""
def __getstate__(self) -> tuple: ...
def __setstate__(self, arg0: tuple) -> None: ...
class _Doc:
def __getstate__(self) -> bytes: ...
def __init__(self) -> None: ...
def __setstate__(self, arg0: bytes) -> None: ...
def field_names(self) -> list[str]: ...
def get_any(self, arg0: str, arg1: typing.DataType) -> typing.Any: ...
def has_field(self, arg0: str) -> bool: ...
def pk(self) -> str: ...
def score(self) -> float: ...
def set_any(self, arg0: str, arg1: typing.DataType, arg2: typing.Any) -> bool: ...
def set_pk(self, arg0: str) -> None: ...
def set_score(self, arg0: typing.SupportsFloat) -> None: ...
class _DocOp:
"""
Members:
INSERT
UPDATE
DELETE
UPSERT
"""
DELETE: typing.ClassVar[_DocOp] # value = <_DocOp.DELETE: 3>
INSERT: typing.ClassVar[_DocOp] # value = <_DocOp.INSERT: 0>
UPDATE: typing.ClassVar[_DocOp] # value = <_DocOp.UPDATE: 2>
UPSERT: typing.ClassVar[_DocOp] # value = <_DocOp.UPSERT: 1>
__members__: typing.ClassVar[
dict[str, _DocOp]
] # value = {'INSERT': <_DocOp.INSERT: 0>, 'UPDATE': <_DocOp.UPDATE: 2>, 'DELETE': <_DocOp.DELETE: 3>, 'UPSERT': <_DocOp.UPSERT: 1>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
python/zvec/common/__init__.py
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .constants import DenseVectorType, SparseVectorType, VectorType
__all__ = ["DenseVectorType", "SparseVectorType", "VectorType"]
python/zvec/common/constants.py
+33
View File
@@ -0,0 +1,33 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Optional, TypeVar, Union
import numpy as np
# VectorType: DenseVectorType | SparseVectorType
DenseVectorType = Union[list[float], list[int], np.ndarray]
SparseVectorType = dict[int, float]
VectorType = Optional[Union[DenseVectorType, SparseVectorType]]
# Embeddable: Text | Image | Audio
TEXT = str
IMAGE = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
AUDIO = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
Embeddable = Optional[Union[TEXT, IMAGE, AUDIO]]
# Multimodal Embeddable
MD = TypeVar("MD", bound=Embeddable, contravariant=True)
python/zvec/executor/__init__.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .query_executor import (
QueryContext,
QueryExecutor,
)
__all__ = [
"QueryContext",
"QueryExecutor",
]
python/zvec/executor/query_executor.py
+266
View File
@@ -0,0 +1,266 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Optional, Union
import numpy as np
from zvec._zvec import _Collection, _MultiQuery
from zvec._zvec.param import _Fts, _SearchQuery, _SubQuery
from ..extension import CallbackReRanker, ReRanker, RrfReRanker, WeightedReRanker
from ..model.convert import convert_to_py_doc
from ..model.doc import DocList
from ..model.param.query import Query
from ..model.schema import CollectionSchema
from ..typing import DataType
__all__ = [
"QueryContext",
"QueryExecutor",
]
DTYPE_MAP = {
DataType.VECTOR_FP16.value: np.float16,
DataType.VECTOR_FP32.value: np.float32,
DataType.VECTOR_FP64.value: np.float64,
DataType.VECTOR_INT8.value: np.int8,
}
def convert_to_numpy(vec: Union[list, np.ndarray], dtype: np.dtype) -> np.ndarray:
if isinstance(vec, np.ndarray):
if vec.dtype == dtype and vec.ndim == 1:
return vec
return np.asarray(vec, dtype=dtype).flatten()
try:
arr = np.asarray(vec, dtype=dtype)
if arr.ndim != 1:
arr = arr.flatten()
return arr
except (ValueError, TypeError) as e:
raise TypeError(
f"Cannot convert input to 1D numpy array with dtype={dtype}: {type(vec)}"
) from e
class QueryContext:
def __init__(
self,
topk: int,
filter: Optional[str] = None,
include_vector: bool = False,
queries: Optional[list[Query]] = None,
output_fields: Optional[list[str]] = None,
reranker: Optional[ReRanker] = None,
):
# query param
self._filter = filter
self._queries = queries or []
self._topk = topk
self._include_vector = include_vector
self._output_fields = output_fields
# reranker
self._reranker = reranker
@property
def topk(self):
return self._topk
@property
def queries(self):
return self._queries
@property
def filter(self):
return self._filter
@property
def reranker(self):
return self._reranker
@property
def output_fields(self):
return self._output_fields
@property
def include_vector(self):
return self._include_vector
class QueryExecutor:
"""Unified query executor that routes based on query count and reranker type."""
def __init__(self, schema: CollectionSchema):
self._schema = schema
def _build_queries(
self, ctx: QueryContext, collection: _Collection
) -> list[_SearchQuery]:
"""Build query vector list (no validation, conversion only)."""
if not ctx.queries:
return [self._build_base_search_query(ctx)]
return [
self._build_search_query(ctx, query, collection) for query in ctx.queries
]
def execute(self, ctx: QueryContext, collection: _Collection) -> DocList:
"""Execute a query, routing by query count.
A single (or vector-less) query is sent to C++ as a ``_SearchQuery``;
multiple queries are assembled into a ``_MultiQuery``.
"""
queries = self._build_queries(ctx, collection)
if not queries:
raise ValueError("No query to execute")
if len(queries) == 1:
return self._execute_single_query(queries[0], collection)
return self._execute_multi_query(ctx, queries, collection)
def _execute_single_query(
self, query: _SearchQuery, collection: _Collection
) -> DocList:
"""Single/vector-less query: send a ``_SearchQuery`` to C++."""
docs = collection.Query(query)
return [convert_to_py_doc(doc, self._schema) for doc in docs]
def _execute_multi_query(
self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection
) -> DocList:
"""Multiple queries: send a ``_MultiQuery`` to C++.
A Python-only reranker (e.g. a model/API-based one) cannot run inside
the C++ MultiQuery, so each route is executed individually and merged by
the reranker in Python. The built-in RRF/Weighted/Callback rerankers use
the C++ variant-based fast path.
"""
reranker = ctx.reranker
if reranker is None:
raise ValueError(
"A reranker is required to merge results from multiple queries; "
"specify the 'reranker' argument."
)
if not isinstance(reranker, (RrfReRanker, WeightedReRanker, CallbackReRanker)):
docs_list = self._execute_python_pipeline(queries, collection)
return self._merge_and_rerank(ctx, docs_list)
multi_query = self._build_multi_query(ctx, queries)
docs = collection.Query(multi_query)
return [convert_to_py_doc(doc, self._schema) for doc in docs]
def _build_multi_query(
self, ctx: QueryContext, queries: list[_SearchQuery]
) -> _MultiQuery:
"""Assemble a C++ ``_MultiQuery`` from per-route ``_SearchQuery`` objects."""
multi_query = _MultiQuery()
multi_query.queries = [_SubQuery.from_search_query(query) for query in queries]
# num_candidates controls per-sub-query candidate count for reranking pool.
# It must NOT be limited to the final output topk; use at least the C++
# SubQuery default of 10 to ensure sufficient candidates for reranking.
_DEFAULT_NUM_CANDIDATES = 10
for sub in multi_query.queries:
sub.num_candidates = max(ctx.topk, _DEFAULT_NUM_CANDIDATES)
multi_query.topk = ctx.topk
if ctx.filter:
multi_query.filter = ctx.filter
multi_query.include_vector = ctx.include_vector
if ctx.output_fields is not None:
multi_query.output_fields = ctx.output_fields
# Set rerank strategy via the C++ variant-based API.
reranker = ctx.reranker
if isinstance(reranker, RrfReRanker):
multi_query.set_rerank_rrf(reranker.rank_constant)
elif isinstance(reranker, WeightedReRanker):
multi_query.set_rerank_weighted(reranker.weights)
elif isinstance(reranker, CallbackReRanker):
multi_query.set_rerank_callback(reranker._callback)
return multi_query
def _execute_python_pipeline(
self, vectors: list[_SearchQuery], collection: _Collection
) -> list[DocList]:
"""Execute queries serially for the Python-only reranker path."""
return [self._execute_single_query(query, collection) for query in vectors]
def _merge_and_rerank(self, ctx: QueryContext, docs_list: list[DocList]) -> DocList:
"""Merge and rerank results from the Python pipeline path."""
if not docs_list:
raise ValueError("Query results is empty")
if len(docs_list) == 1 and not ctx.reranker:
return docs_list[0]
return ctx.reranker.rerank(docs_list, ctx.topk)
def _build_base_search_query(self, ctx: QueryContext) -> _SearchQuery:
search_query = _SearchQuery()
search_query.topk = ctx.topk
search_query.include_vector = ctx.include_vector
if ctx.filter:
search_query.filter = ctx.filter
if ctx.output_fields is not None:
search_query.output_fields = ctx.output_fields
return search_query
def _apply_fts(self, query: Query, search_query: _SearchQuery) -> None:
"""Set FTS query on search_query if the query has FTS parameters."""
if query.has_fts():
fts = _Fts()
fts.query_string = query.fts.query_string or ""
fts.match_string = query.fts.match_string or ""
search_query.fts = fts
def _build_search_query(
self, ctx: QueryContext, query: Query, collection: _Collection
) -> _SearchQuery:
query._validate()
search_query = self._build_base_search_query(ctx)
search_query.field_name = query.field_name
if query.param:
search_query.query_params = query.param
# set FTS query if provided
self._apply_fts(query, search_query)
vector_schema = None
if query.has_vector() or query.has_id():
vector_schema = (
self._schema.vector(query.field_name)
if query
else self._schema.vectors[0]
)
if vector_schema is None:
raise ValueError("No vector field found")
# set vector
if query.has_vector():
vec_data = query.vector
elif query.has_id():
fetched = collection.Fetch([query.id])
doc = next(iter(fetched.values()), None)
if not doc:
raise ValueError(f"Document with id '{query.id}' not found")
vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)
else:
return search_query
target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)
search_query.set_vector(
vector_schema._get_object(),
convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,
)
return search_query
python/zvec/extension/__init__.py
+58
View File
@@ -0,0 +1,58 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .bm25_embedding_function import BM25EmbeddingFunction
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .http_embedding_function import HTTPDenseEmbedding
from .jina_embedding_function import JinaDenseEmbedding
from .jina_function import JinaFunctionBase
from .multi_vector_reranker import CallbackReRanker, RrfReRanker, WeightedReRanker
from .openai_embedding_function import OpenAIDenseEmbedding
from .openai_function import OpenAIFunctionBase
from .qwen_embedding_function import QwenDenseEmbedding, QwenSparseEmbedding
from .qwen_function import QwenFunctionBase
from .qwen_rerank_function import QwenReRanker
from .rerank_function import RerankFunction
from .rerank_function import RerankFunction as ReRanker
from .sentence_transformer_embedding_function import (
DefaultLocalDenseEmbedding,
DefaultLocalSparseEmbedding,
)
from .sentence_transformer_function import SentenceTransformerFunctionBase
from .sentence_transformer_rerank_function import DefaultLocalReRanker
__all__ = [
"BM25EmbeddingFunction",
"CallbackReRanker",
"DefaultLocalDenseEmbedding",
"DefaultLocalReRanker",
"DefaultLocalSparseEmbedding",
"DenseEmbeddingFunction",
"HTTPDenseEmbedding",
"JinaDenseEmbedding",
"JinaFunctionBase",
"OpenAIDenseEmbedding",
"OpenAIFunctionBase",
"QwenDenseEmbedding",
"QwenFunctionBase",
"QwenReRanker",
"QwenSparseEmbedding",
"ReRanker",
"RerankFunction",
"RrfReRanker",
"SentenceTransformerFunctionBase",
"SparseEmbeddingFunction",
"WeightedReRanker",
]

Some files were not shown because too many files have changed in this diff Show More