From 058e10b5bc3bcb715a853e540fc27501f9eb366d Mon Sep 17 00:00:00 2001 From: lc285652 Date: Fri, 13 Nov 2026 11:22:29 +0800 Subject: [PATCH 1/3] chore(vendor): track zvec v0.8.1 8397071 (v0.8.0-beta2) -> 30d662d (v0.8.1), 30 upstream commits. c_api.h is not byte-identical: 4470 -> 4540 lines, +4 functions, +2 types, +2 constants, nothing removed and no breaking signature change. - feat(turbo): support product quantization end to end (#833) adds zvec_index_params_set_quantizer_pq / _get_quantizer_pq, zvec_pq_quantizer_params_t, ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT and ZVEC_QUANTIZE_TYPE_PQ, and rewords the enable_rotate doc: with PQ the same flag now selects a learned OPQ rotation. - feat(c-api): expose grouped vector search (#845) adds zvec_collection_group_by_query, zvec_group_results_free and zvec_group_result_t. zvec_group_by_vector_query_* was already bound, but until now a grouped query could only be built, never executed. - zvec_collection_multi_query's doc now spells out the NULL / zero-result contract. Comment only, no signature change. Beyond the header: DiskANN is now built for Intel macOS and Windows ARM64 (#874), and RaBitQ for Linux/macOS/Android/iOS ARM64 plus Intel macOS (#832), which leaves Windows ARM64 the only prebuilt target without RaBitQ. zvec/NOTICE is unchanged, so no SDK NOTICE needs regenerating. The Rust adaptation follows in the next commit. --- vendor/zvec | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vendor/zvec b/vendor/zvec index 8397071..30d662d 160000 --- a/vendor/zvec +++ b/vendor/zvec @@ -1 +1 @@ -Subproject commit 8397071d815cd842be2068339e00a94e69f069d1 +Subproject commit 30d662db4d2467937d843443074b34743bcc9baa From 2d652f630b557ccb7e311489cd2f748d8e1cac5b Mon Sep 17 00:00:00 2001 From: lc285652 Date: Fri, 13 Nov 2026 11:23:49 +0800 Subject: [PATCH 2/3] feat: adapt to zvec v0.8.1 grouped search and IVF product quantization Grouped search. GroupBySearchQuery could be built but never executed, because the C API only gained zvec_collection_group_by_query in v0.8.1. Collection::group_by_query now runs it and returns Vec, each group carrying its grouping value and its documents in similarity order. zvec_group_results_free is the only release path for a zvec_group_result_t array, and that array owns everything inside it: the grouping values, each group's document pointer array, and the documents themselves. collect_group_results therefore moves every document into an owning Doc, releases the pointer array with zvec_free, and nulls docs/doc_count before handing the shell back, so the engine is left with only the strings and the outer array to free and cannot destroy a document twice. grouped_results_outlive_the_query_and_the_collection pins that contract down: the groups stay readable after both handles are gone. Product quantization. PqQuantizerParams mirrors zvec_pq_quantizer_params_t, its Default reproduces ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT, and IndexParams::set_quantizer_pq / quantizer_pq wrap the new setter and getter. QuantizeType::Pq (8) joins the enum mirror, its From mapping and the roundtrip test. The C API accepts PQ on IVF params only, which the wrapper documents and the tests assert along with the num_chunk / num_bits / fast_scan / OPQ validation rules. set_quantizer_enable_rotate's docs said rotation "has no effect with other quantize types"; since #833 it selects learned OPQ with PQ, and it writes the same QuantizerParam field that set_quantizer_pq writes, so both docs now say which call wins. Platform gates moved without the header (#874, #832). The DiskANN and IVF RaBitQ end-to-end tests already skip on NotSupported, so they widen their own coverage with no code change -- on Apple Silicon the RaBitQ test now runs instead of skipping. Their comments and both READMEs' availability prose are updated to match, and the Windows ARM64 note now says RaBitQ is the only thing missing there. Also adds the group_by_search example (make example-group-by), an IVF PQ section to schema_builder, and grouped-search plus PQ docs to README.md and README_CN.md. Verified against the zvec v0.8.1 osx-arm64 SDK, whose c_api.h is byte-identical to the submodule's: cargo fmt --check, clippy --workspace --all-targets -D warnings, build --workspace --all-targets and make test-all (unit + integration + doc) all pass. --- Makefile | 4 + README.md | 63 ++++++++- README_CN.md | 58 ++++++++- zvec-sys/src/lib.rs | 42 ++++++ zvec/Cargo.toml | 6 +- zvec/examples/group_by_search.rs | 121 ++++++++++++++++++ zvec/examples/schema_builder.rs | 23 ++++ zvec/src/collection.rs | 138 +++++++++++++++++++- zvec/src/lib.rs | 8 +- zvec/src/schema.rs | 205 +++++++++++++++++++++++++++++- zvec/src/types.rs | 9 ++ zvec/tests/group_by_query_test.rs | 201 +++++++++++++++++++++++++++++ zvec/tests/integration_test.rs | 14 +- 13 files changed, 864 insertions(+), 28 deletions(-) create mode 100644 zvec/examples/group_by_search.rs create mode 100644 zvec/tests/group_by_query_test.rs diff --git a/Makefile b/Makefile index 79caf65..35d3a88 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,7 @@ bench coverage coverage-html coverage-lcov coverage-text \ fuzz fuzz-types fuzz-schema fuzz-doc fuzz-query fuzz-collection \ example example-basic example-schema example-search example-crud example-config \ + example-group-by \ lint setup help \ submodule-init submodule-update zvec-build zvec-clean zvec-rebuild @@ -158,6 +159,9 @@ example-crud: ## Run crud_operations example example-config: ## Run config_logging example $(CARGO) run --example config_logging +example-group-by: ## Run group_by_search example + $(CARGO) run --example group_by_search + # --------------------------------------------------------------------------- # Documentation # --------------------------------------------------------------------------- diff --git a/README.md b/README.md index 9e58f8d..31c25d3 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ Safe, idiomatic Rust bindings for the [zvec](https://github.com/alibaba/zvec) ve - **Zero-Copy Where Possible** — Minimizes data copying across the FFI boundary - **Prebuilt Libraries** — Automatically downloads prebuilt `libzvec_c_api` from GitHub Releases; advanced users can override with `ZVEC_LIB_DIR` - **FTS with jieba Out of the Box** — `libzvec_c_api` embeds cppjieba's default dictionary and HMM model (LZ4-compressed), so the `jieba` FTS tokenizer needs no dictionary files on disk and no registration step +- **Grouped Vector Search** — `collection.group_by_query(&query)` returns the best matches per distinct value of a scalar field, so one dense category cannot sweep the whole result page ## Supported Platforms @@ -41,10 +42,11 @@ Safe, idiomatic Rust bindings for the [zvec](https://github.com/alibaba/zvec) ve > The dynamic library name varies by platform: `libzvec_c_api.dylib` (macOS), `libzvec_c_api.so` (Linux), `zvec_c_api.dll` (Windows). -> The Windows ARM64 prebuilt ships without RaBitQ and DiskANN: upstream zvec -> gates RaBitQ on x86_64 (its kernels are AVX2/AVX-512) and its Windows -> DiskANN backend on x86_64 too, so those index types are unavailable on that -> target. +> The Windows ARM64 prebuilt ships without RaBitQ: since zvec v0.8.1 upstream +> builds RaBitQ for ARM64 (NEON) on Linux, macOS, Android and iOS and for +> x86-64 (runtime AVX2/AVX-512 dispatch) on Linux, macOS and Windows, but not +> for Windows ARM64, so that index type is unavailable on that target. DiskANN +> is built for every prebuilt target listed above. ## Architecture @@ -229,6 +231,7 @@ Run any example with `cargo run --example `: | `basic` | End-to-end workflow: schema → insert → query → fetch → delete | | `schema_builder` | Various schema configurations: field types, index types, quantization | | `vector_search` | Vector query patterns: simple, builder, filter, output fields, HNSW params | +| `group_by_search` | Grouped vector search: best matches per category, filters, and a flat top-k comparison | | `crud_operations` | Full CRUD: insert, fetch, update, upsert, delete, stats, flush | | `config_logging` | Library configuration: memory limits, thread counts, logging | @@ -313,6 +316,7 @@ let schema = CollectionSchema::builder("name") | `collection.delete(&pks)` | Delete by primary keys | | `collection.delete_by_filter(filter)` | Delete documents matching a filter expression | | `collection.query(&query)` | Scalar filtering, full-text, or vector similarity search | +| `collection.group_by_query(&query)` | Vector search grouped by a scalar field (best matches per group) | | `collection.multi_query(&query)` | Multi-route search with RRF / weighted rerank | | `collection.fetch(&pks)` | Fetch by primary keys | | `collection.fetch_with_options(&pks, fields, include_vector)` | Fetch with output-field control | @@ -385,6 +389,31 @@ vector query (adding an FTS clause makes it a hybrid search), an FTS clause without a vector builds a keyword-only query, and neither builds a scalar query. `SearchQuery::fts(field, fts, topk)` is the direct keyword-only constructor. +### Grouped Search + +`GroupBySearchQuery` runs a vector search grouped by a scalar field: instead of one flat top-k +list it returns the best matches for every distinct value of that field. + +```rust +// The 2 best matches in each of up to 3 categories +let mut query = GroupBySearchQuery::new("embedding", "category", &query_vec, 3, 2)?; +query.set_output_fields(&["title", "price"])?; +query.set_filter("price <= 100.0")?; // applied before grouping +query.set_include_vector(true)?; // return the vectors as well + +for group in collection.group_by_query(&query)? { + println!("{}: {} hit(s)", group.group_by_value, group.docs.len()); + for doc in &group.docs { + println!(" {} scored {:.4}", doc.get_pk().unwrap_or(""), doc.get_score()); + } +} +``` + +Groups come back ordered by their best match, each holding at most `group_topk` documents in +similarity order. A group the filter empties out is dropped rather than returned empty, and a +query that matches nothing returns an empty `Vec` instead of an error. The documents are owned +by Rust, so the groups stay usable after the query — and the collection — are dropped. + ### Multi-Query (Hybrid Search) `MultiQuery` combines multiple sub-queries (dense vector, sparse vector, or FTS) with RRF or weighted reranking: @@ -434,15 +463,35 @@ Available distance metrics: `L2`, `Ip`, `Cosine`, `MipsL2`. | HNSW | `IndexParams::hnsw(metric, m, ef)` | Graph index (recommended) | | HNSW+Q | `IndexParams::hnsw_with_quantize(...)` | HNSW with quantization | | IVF | `IndexParams::ivf(metric, nlist, niters, soar)` | Inverted file index | +| IVF PQ | `IndexParams::ivf(...)` + `set_quantize_type(Pq)` + `set_quantizer_pq(&pq)` | IVF compressed with PQ4/PQ8 product quantization, optionally learned OPQ | | IVF RaBitQ | `IndexParams::ivf_rabitq(metric, nlist, total_bits, sample_count)` | IVF with RaBitQ quantization | | Flat | `IndexParams::flat(metric)` | Brute-force index | -| DiskANN | `IndexParams::diskann(metric, max_degree, list_size, pq_chunk_num)` | Disk-based graph index for large datasets (Linux x86_64/ARM64, macOS ARM64, Windows x86_64, 64-bit Android/iOS) | +| DiskANN | `IndexParams::diskann(metric, max_degree, list_size, pq_chunk_num)` | Disk-based graph index for large datasets (Linux x86_64/ARM64, macOS x86_64/ARM64, Windows x86_64/ARM64, 64-bit Android/iOS) | | Invert | `IndexParams::invert(range, wildcard)` | Scalar field index | | FTS | `IndexParams::fts(tokenizer, filters, extra)` | Full-text search index | With `Int8` / `Int4` quantization, `params.set_quantizer_enable_rotate(true)` randomly rotates -vectors before quantization to reduce the quantization error. The rotation matrix is stored with -the index and applied to query vectors at search time. +vectors before quantization to reduce the quantization error; with `Pq` the same flag selects a +learned OPQ rotation instead. The rotation matrix is stored with the index and applied to query +vectors at search time. + +An IVF index can also be compressed with product quantization. Set the quantize type to +`QuantizeType::Pq`, then describe the layout with `PqQuantizerParams`: + +```rust +let mut params = IndexParams::ivf(MetricType::L2, 100, 10, false)?; +params.set_quantize_type(QuantizeType::Pq)?; +params.set_quantizer_pq(&PqQuantizerParams { + num_chunk: 16, // the vector dimension must be divisible by it + num_bits: 8, // 4 or 8 bits per sub-quantizer code + fast_scan: false, // SIMD layout, requires num_bits == 4 + enable_rotate: true, // learned OPQ instead of plain PQ + ..Default::default() // opq_iter: 5, opq_pq_iter: 4 +})?; +``` + +PQ supports FP32 IVF fields with the `L2`, `Ip` and `Cosine` metrics; the dimension is checked +against `num_chunk` when the collection is created. ## Testing diff --git a/README_CN.md b/README_CN.md index 1b9a624..04da53b 100644 --- a/README_CN.md +++ b/README_CN.md @@ -17,6 +17,7 @@ - **尽可能零拷贝** — 最小化 FFI 边界的数据拷贝 - **预编译库支持** — 自动从 GitHub Releases 下载预编译的 `libzvec_c_api`；高阶用户可通过 `ZVEC_LIB_DIR` 覆盖 - **FTS jieba 开箱即用** — `libzvec_c_api` 内嵌 cppjieba 默认词表与 HMM 模型（LZ4 压缩），`jieba` 分词器不再依赖磁盘上的词表文件，也无需注册步骤 +- **分组向量检索** — `collection.group_by_query(&query)` 按标量字段的每个取值分别返回最佳匹配，避免单一类别占满整个结果页 ## 支持的平台 @@ -40,9 +41,10 @@ > 动态库文件名因平台而异：`libzvec_c_api.dylib`（macOS）、`libzvec_c_api.so`（Linux）、`zvec_c_api.dll`（Windows）。 -> Windows ARM64 预编译产物不含 RaBitQ 与 DiskANN：上游 zvec 将 RaBitQ 限定在 -> x86_64（其内核为 AVX2/AVX-512），Windows 的 DiskANN 后端同样只支持 x86_64， -> 因此该目标上这两类索引不可用。 +> Windows ARM64 预编译产物不含 RaBitQ：自 zvec v0.8.1 起，上游在 Linux、macOS、 +> Android、iOS 的 ARM64（NEON）以及 Linux、macOS、Windows 的 x86-64（运行时 +> AVX2/AVX-512 分发）上构建 RaBitQ，但不包含 Windows ARM64，因此该目标上这类 +> 索引不可用。DiskANN 在上表所有预编译目标上均可用。 ## 架构 @@ -221,6 +223,7 @@ fn main() -> zvec_rust::Result<()> { | `basic` | 端到端流程：schema → 插入 → 查询 → 获取 → 删除 | | `schema_builder` | 各种 Schema 配置：字段类型、索引类型、量化 | | `vector_search` | 向量查询模式：简单查询、Builder、过滤器、输出字段、HNSW 参数 | +| `group_by_search` | 分组向量检索：按类别取最佳匹配、过滤器，以及与扁平 top-k 的对比 | | `crud_operations` | 完整 CRUD：插入、获取、更新、upsert、删除、统计、刷新 | | `config_logging` | 库配置：内存限制、线程数、日志 | @@ -303,6 +306,7 @@ let schema = CollectionSchema::builder("name") | `collection.delete(&pks)` | 按主键删除 | | `collection.delete_by_filter(filter)` | 按过滤表达式批量删除 | | `collection.query(&query)` | 标量过滤、全文搜索或向量相似度搜索 | +| `collection.group_by_query(&query)` | 按标量字段分组的向量检索（每组取最佳匹配） | | `collection.multi_query(&query)` | 多路检索 + RRF / 加权重排 | | `collection.fetch(&pks)` | 按主键获取 | | `collection.fetch_with_options(&pks, fields, include_vector)` | 按主键获取并控制输出字段 | @@ -373,6 +377,30 @@ Builder 会根据已设置的字段推断查询类型：设置 `vector(..)` 构 两者都不设置则构建标量查询。纯关键词查询也可以直接用 `SearchQuery::fts(field, fts, topk)` 构造。 +### 分组检索 + +`GroupBySearchQuery` 按某个标量字段分组执行向量检索：返回的不是一个扁平的 top-k 列表， +而是该字段每个取值下的最佳匹配。 + +```rust +// 最多 3 个类别，每个类别取最佳的 2 条 +let mut query = GroupBySearchQuery::new("embedding", "category", &query_vec, 3, 2)?; +query.set_output_fields(&["title", "price"])?; +query.set_filter("price <= 100.0")?; // 分组前生效 +query.set_include_vector(true)?; // 同时返回向量 + +for group in collection.group_by_query(&query)? { + println!("{}: {} 条命中", group.group_by_value, group.docs.len()); + for doc in &group.docs { + println!(" {} 得分 {:.4}", doc.get_pk().unwrap_or(""), doc.get_score()); + } +} +``` + +分组按组内最佳匹配排序返回，每组最多 `group_topk` 条文档并按相似度排序。被过滤器清空的组 +会被直接丢弃，而不是返回一个空组；没有任何命中时返回空 `Vec` 而不是错误。文档所有权归 +Rust，因此查询对象乃至集合被 drop 之后，分组结果依然可用。 + ### 多路检索（混合搜索） `MultiQuery` 可以组合多个子查询（稠密向量、稀疏向量或全文检索），并使用 RRF 或加权策略重排： @@ -422,14 +450,34 @@ mq.set_rerank_weighted(&[0.7, 0.3])?; // 每个子查询的权重 | HNSW | `IndexParams::hnsw(metric, m, ef)` | 图索引（推荐） | | HNSW+量化 | `IndexParams::hnsw_with_quantize(...)` | 带量化的 HNSW | | IVF | `IndexParams::ivf(metric, nlist, niters, soar)` | 倒排文件索引 | +| IVF PQ | `IndexParams::ivf(...)` + `set_quantize_type(Pq)` + `set_quantizer_pq(&pq)` | 使用 PQ4/PQ8 乘积量化压缩的 IVF 索引，可选学习式 OPQ | | IVF RaBitQ | `IndexParams::ivf_rabitq(metric, nlist, total_bits, sample_count)` | 带 RaBitQ 量化的 IVF 索引 | | Flat | `IndexParams::flat(metric)` | 暴力搜索索引 | -| DiskANN | `IndexParams::diskann(metric, max_degree, list_size, pq_chunk_num)` | 基于磁盘的图索引，适合大规模数据集（Linux x86_64/ARM64、macOS ARM64、Windows x86_64、64 位 Android/iOS） | +| DiskANN | `IndexParams::diskann(metric, max_degree, list_size, pq_chunk_num)` | 基于磁盘的图索引，适合大规模数据集（Linux x86_64/ARM64、macOS x86_64/ARM64、Windows x86_64/ARM64、64 位 Android/iOS） | | Invert | `IndexParams::invert(range, wildcard)` | 标量字段索引 | | FTS | `IndexParams::fts(tokenizer, filters, extra)` | 全文检索索引 | 使用 `Int8` / `Int4` 量化时，`params.set_quantizer_enable_rotate(true)` 会在量化前对向量做随机旋转， -以降低量化误差。旋转矩阵随索引一起保存，并在检索时自动作用于查询向量。 +以降低量化误差；使用 `Pq` 时，同一开关改为启用学习式 OPQ 旋转。旋转矩阵随索引一起保存， +并在检索时自动作用于查询向量。 + +IVF 索引还可以用乘积量化（PQ）压缩：先把量化类型设为 `QuantizeType::Pq`，再用 +`PqQuantizerParams` 描述编码布局： + +```rust +let mut params = IndexParams::ivf(MetricType::L2, 100, 10, false)?; +params.set_quantize_type(QuantizeType::Pq)?; +params.set_quantizer_pq(&PqQuantizerParams { + num_chunk: 16, // 向量维度必须能被它整除 + num_bits: 8, // 每个子量化器 4 或 8 bit + fast_scan: false, // SIMD 布局，要求 num_bits == 4 + enable_rotate: true, // 使用学习式 OPQ 而非普通 PQ + ..Default::default() // opq_iter: 5, opq_pq_iter: 4 +})?; +``` + +PQ 支持使用 `L2`、`Ip`、`Cosine` 度量的 FP32 IVF 字段；创建集合时会校验向量维度与 +`num_chunk` 是否匹配。 ## 测试 diff --git a/zvec-sys/src/lib.rs b/zvec-sys/src/lib.rs index 32eaae3..9395a52 100644 --- a/zvec-sys/src/lib.rs +++ b/zvec-sys/src/lib.rs @@ -349,6 +349,8 @@ pub const ZVEC_QUANTIZE_TYPE_INT8: zvec_quantize_type_t = 2; pub const ZVEC_QUANTIZE_TYPE_INT4: zvec_quantize_type_t = 3; /// RaBitQ quantization. pub const ZVEC_QUANTIZE_TYPE_RABITQ: zvec_quantize_type_t = 4; +/// Product quantization (PQ), optionally with learned OPQ rotation. +pub const ZVEC_QUANTIZE_TYPE_PQ: zvec_quantize_type_t = 8; // ============================================================================= // Log Level Constants @@ -462,6 +464,31 @@ pub struct zvec_write_result_t { pub message: *const c_char, } +/// Product quantization (PQ) configuration of an IVF index. +/// +/// The C default initializer `ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT` expands to +/// `{8, 8, false, false, 5, 4}`. +#[repr(C)] +pub struct zvec_pq_quantizer_params_t { + pub num_chunk: c_int, + pub num_bits: c_int, + pub enable_rotate: bool, + pub fast_scan: bool, + pub opq_iter: u32, + pub opq_pq_iter: u32, +} + +/// One group of grouped vector search results. +/// +/// Every member is owned by the result array it came from: release the whole +/// array with `zvec_group_results_free` and never free a member on its own. +#[repr(C)] +pub struct zvec_group_result_t { + pub group_by_value: *mut c_char, + pub docs: *mut *mut zvec_doc_t, + pub doc_count: usize, +} + /// Detailed error information including source location. #[repr(C)] pub struct zvec_error_details_t { @@ -664,6 +691,14 @@ extern "C" { pub fn zvec_index_params_get_quantizer_enable_rotate( params: *const zvec_index_params_t, ) -> bool; + pub fn zvec_index_params_set_quantizer_pq( + params: *mut zvec_index_params_t, + pq: *const zvec_pq_quantizer_params_t, + ) -> zvec_error_code_t; + pub fn zvec_index_params_get_quantizer_pq( + params: *const zvec_index_params_t, + pq: *mut zvec_pq_quantizer_params_t, + ) -> zvec_error_code_t; pub fn zvec_index_params_set_hnsw_params( params: *mut zvec_index_params_t, m: c_int, @@ -1087,6 +1122,13 @@ extern "C" { results: *mut *mut *mut zvec_doc_t, result_count: *mut usize, ) -> zvec_error_code_t; + pub fn zvec_collection_group_by_query( + collection: *const zvec_collection_t, + query: *const zvec_group_by_vector_query_t, + results: *mut *mut zvec_group_result_t, + result_count: *mut usize, + ) -> zvec_error_code_t; + pub fn zvec_group_results_free(results: *mut zvec_group_result_t, result_count: usize); pub fn zvec_collection_multi_query( collection: *const zvec_collection_t, query: *const zvec_multi_query_t, diff --git a/zvec/Cargo.toml b/zvec/Cargo.toml index 5b68078..114eded 100644 --- a/zvec/Cargo.toml +++ b/zvec/Cargo.toml @@ -46,6 +46,10 @@ path = "examples/crud_operations.rs" name = "config_logging" path = "examples/config_logging.rs" +[[example]] +name = "group_by_search" +path = "examples/group_by_search.rs" + [[bench]] name = "zvec_bench" -harness = false \ No newline at end of file +harness = false diff --git a/zvec/examples/group_by_search.rs b/zvec/examples/group_by_search.rs new file mode 100644 index 0000000..0737902 --- /dev/null +++ b/zvec/examples/group_by_search.rs @@ -0,0 +1,121 @@ +//! Grouped vector search: the best matches per distinct value of a scalar +//! field, instead of a single flat top-k list. +//! +//! A plain top-k search can be swept by one dominant category. Grouping on a +//! scalar field guarantees every group is represented, which is what a +//! "related items" shelf or a diversified result page needs. + +use zvec_rust::*; + +fn main() -> zvec_rust::Result<()> { + println!("zvec version: {}", version()); + initialize(None)?; + + let data_path = "./zvec_group_by_search_example_data"; + + let schema = CollectionSchema::builder("group_by_collection") + .add_field(FieldSchema::new("id", DataType::String, false, 0)?) + .add_field(FieldSchema::new("title", DataType::String, true, 0)?) + .add_field(FieldSchema::new("category", DataType::String, false, 0)?) + .add_field(FieldSchema::new("price", DataType::Float, true, 0)?) + .add_vector_field( + "embedding", + DataType::VectorFp32, + 4, + IndexParams::flat(MetricType::L2)?, + ) + .build()?; + + let collection = Collection::create_and_open(data_path, &schema, None)?; + println!("Collection created successfully!"); + + let categories = [ + "electronics", + "books", + "clothing", + "electronics", + "books", + "clothing", + "electronics", + "books", + "clothing", + ]; + + let mut docs = Vec::new(); + for (i, category) in categories.iter().enumerate() { + let mut doc = Doc::new()?; + let pk = format!("doc_{}", i); + doc.set_pk(&pk); + doc.add_string("id", &pk)?; + doc.add_string("title", &format!("Product {}", i))?; + doc.add_string("category", category)?; + doc.add_f32("price", (i as f32 + 1.0) * 10.0)?; + // Every document sits on a line, so similarity order follows `i`. + doc.add_vector_f32("embedding", &[(i as f32) * 0.1, 0.2, 0.3, 0.4])?; + docs.push(doc); + } + + let doc_refs: Vec<&Doc> = docs.iter().collect(); + collection.insert(&doc_refs)?; + collection.flush()?; + println!("Inserted {} sample documents", docs.len()); + + let query_vector = [0.0, 0.2, 0.3, 0.4]; + + // Example 1: the two best matches of each of the three categories + println!("\n=== Example 1: Grouped Search ==="); + let mut query = GroupBySearchQuery::new("embedding", "category", &query_vector, 3, 2)?; + query.set_output_fields(&["title", "category", "price"])?; + + let groups = collection.group_by_query(&query)?; + println!("{} group(s) returned:", groups.len()); + for group in &groups { + println!(" category '{}':", group.group_by_value); + for doc in &group.docs { + println!( + " pk={}, title='{}', price={:.2}, distance={:.4}", + doc.get_pk().unwrap_or(""), + doc.get_string("title")?.unwrap_or_default(), + doc.get_f32("price")?.unwrap_or(0.0), + doc.get_score() + ); + } + } + + // Example 2: filter before grouping, and ask for the vectors back + println!("\n=== Example 2: Grouped Search with a Filter ==="); + let mut filtered = GroupBySearchQuery::new("embedding", "category", &query_vector, 3, 1)?; + filtered.set_output_fields(&["title"])?; + filtered.set_filter("price <= 30.0")?; + filtered.set_include_vector(true)?; + + for group in collection.group_by_query(&filtered)? { + let doc = &group.docs[0]; + println!( + " category '{}': pk={}, vector={:?}", + group.group_by_value, + doc.get_pk().unwrap_or(""), + doc.get_vector_f32("embedding")? + ); + } + + // Example 3: what a flat top-3 would have returned instead + println!("\n=== Example 3: Flat Search for Comparison ==="); + let flat = SearchQuery::new("embedding", &query_vector, 3)?; + for doc in collection.query(&flat)? { + println!( + " pk={}, category='{}'", + doc.get_pk().unwrap_or(""), + doc.get_string("category")?.unwrap_or_default() + ); + } + println!(" -> the flat top-3 misses whole categories the grouped search keeps"); + + collection.close()?; + shutdown()?; + + let _ = std::fs::remove_dir_all(data_path); + println!("\nDone! All grouped search examples demonstrated successfully."); + + Ok(()) +} diff --git a/zvec/examples/schema_builder.rs b/zvec/examples/schema_builder.rs index 7d6fc4d..0f07768 100644 --- a/zvec/examples/schema_builder.rs +++ b/zvec/examples/schema_builder.rs @@ -148,6 +148,29 @@ fn main() -> zvec_rust::Result<()> { println!("Schema created with FP32, FP16, and INT8 vector types"); + // Example 9: IVF index with product quantization (PQ) + println!("\n=== Example 9: IVF Product Quantization ==="); + let mut pq_params = IndexParams::ivf(MetricType::L2, 100, 10, false)?; + pq_params.set_quantize_type(QuantizeType::Pq)?; + pq_params.set_quantizer_pq(&PqQuantizerParams { + num_chunk: 16, + num_bits: 8, + ..Default::default() + })?; + + let pq = pq_params.quantizer_pq()?; + println!( + "PQ configured: {} chunk(s) x {} bit(s), fast_scan={}, OPQ={} ({} outer, {} inner iterations)", + pq.num_chunk, pq.num_bits, pq.fast_scan, pq.enable_rotate, pq.opq_iter, pq.opq_pq_iter + ); + + let _pq_schema = CollectionSchema::builder("pq_collection") + .add_field(FieldSchema::new("id", DataType::String, false, 0)?) + .add_vector_field("embedding", DataType::VectorFp32, 128, pq_params) + .build()?; + + println!("Schema created with an IVF index compressed to 16 x 8-bit PQ codes"); + // Clean up shutdown()?; diff --git a/zvec/src/collection.rs b/zvec/src/collection.rs index 44a6c57..0b31b97 100644 --- a/zvec/src/collection.rs +++ b/zvec/src/collection.rs @@ -1,11 +1,12 @@ use std::ffi::CStr; +use std::fmt; use std::marker::PhantomData; use std::ptr; use crate::doc::Doc; use crate::error::{check_error, to_cstring, Error, ErrorCode, Result}; use crate::multi_query::MultiQuery; -use crate::query::SearchQuery; +use crate::query::{GroupBySearchQuery, SearchQuery}; use crate::schema::{CollectionSchema, FieldSchema, IndexParams}; /// Options for creating or opening a collection. @@ -108,6 +109,28 @@ pub struct WriteResult { pub results: Vec, } +/// One group of grouped vector search results. +/// +/// Returned by [`Collection::group_by_query`]. +pub struct SearchGroup { + /// String representation of the grouping value, as the engine renders it. + pub group_by_value: String, + /// Documents in this group, ordered by vector similarity. + pub docs: Vec, +} + +impl fmt::Debug for SearchGroup { + // Prints the grouping value and how many documents it matched. Hand-written + // because `Doc` has no `Debug` implementation; the documents themselves are + // reachable through `SearchGroup::docs`. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SearchGroup") + .field("group_by_value", &self.group_by_value) + .field("doc_count", &self.docs.len()) + .finish() + } +} + /// A zvec collection for storing and querying vector data. /// /// Collections are the primary data container in zvec. They hold documents @@ -335,6 +358,44 @@ impl Collection { Ok(docs) } + /// Performs a vector similarity search grouped by a scalar field. + /// + /// The results come back as one [`SearchGroup`] per distinct value of the + /// grouping field, each holding at most `group_topk` documents ordered by + /// similarity. At most `group_count` groups are returned, starting with the + /// group containing the best match. + /// + /// A query that matches nothing returns an empty vector rather than an + /// error. The groups stay valid after the query — and after the collection + /// — is dropped. + /// + /// ```no_run + /// # use zvec_rust::*; + /// # fn f(collection: &Collection) -> Result<()> { + /// let query = GroupBySearchQuery::new("embedding", "category", &[0.1, 0.2], 3, 5)?; + /// for group in collection.group_by_query(&query)? { + /// println!("{} matched {} doc(s)", group.group_by_value, group.docs.len()); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn group_by_query(&self, query: &GroupBySearchQuery) -> Result> { + let mut results: *mut zvec_rust_sys::zvec_group_result_t = ptr::null_mut(); + let mut result_count: usize = 0; + + check_error(unsafe { + zvec_rust_sys::zvec_collection_group_by_query( + self.handle, + query.handle, + &mut results, + &mut result_count, + ) + })?; + + let groups = unsafe { collect_group_results(results, result_count) }; + Ok(groups) + } + /// Performs a multi-query that combines several sub-queries with a rerank /// strategy (RRF or weighted). pub fn multi_query(&self, query: &MultiQuery) -> Result> { @@ -806,6 +867,67 @@ unsafe fn collect_docs(results: *mut *mut zvec_rust_sys::zvec_doc_t, count: usiz docs } +/// Collects grouped results from a C array into a `Vec`. +/// +/// The C API releases a `zvec_group_result_t` array only as a whole, through +/// `zvec_group_results_free`, and that array owns everything inside it: the +/// grouping values, each group's document pointer array, and the documents +/// themselves. Every document is therefore moved into an owning [`Doc`], its +/// pointer array is released here, and the group is handed back as an empty +/// shell — so `zvec_group_results_free` is left with only the grouping strings +/// and the outer array to free, and cannot destroy a document twice. +/// +/// # Safety +/// `results` must be null, or a `count`-element array returned by +/// `zvec_collection_group_by_query` that has not been freed yet. +unsafe fn collect_group_results( + results: *mut zvec_rust_sys::zvec_group_result_t, + count: usize, +) -> Vec { + if results.is_null() { + return Vec::new(); + } + + let mut groups = Vec::with_capacity(count); + for i in 0..count { + let group = &mut *results.add(i); + + let group_by_value = if group.group_by_value.is_null() { + String::new() + } else { + CStr::from_ptr(group.group_by_value) + .to_string_lossy() + .into_owned() + }; + + let mut docs = Vec::with_capacity(group.doc_count); + if !group.docs.is_null() { + for j in 0..group.doc_count { + let doc_ptr = *group.docs.add(j); + if !doc_ptr.is_null() { + // Take ownership: Rust will call zvec_doc_destroy on drop. + docs.push(Doc::from_raw(doc_ptr)); + } + } + // Free the pointer array only; the documents it listed are owned by + // the Doc wrappers above now. + zvec_rust_sys::zvec_free(group.docs as *mut std::os::raw::c_void); + group.docs = ptr::null_mut(); + group.doc_count = 0; + } + + groups.push(SearchGroup { + group_by_value, + docs, + }); + } + + // Releases every group_by_value string and the result array itself. + zvec_rust_sys::zvec_group_results_free(results, count); + + groups +} + #[cfg(test)] mod tests { use super::*; @@ -816,6 +938,20 @@ mod tests { assert!(docs.is_empty()); } + #[test] + fn collect_group_results_handles_null_pointer() { + let groups = unsafe { collect_group_results(ptr::null_mut(), 0) }; + assert!(groups.is_empty()); + } + + #[test] + fn collect_group_results_null_array_wins_over_a_nonzero_count() { + // The engine reports "no groups" as a null array and can leave a stale + // count next to it; the null must win instead of being dereferenced. + let groups = unsafe { collect_group_results(ptr::null_mut(), 99) }; + assert!(groups.is_empty()); + } + #[test] fn collect_docs_handles_zero_count() { // Even with a non-null pointer, zero count should return empty diff --git a/zvec/src/lib.rs b/zvec/src/lib.rs index 9ee3f8f..43d2c3d 100644 --- a/zvec/src/lib.rs +++ b/zvec/src/lib.rs @@ -47,7 +47,7 @@ pub mod types; pub use collection::{ Collection, CollectionOptions, CollectionStats, DocIterator, DocWriteResult, IndexStat, - WriteResult, + SearchGroup, WriteResult, }; pub use config::{ initialize, io_backend_description, io_backend_type, io_backend_type_name, is_initialized, @@ -60,7 +60,7 @@ pub use query::{ DiskannQueryParams, FlatQueryParams, Fts, FtsQueryParams, GroupBySearchQuery, HnswQueryParams, IvfQueryParams, IvfRabitqQueryParams, SearchQuery, SearchQueryBuilder, }; -pub use schema::{CollectionSchema, FieldSchema, IndexParams}; +pub use schema::{CollectionSchema, FieldSchema, IndexParams, PqQuantizerParams}; pub use types::{DataType, DocOperator, IndexType, LogLevel, MetricType, QuantizeType}; /// Re-export the raw FFI bindings for advanced users. @@ -75,7 +75,7 @@ pub mod prelude { pub use crate::{ initialize, is_initialized, version, Collection, CollectionOptions, CollectionSchema, CollectionStats, ConfigBuilder, DataType, Doc, DocIterator, Error, ErrorCode, FieldSchema, - IndexParams, IndexStat, MetricType, MultiQuery, QuantizeType, Result, SearchQuery, - SubQuery, WriteResult, + GroupBySearchQuery, IndexParams, IndexStat, MetricType, MultiQuery, PqQuantizerParams, + QuantizeType, Result, SearchGroup, SearchQuery, SubQuery, WriteResult, }; } diff --git a/zvec/src/schema.rs b/zvec/src/schema.rs index e8fbce1..e0c85cf 100644 --- a/zvec/src/schema.rs +++ b/zvec/src/schema.rs @@ -3,6 +3,44 @@ use std::ffi::CStr; use crate::error::{check_error, to_cstring, Error, ErrorCode, Result}; use crate::types::{DataType, IndexType, MetricType, QuantizeType}; +/// Product quantization (PQ) configuration of an IVF index. +/// +/// PQ splits every vector into `num_chunk` sub-vectors and encodes each one +/// with `num_bits` bits, trading a little recall for a much smaller index. +/// Pair it with [`QuantizeType::Pq`] and [`IndexParams::set_quantizer_pq`]. +/// +/// [`Default`] reproduces the C initializer `ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PqQuantizerParams { + /// Number of sub-quantizers (chunks). The vector dimension must be + /// divisible by it, and it must be positive. + pub num_chunk: i32, + /// Bits per sub-quantizer code: `4` or `8`. + pub num_bits: i32, + /// Learn an OPQ rotation instead of quantizing the raw vectors. Both + /// iteration counts must then be non-zero. + pub enable_rotate: bool, + /// Use the FastScan SIMD layout. Requires `num_bits == 4`. + pub fast_scan: bool, + /// Outer OPQ iterations. + pub opq_iter: u32, + /// Inner PQ iterations performed per OPQ step. + pub opq_pq_iter: u32, +} + +impl Default for PqQuantizerParams { + fn default() -> Self { + PqQuantizerParams { + num_chunk: 8, + num_bits: 8, + enable_rotate: false, + fast_scan: false, + opq_iter: 5, + opq_pq_iter: 4, + } + } +} + /// Parameters for configuring an index on a field. pub struct IndexParams { pub(crate) handle: *mut zvec_rust_sys::zvec_index_params_t, @@ -312,10 +350,15 @@ impl IndexParams { /// Enables or disables random rotation before quantization. /// - /// Rotating vectors before `Int8` / `Int4` quantization reduces the - /// quantization error. The rotation matrix is stored with the index and - /// applied to query vectors at search time, so it has no effect with other - /// quantize types. + /// With `Int8` / `Int4` the vectors are randomly rotated before + /// quantization to reduce the quantization error; with `Pq` a learned OPQ + /// rotation is used instead. The rotation matrix is stored with the index + /// and applied to query vectors at search time, so it has no effect with + /// other quantize types. + /// + /// This is the same flag as [`PqQuantizerParams::enable_rotate`]: + /// [`IndexParams::set_quantizer_pq`] overwrites it, so call it after that + /// method to keep the value set here. /// /// Only vector index params accept this option; scalar indexes (`invert`, /// `fts`) return [`ErrorCode::InvalidArgument`]. @@ -324,6 +367,69 @@ impl IndexParams { zvec_rust_sys::zvec_index_params_set_quantizer_enable_rotate(self.handle, enable_rotate) }) } + + /// Configures product quantization (PQ) for an IVF index. + /// + /// ```no_run + /// # use zvec_rust::*; + /// # fn f() -> Result<()> { + /// let mut params = IndexParams::ivf(MetricType::L2, 100, 10, false)?; + /// params.set_quantize_type(QuantizeType::Pq)?; + /// params.set_quantizer_pq(&PqQuantizerParams { + /// num_chunk: 4, + /// num_bits: 4, + /// fast_scan: true, + /// ..Default::default() + /// })?; + /// # Ok(()) + /// # } + /// ``` + /// + /// PQ supports FP32 IVF fields with the `L2`, `Ip` and `Cosine` metrics. + /// `num_bits` is `4` or `8`, `fast_scan` requires 4 bits, and the vector + /// dimension must be divisible by `num_chunk` — the schema is validated + /// when the collection is created. + /// + /// Only IVF index params accept PQ configuration; every other index type + /// returns [`ErrorCode::InvalidArgument`], as does an invalid combination + /// of fields. + pub fn set_quantizer_pq(&mut self, pq: &PqQuantizerParams) -> Result<()> { + let raw = zvec_rust_sys::zvec_pq_quantizer_params_t { + num_chunk: pq.num_chunk, + num_bits: pq.num_bits, + enable_rotate: pq.enable_rotate, + fast_scan: pq.fast_scan, + opq_iter: pq.opq_iter, + opq_pq_iter: pq.opq_pq_iter, + }; + check_error(unsafe { zvec_rust_sys::zvec_index_params_set_quantizer_pq(self.handle, &raw) }) + } + + /// Returns the product quantization (PQ) configuration of an IVF index. + /// + /// Only IVF index params carry PQ configuration; every other index type + /// returns [`ErrorCode::InvalidArgument`]. + pub fn quantizer_pq(&self) -> Result { + let mut raw = zvec_rust_sys::zvec_pq_quantizer_params_t { + num_chunk: 0, + num_bits: 0, + enable_rotate: false, + fast_scan: false, + opq_iter: 0, + opq_pq_iter: 0, + }; + check_error(unsafe { + zvec_rust_sys::zvec_index_params_get_quantizer_pq(self.handle, &mut raw) + })?; + Ok(PqQuantizerParams { + num_chunk: raw.num_chunk, + num_bits: raw.num_bits, + enable_rotate: raw.enable_rotate, + fast_scan: raw.fast_scan, + opq_iter: raw.opq_iter, + opq_pq_iter: raw.opq_pq_iter, + }) + } } impl Drop for IndexParams { @@ -867,6 +973,97 @@ mod tests { assert_eq!(err.code, ErrorCode::InvalidArgument); } + #[test] + fn test_pq_quantizer_params_default_matches_c_macro() { + // ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT = {8, 8, false, false, 5, 4} + assert_eq!( + PqQuantizerParams::default(), + PqQuantizerParams { + num_chunk: 8, + num_bits: 8, + enable_rotate: false, + fast_scan: false, + opq_iter: 5, + opq_pq_iter: 4, + } + ); + } + + #[test] + fn test_index_params_quantizer_pq_roundtrip() { + let mut params = IndexParams::ivf(MetricType::L2, 100, 10, false).unwrap(); + params.set_quantize_type(QuantizeType::Pq).unwrap(); + let pq = PqQuantizerParams { + num_chunk: 4, + num_bits: 4, + enable_rotate: true, + fast_scan: true, + opq_iter: 3, + opq_pq_iter: 2, + }; + params.set_quantizer_pq(&pq).unwrap(); + + assert_eq!(params.quantize_type(), QuantizeType::Pq); + assert_eq!(params.quantizer_pq().unwrap(), pq); + } + + #[test] + fn test_index_params_quantizer_pq_rejects_invalid_combinations() { + let mut params = IndexParams::ivf(MetricType::L2, 100, 10, false).unwrap(); + + // FastScan is a 4-bit SIMD layout. + let err = params + .set_quantizer_pq(&PqQuantizerParams { + num_bits: 8, + fast_scan: true, + ..Default::default() + }) + .expect_err("fast_scan must require 4-bit codes"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + // num_bits is 4 or 8. + let err = params + .set_quantizer_pq(&PqQuantizerParams { + num_bits: 3, + ..Default::default() + }) + .expect_err("num_bits must be 4 or 8"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + // num_chunk must be positive. + let err = params + .set_quantizer_pq(&PqQuantizerParams { + num_chunk: 0, + ..Default::default() + }) + .expect_err("num_chunk must be positive"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + // OPQ needs positive iteration counts. + let err = params + .set_quantizer_pq(&PqQuantizerParams { + enable_rotate: true, + opq_iter: 0, + ..Default::default() + }) + .expect_err("enable_rotate must require positive OPQ iterations"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + } + + #[test] + fn test_index_params_quantizer_pq_rejects_non_ivf_index() { + let mut params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap(); + let err = params + .set_quantizer_pq(&PqQuantizerParams::default()) + .expect_err("PQ configuration is IVF-only"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + let err = params + .quantizer_pq() + .expect_err("PQ configuration is IVF-only"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + } + #[test] fn test_index_params_ivf_cosine() { let params = IndexParams::ivf(MetricType::Cosine, 100, 10, false).unwrap(); diff --git a/zvec/src/types.rs b/zvec/src/types.rs index a7dda61..1eec7f4 100644 --- a/zvec/src/types.rs +++ b/zvec/src/types.rs @@ -223,6 +223,10 @@ pub enum QuantizeType { Int4 = 3, /// RaBitQ quantization. Rabitq = 4, + /// Product quantization (PQ), optionally with learned OPQ rotation. + /// + /// Configure it with [`crate::schema::IndexParams::set_quantizer_pq`]. + Pq = 8, } impl From for QuantizeType { @@ -232,6 +236,7 @@ impl From for QuantizeType { 2 => QuantizeType::Int8, 3 => QuantizeType::Int4, 4 => QuantizeType::Rabitq, + 8 => QuantizeType::Pq, _ => QuantizeType::Undefined, } } @@ -540,11 +545,13 @@ mod tests { assert_eq!(QuantizeType::from(2), QuantizeType::Int8); assert_eq!(QuantizeType::from(3), QuantizeType::Int4); assert_eq!(QuantizeType::from(4), QuantizeType::Rabitq); + assert_eq!(QuantizeType::from(8), QuantizeType::Pq); } #[test] fn quantize_type_from_u32_unknown() { assert_eq!(QuantizeType::from(5), QuantizeType::Undefined); + assert_eq!(QuantizeType::from(7), QuantizeType::Undefined); assert_eq!(QuantizeType::from(99), QuantizeType::Undefined); } @@ -556,6 +563,7 @@ mod tests { QuantizeType::Int8, QuantizeType::Int4, QuantizeType::Rabitq, + QuantizeType::Pq, ]; for qt in all { let numeric: u32 = qt.into(); @@ -568,6 +576,7 @@ mod tests { fn quantize_type_display() { assert_eq!(QuantizeType::Fp16.to_string(), "Fp16"); assert_eq!(QuantizeType::Int4.to_string(), "Int4"); + assert_eq!(QuantizeType::Pq.to_string(), "Pq"); } // ========================================================================= diff --git a/zvec/tests/group_by_query_test.rs b/zvec/tests/group_by_query_test.rs new file mode 100644 index 0000000..2c089db --- /dev/null +++ b/zvec/tests/group_by_query_test.rs @@ -0,0 +1,201 @@ +//! Coverage for grouped vector search, which the C API exposes as of zvec +//! v0.8.1 through `zvec_collection_group_by_query`. + +use std::path::Path; +use std::sync::Once; + +use zvec_rust::{ + initialize, Collection, CollectionSchema, DataType, Doc, ErrorCode, FieldSchema, + FlatQueryParams, GroupBySearchQuery, IndexParams, MetricType, SearchGroup, +}; + +static INIT: Once = Once::new(); + +fn ensure_initialized() { + INIT.call_once(|| { + initialize(None).expect("failed to initialize zvec"); + }); +} + +/// Opens a collection shaped for grouped search: two scalar fields to group and +/// filter on, plus a brute-force vector index so the expected order is exact. +fn open_grouped_collection(path: &Path) -> Collection { + let schema = CollectionSchema::builder("group_by_collection") + .add_field(FieldSchema::new("id", DataType::Int64, false, 0).expect("id field")) + .add_field(FieldSchema::new("group_id", DataType::Int64, false, 0).expect("group_id field")) + .add_vector_field( + "embedding", + DataType::VectorFp32, + 4, + IndexParams::flat(MetricType::L2).expect("flat index params"), + ) + .build() + .expect("grouped search schema"); + + Collection::create_and_open(path.to_str().expect("database path"), &schema, None) + .expect("create grouped search collection") +} + +/// Inserts nine documents in three groups of three, placed on the x axis so +/// that the L2 distance of `doc_` to the origin is `id * id`. +fn insert_grouped_docs(collection: &Collection) { + let mut docs = Vec::new(); + for id in 0..9i64 { + let mut doc = Doc::new().expect("document"); + doc.set_pk(&format!("doc_{}", id)); + doc.add_i64("id", id).expect("id value"); + doc.add_i64("group_id", id / 3).expect("group_id value"); + doc.add_vector_f32("embedding", &[id as f32, 0.0, 0.0, 0.0]) + .expect("vector value"); + docs.push(doc); + } + + let refs: Vec<&Doc> = docs.iter().collect(); + let result = collection.insert(&refs).expect("insert documents"); + assert_eq!(result.success_count, 9); + assert_eq!(result.error_count, 0); + collection.flush().expect("persist documents"); +} + +/// Asserts that group `index` holds the `topk` closest documents of +/// `group_id`, in similarity order, carrying exactly the requested projection. +fn assert_group(groups: &[SearchGroup], index: usize, group_id: i64, topk: usize, vectors: bool) { + let group = groups + .get(index) + .unwrap_or_else(|| panic!("expected a group at index {} of {}", index, groups.len())); + + assert_eq!(group.group_by_value, group_id.to_string()); + assert_eq!(group.docs.len(), topk); + + for (offset, doc) in group.docs.iter().enumerate() { + // Documents are ordered by similarity, so with this data set the offsets + // walk the group from its closest to its farthest member. + let id = group_id * 3 + offset as i64; + assert_eq!(doc.get_pk(), Some(format!("doc_{}", id).as_str())); + assert_eq!(doc.get_i64("id").expect("id lookup"), Some(id)); + assert_eq!(doc.has_field("embedding"), vectors); + // The projection asked for "id" only, so the grouping field itself is + // not part of the returned document. + assert!(!doc.has_field("group_id")); + assert!( + (doc.get_score() - (id * id) as f32).abs() < 1e-3, + "doc_{} scored {}, expected the L2 distance {}", + id, + doc.get_score(), + id * id + ); + } +} + +#[test] +fn group_by_query_returns_the_best_documents_of_each_group() { + ensure_initialized(); + let temporary = tempfile::tempdir().expect("temporary directory"); + let collection = open_grouped_collection(&temporary.path().join("grouped")); + insert_grouped_docs(&collection); + + let mut query = + GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 2, 2).expect("grouped query"); + query.set_output_fields(&["id"]).expect("output fields"); + // The query takes ownership of the index specific search parameters. + query + .set_flat_params(FlatQueryParams::new(false, 10.0)) + .expect("flat query params"); + + let groups = collection.group_by_query(&query).expect("grouped search"); + assert_eq!(groups.len(), 2, "group_count caps the number of groups"); + assert_group(&groups, 0, 0, 2, false); + assert_group(&groups, 1, 1, 2, false); +} + +#[test] +fn group_by_query_honours_filter_group_count_and_include_vector() { + ensure_initialized(); + let temporary = tempfile::tempdir().expect("temporary directory"); + let collection = open_grouped_collection(&temporary.path().join("filtered")); + insert_grouped_docs(&collection); + + let mut query = + GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 3, 1).expect("grouped query"); + query.set_output_fields(&["id"]).expect("output fields"); + query.set_include_vector(true).expect("include vectors"); + query.set_filter("id >= 3").expect("filter"); + + let groups = collection.group_by_query(&query).expect("grouped search"); + // Group 0 holds ids 0..3, so the filter drops the whole group rather than + // returning it empty. + assert_eq!(groups.len(), 2); + assert_group(&groups, 0, 1, 1, true); + assert_group(&groups, 1, 2, 1, true); + + // A filter that matches nothing is an empty result, not an error. + query.set_filter("id < 0").expect("filter"); + let groups = collection.group_by_query(&query).expect("grouped search"); + assert!(groups.is_empty()); +} + +#[test] +fn group_by_query_on_an_empty_collection_returns_no_groups() { + ensure_initialized(); + let temporary = tempfile::tempdir().expect("temporary directory"); + let collection = open_grouped_collection(&temporary.path().join("empty")); + + let query = + GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 2, 3).expect("grouped query"); + let groups = collection.group_by_query(&query).expect("grouped search"); + assert!(groups.is_empty()); +} + +#[test] +fn group_by_query_rejects_invalid_parameters() { + ensure_initialized(); + let temporary = tempfile::tempdir().expect("temporary directory"); + let collection = open_grouped_collection(&temporary.path().join("invalid")); + insert_grouped_docs(&collection); + + let zero_groups = + GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 0, 2).expect("grouped query"); + let err = collection + .group_by_query(&zero_groups) + .expect_err("group_count = 0 must be rejected"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + let zero_topk = + GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 2, 0).expect("grouped query"); + let err = collection + .group_by_query(&zero_topk) + .expect_err("group_topk = 0 must be rejected"); + assert_eq!(err.code, ErrorCode::InvalidArgument); + + let missing_field = + GroupBySearchQuery::new("embedding", "missing", &[0.0; 4], 2, 2).expect("grouped query"); + let err = collection + .group_by_query(&missing_field) + .expect_err("an unknown grouping field must be rejected"); + assert_eq!(err.code, ErrorCode::InvalidArgument); +} + +#[test] +fn grouped_results_outlive_the_query_and_the_collection() { + ensure_initialized(); + let temporary = tempfile::tempdir().expect("temporary directory"); + + // The engine documents grouped results as valid after the query and the + // collection are gone: the safe API moves every document into an owning + // `Doc`, so dropping both here must not invalidate the groups. + let groups = { + let collection = open_grouped_collection(&temporary.path().join("outlives")); + insert_grouped_docs(&collection); + + let mut query = GroupBySearchQuery::new("embedding", "group_id", &[0.0; 4], 2, 2) + .expect("grouped query"); + query.set_output_fields(&["id"]).expect("output fields"); + + let groups = collection.group_by_query(&query).expect("grouped search"); + assert_eq!(groups.len(), 2); + groups + }; + + assert_group(&groups, 0, 0, 2, false); + assert_group(&groups, 1, 1, 2, false); +} diff --git a/zvec/tests/integration_test.rs b/zvec/tests/integration_test.rs index fd0a907..f45df9b 100644 --- a/zvec/tests/integration_test.rs +++ b/zvec/tests/integration_test.rs @@ -1396,10 +1396,10 @@ fn test_diskann_collection_end_to_end() { let collection = match Collection::create_and_open(dir.to_str().unwrap(), &schema, None) { Ok(c) => c, - // DiskANN is unavailable on some platforms (upstream supports Linux - // x86_64/ARM64, macOS ARM64, Windows x86_64 and 64-bit Android/iOS). - // Where it is missing the engine reports NotSupported; skip the - // end-to-end flow gracefully. + // Since zvec v0.8.1 DiskANN is built for Linux, macOS and Windows on + // both x86_64 and ARM64, plus 64-bit Android/iOS — every target this + // SDK ships a prebuilt for. Keep the skip as a safety net: anywhere + // else the engine reports NotSupported. Err(e) if e.code == ErrorCode::NotSupported => { eprintln!("skipping diskann end-to-end test: {}", e.message); return; @@ -1607,8 +1607,10 @@ fn test_ivf_rabitq_collection_end_to_end() { let collection = match Collection::create_and_open(dir.to_str().unwrap(), &schema, None) { Ok(c) => c, - // RaBitQ is only supported on Linux x86_64. On other platforms the - // engine reports NotSupported; skip the end-to-end flow gracefully. + // Since zvec v0.8.1 RaBitQ is built for ARM64 (NEON) on Linux, macOS, + // Android and iOS, and for x86-64 on Linux, macOS and Windows; only + // Windows ARM64 is still left out. Where it is missing the engine + // reports NotSupported; skip the end-to-end flow gracefully. Err(e) if e.code == ErrorCode::NotSupported => { eprintln!("skipping ivf rabitq end-to-end test: {}", e.message); return; From 1dab7367b94d43c56bd305832eecf4a5415b6a54 Mon Sep 17 00:00:00 2001 From: lc285652 Date: Fri, 13 Nov 2026 11:24:10 +0800 Subject: [PATCH 3/3] chore(release): 0.8.1 Bump the workspace from 0.8.0 to 0.8.1 to match the v0.8.1 tag the vendor submodule now tracks, along with the zvec-rust-sys path dependency requirement, Cargo.lock and the dependency snippets in both READMEs. Additive release: grouped vector search (Collection::group_by_query, SearchGroup) and IVF product quantization (PqQuantizerParams, IndexParams::set_quantizer_pq / quantizer_pq) are new surface, no C-API symbol was dropped and no wrapper was removed, so consumers keep compiling unchanged. One caveat for the semver purists: QuantizeType gained a Pq variant and the enum is not #[non_exhaustive], so a caller matching it exhaustively without a wildcard arm has to add one. The version still follows the upstream zvec release it tracks, as 0.8.0 did. release.yml and publish-crates.yml re-check these strings against the tag, so this commit has to be on main before v0.8.1 is tagged. --- Cargo.lock | 8 ++++---- Cargo.toml | 2 +- README.md | 6 +++--- README_CN.md | 6 +++--- zvec/Cargo.toml | 2 +- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 62b6488..d38bf90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -915,7 +915,7 @@ checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" [[package]] name = "zvec-rpath-smoke" -version = "0.8.0" +version = "0.8.1" dependencies = [ "zvec-rust", "zvec-rust-build", @@ -923,7 +923,7 @@ dependencies = [ [[package]] name = "zvec-rust" -version = "0.8.0" +version = "0.8.1" dependencies = [ "criterion", "rand", @@ -933,8 +933,8 @@ dependencies = [ [[package]] name = "zvec-rust-build" -version = "0.8.0" +version = "0.8.1" [[package]] name = "zvec-rust-sys" -version = "0.8.0" +version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index e558124..4fe1dcc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ exclude = ["fuzz"] resolver = "2" [workspace.package] -version = "0.8.0" +version = "0.8.1" edition = "2021" license = "Apache-2.0" repository = "https://github.com/zvec-ai/zvec-rust" diff --git a/README.md b/README.md index 31c25d3..7158977 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ Add `zvec-rust` to your `Cargo.toml`. The default `bundled` feature automaticall ```toml [dependencies] -zvec-rust = "0.8.0" +zvec-rust = "0.8.1" ``` `cargo run` / `cargo test` work out of the box because Cargo passes the @@ -83,10 +83,10 @@ the executable: ```toml [dependencies] -zvec-rust = "0.8.0" +zvec-rust = "0.8.1" [build-dependencies] -zvec-rust-build = "0.8.0" +zvec-rust-build = "0.8.1" ``` ```rust diff --git a/README_CN.md b/README_CN.md index 04da53b..6a787cc 100644 --- a/README_CN.md +++ b/README_CN.md @@ -68,7 +68,7 @@ Rust SDK 依赖 zvec C 库（`libzvec_c_api`）。可通过以下任一方式提 ```toml [dependencies] -zvec-rust = "0.8.0" +zvec-rust = "0.8.1" ``` `cargo run` / `cargo test` 可开箱即用，因为 Cargo 会把解析出的库目录传给链接器。 @@ -79,10 +79,10 @@ zvec-rust = "0.8.0" ```toml [dependencies] -zvec-rust = "0.8.0" +zvec-rust = "0.8.1" [build-dependencies] -zvec-rust-build = "0.8.0" +zvec-rust-build = "0.8.1" ``` ```rust diff --git a/zvec/Cargo.toml b/zvec/Cargo.toml index 114eded..2ffbb5e 100644 --- a/zvec/Cargo.toml +++ b/zvec/Cargo.toml @@ -12,7 +12,7 @@ categories = ["database", "api-bindings"] links = "zvec_rust" [dependencies] -zvec-rust-sys = { path = "../zvec-sys", version = "0.8.0" } +zvec-rust-sys = { path = "../zvec-sys", version = "0.8.1" } [features] default = ["bundled"]