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/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..7158977 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 @@ -68,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 @@ -81,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 @@ -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..6a787cc 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 在上表所有预编译目标上均可用。 ## 架构 @@ -66,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 会把解析出的库目录传给链接器。 @@ -77,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 @@ -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/vendor/zvec b/vendor/zvec index 8397071..30d662d 160000 --- a/vendor/zvec +++ b/vendor/zvec @@ -1 +1 @@ -Subproject commit 8397071d815cd842be2068339e00a94e69f069d1 +Subproject commit 30d662db4d2467937d843443074b34743bcc9baa 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..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"] @@ -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;