From ccee5e4ede8ca05f3d1167aca7d8a8c15a75cc9d Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 9 Jul 2026 18:49:02 +0800 Subject: [PATCH 1/4] Add IVF RQ query quantization support --- README.md | 20 +- STORAGE_FORMAT.md | 68 +- c/test_vindex.c | 62 +- core/benches/ann_bench.rs | 85 ++ core/src/index.rs | 176 ++- core/src/ivfrq.rs | 314 +++++ core/src/ivfrq_io.rs | 1115 +++++++++++++++++ core/src/lib.rs | 3 + core/src/rq.rs | 682 ++++++++++ core/tests/fixtures/ivf_rq_v1.hex | 8 + core/tests/storage_format_fixtures.rs | 42 + cpp/test_vindex.cpp | 69 +- ffi/src/lib.rs | 127 ++ include/paimon_vindex.hpp | 46 +- .../index/vector/VectorIndexNative.java | 31 + .../index/vector/VectorIndexReader.java | 53 +- .../index/vector/VectorIndexJavaApiTest.java | 11 + .../VectorIndexNativeValidationTest.java | 64 +- jni/src/lib.rs | 204 ++- python/paimon_vindex/__init__.py | 17 +- python/paimon_vindex/_ffi.py | 58 + python/tests/test_vindex.py | 40 + 22 files changed, 3229 insertions(+), 66 deletions(-) create mode 100644 core/src/ivfrq.rs create mode 100644 core/src/ivfrq_io.rs create mode 100644 core/src/rq.rs create mode 100644 core/tests/fixtures/ivf_rq_v1.hex diff --git a/README.md b/README.md index f04a838..28e6853 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ support multiple index families across Rust, C FFI, Java/JNI, and Python: | --- | --- | --- | | `IVF_FLAT` | IVF partitioning with uncompressed vectors. | Baseline recall and simple storage. | | `IVF_PQ` | IVF with product quantization and optional OPQ rotation. | Compact indexes with fast approximate scans. | +| `IVF_RQ` | IVF with 1-bit RaBitQ-style rotated residual quantization. | Very compact high-dimensional indexes with low training cost. | | `IVF_HNSW_FLAT` | IVF partitioning with an HNSW graph inside each list over raw vectors. | Higher recall within probed IVF lists. | | `IVF_HNSW_SQ` | IVF partitioning with per-list HNSW and scalar-quantized vectors. | HNSW-style search with smaller vector storage. | @@ -67,6 +68,14 @@ types: - `nprobe`: number of IVF lists to probe. - `ef_search`: optional HNSW search breadth for `IVF_HNSW_FLAT` and `IVF_HNSW_SQ`. A value of `0` uses the default. +- `rq.query_bits` / `query_bits`: optional `IVF_RQ` query quantization bits. + Supported values are `0`, `4`, and `8`; `0` keeps the default float-query + byte-LUT path. + +`IVF_RQ` currently writes 1-bit codes by default. Its v1 on-disk header already +records `num_bits`, `rotation_type`, `factor_layout`, and optional-section +flags for future multi-bit codes, but current readers intentionally reject +reserved multi-bit payloads until the corresponding scanner is implemented. Readers also expose an optional search warm-up API. Call `optimize_for_search` in Rust, C++, and Python, @@ -75,8 +84,8 @@ Readers also expose an optional search warm-up API. Call searches. The call builds in-memory search caches and does not change the serialized index format or search results. Currently `IVF_PQ` builds residual L2 precomputed tables for repeated PQ searches, `IVF_HNSW_SQ` builds SQ decode -LUTs for filtered-search SQ scan and fallback paths, and other index types -preload metadata. The `IVF_HNSW_SQ` LUTs are not expected to speed up the +LUTs for filtered-search SQ scan and fallback paths, and `IVF_RQ` plus other +index types preload metadata. The `IVF_HNSW_SQ` LUTs are not expected to speed up the normal unfiltered HNSW graph-search path. ### Rust @@ -111,6 +120,9 @@ let mut reader = VectorIndexReader::open(file)?; reader.optimize_for_search()?; let params = VectorSearchParams::with_ef_search(10, 16, 80); let (ids, distances) = reader.search(&query, params)?; + +let rq_params = VectorSearchParams::with_query_bits(10, 16, 4); +let (ids, distances) = reader.search(&query, rq_params)?; ``` Other Rust configs follow the same shape: @@ -300,6 +312,7 @@ writer.write(output) reader = VectorIndexReader(VectorIndexInput(index_bytes)) reader.optimize_for_search() ids, distances = reader.search(query, top_k=10, nprobe=16, ef_search=80) +ids, distances = reader.search(query, top_k=10, nprobe=16, query_bits=4) ``` The Python package is pure Python and uses `ctypes` to load @@ -332,7 +345,7 @@ Row IDs must be non-negative to map directly into `RoaringTreemap`'s `u64` domai ## ANN Benchmark The core crate includes an ANN-style benchmark for comparing Paimon's -`IVF_PQ`, `IVF_HNSW_FLAT`, and `IVF_HNSW_SQ` implementations. It reports build +`IVF_PQ`, `IVF_RQ`, `IVF_HNSW_FLAT`, and `IVF_HNSW_SQ` implementations. It reports build time, reader open/load time, first-query latency, batch query throughput, and serialized index size: @@ -345,6 +358,7 @@ The benchmark is configured with environment variables: ```bash ANN_N=100000 ANN_NQ=1000 ANN_D=128 ANN_K=10 ANN_NLIST=256 ANN_NPROBE=16 \ ANN_PQ_M=16 ANN_HNSW_M=20 ANN_HNSW_EF_CONSTRUCTION=150 ANN_HNSW_EF_SEARCH=80 \ +ANN_RQ_QUERY_BITS=0 \ cargo bench -p paimon-vindex-core --bench ann_bench -- --nocapture ``` diff --git a/STORAGE_FORMAT.md b/STORAGE_FORMAT.md index d66699c..a24b084 100644 --- a/STORAGE_FORMAT.md +++ b/STORAGE_FORMAT.md @@ -50,7 +50,7 @@ of the compatibility contract. ### Delta-Varint IDs -IVF-PQ and IVF-FLAT v1 sort each non-empty list by signed row id before writing. +IVF-PQ, IVF-FLAT, and IVF-RQ v1 sort each non-empty list by signed row id before writing. The first id is stored as `base_id: i64`. The id stream then stores one unsigned LEB128 varint per id, including the first id's zero delta. Each delta is computed with wrapping unsigned subtraction from the previous signed id. Readers reject a @@ -162,6 +162,72 @@ For each non-empty list payload: | `id_bytes` | bytes | delta-varint ids | | `vectors` | `count * d` `f32` | raw stored vectors | +## IVF-RQ v1 + +Magic: `IVRQ` (`0x49565251`). Version: `1`. Header size: 64 bytes. + +| Offset | Size | Type | Field | +| ---: | ---: | --- | --- | +| 0 | 4 | `u32` | magic | +| 4 | 4 | `u32` | version | +| 8 | 4 | `i32` | dimension `d` | +| 12 | 4 | `i32` | IVF list count `nlist` | +| 16 | 4 | `u32` | metric (`0=L2`, `1=InnerProduct`, `2=Cosine`) | +| 20 | 4 | `u32` | flags | +| 24 | 8 | `i64` | total vector count | +| 32 | 8 | `u64` | deterministic rotation seed | +| 40 | 4 | `u32` | deterministic rotation rounds | +| 44 | 4 | `i32` | bytes per primary sign-code plane, `ceil(d / 8)` | +| 48 | 4 | `u32` | RQ `num_bits`; currently `1`, reserved values `2`, `4`, and `8` | +| 52 | 4 | `u32` | `rotation_type`; currently `1` for deterministic Kac rotation | +| 56 | 4 | `u32` | `factor_layout`; currently `1` for three RaBitQ `f32` factors | +| 60 | 4 | `u32` | RQ format flags | + +Flags: + +| Bit | Meaning | +| ---: | --- | +| 0 | delta-varint ids are used; required in v1 | + +RQ format flags: + +| Bit | Meaning | +| ---: | --- | +| 0 | optional `ex_codes` section is present | +| 1 | optional `error_factor` section is present | + +Sections after the header: + +1. IVF coarse centroids: `nlist * d` `f32` values. +2. Offset table: `nlist` entries of `(offset: i64, count: i32, id_bytes_len: i32)`. +3. List payloads. + +For each non-empty list payload: + +| Field | Type | Notes | +| --- | --- | --- | +| `base_id` | `i64` | first sorted row id | +| `id_bytes_len` | `i32` | byte length of encoded id stream | +| `id_bytes` | bytes | delta-varint ids | +| `codes` | `count * ceil(d / 8)` bytes | primary RaBitQ sign-code plane over rotated residuals | +| `ex_codes` | optional bytes | reserved for additional code planes when `num_bits > 1`; present when RQ format flag bit 0 is set | +| `factors` | `count * 3` `f32` | per-vector `(residual_norm_sqr, vector_norm_sqr, dp_multiplier)` correction factors | +| `error_factor` | optional `count` `f32` | reserved per-vector refinement factor; present when RQ format flag bit 1 is set | + +The deterministic rotation is derived from `(d, rotation_seed, rotation_rounds)` +and applied to both indexed residuals and query residuals. The rotation parameters +are stored in the header so independently written files can use the same +distance estimator after being read by any v1 reader. + +Current writers emit `num_bits=1`, `rotation_type=1`, `factor_layout=1`, and no +optional RQ sections. Current readers validate these format fields and reject +reserved multi-bit or optional-section encodings until the corresponding scanner +is implemented. + +`query_bits` is a search-time IVF-RQ parameter and is not serialized. `0` uses +the default float-query byte-LUT estimator; `4` and `8` quantize the rotated +query residual into sign and magnitude bit planes for bitwise/popcount scanning. + ## IVF-HNSW-FLAT v1 Magic: `IHFL` (`0x4948464C`). Version: `1`. Header size: 64 bytes. diff --git a/c/test_vindex.c b/c/test_vindex.c index 99f3f63..6792139 100644 --- a/c/test_vindex.c +++ b/c/test_vindex.c @@ -50,7 +50,7 @@ struct MemBuffer { }; enum { - ROUNDTRIP_DIMENSION = 2, + ROUNDTRIP_DIMENSION = 8, ROUNDTRIP_NLIST = 4, ROUNDTRIP_PER_LIST = 128, ROUNDTRIP_VECTOR_COUNT = ROUNDTRIP_NLIST * ROUNDTRIP_PER_LIST, @@ -156,12 +156,20 @@ static void fill_roundtrip_data(float *data, int64_t *ids) { size_t cluster = i / ROUNDTRIP_PER_LIST; size_t local = i % ROUNDTRIP_PER_LIST; float center = (float)cluster * 20.0f; - data[i * ROUNDTRIP_DIMENSION] = center + (float)(local % 16) * 0.001f; - data[i * ROUNDTRIP_DIMENSION + 1] = center + (float)(local / 16) * 0.001f; + for (size_t dim = 0; dim < ROUNDTRIP_DIMENSION; dim++) { + data[i * ROUNDTRIP_DIMENSION + dim] = + center + (float)dim * 0.01f + (float)(local % 16) * 0.001f; + } ids[i] = cluster_base_id(cluster) + (int64_t)local; } } +static void fill_query(float *query, float center) { + for (size_t dim = 0; dim < ROUNDTRIP_DIMENSION; dim++) { + query[dim] = center + (float)dim * 0.01f; + } +} + static void assert_id_in_cluster(int64_t id, size_t cluster) { int64_t base = cluster_base_id(cluster); ASSERT_TRUE(id >= base); @@ -186,7 +194,7 @@ static void run_roundtrip( if (paimon_vindex_trainer_dimension(trainer, &dimension) != 0) { fail_ffi("trainer dimension failed"); } - ASSERT_EQ_I64(dimension, 2); + ASSERT_EQ_I64(dimension, ROUNDTRIP_DIMENSION); float *data = (float *)malloc(sizeof(float) * ROUNDTRIP_VECTOR_COUNT * ROUNDTRIP_DIMENSION); int64_t *ids = (int64_t *)malloc(sizeof(int64_t) * ROUNDTRIP_VECTOR_COUNT); @@ -240,7 +248,7 @@ static void run_roundtrip( } ASSERT_EQ_I64(metadata.index_type, expected_index_type); ASSERT_EQ_I64(metadata.metric, PAIMON_VINDEX_METRIC_L2); - ASSERT_EQ_I64(metadata.dimension, 2); + ASSERT_EQ_I64(metadata.dimension, ROUNDTRIP_DIMENSION); ASSERT_EQ_I64(metadata.nlist, 4); ASSERT_EQ_I64(metadata.total_vectors, ROUNDTRIP_VECTOR_COUNT); ASSERT_EQ_I64(metadata.pq_m, expected_pq_m); @@ -250,7 +258,8 @@ static void run_roundtrip( fail_ffi("reader optimize_for_search failed"); } - const float query[] = {0.0f, 0.0f}; + float query[ROUNDTRIP_DIMENSION]; + fill_query(query, 0.0f); int64_t result_ids[2] = {0}; float result_distances[2] = {0}; if (paimon_vindex_reader_search( @@ -259,8 +268,18 @@ static void run_roundtrip( } assert_id_in_cluster(result_ids[0], 0); ASSERT_TRUE(isfinite(result_distances[0])); + if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { + if (paimon_vindex_reader_search_with_query_bits( + reader, query, 2, 4, 16, 4, result_ids, result_distances, 2) != 0) { + fail_ffi("reader search with query bits failed"); + } + assert_id_in_cluster(result_ids[0], 0); + ASSERT_TRUE(isfinite(result_distances[0])); + } - const float queries[] = {0.0f, 0.0f, 20.0f, 20.0f}; + float queries[2 * ROUNDTRIP_DIMENSION]; + fill_query(queries, 0.0f); + fill_query(queries + ROUNDTRIP_DIMENSION, 20.0f); int64_t batch_ids[2] = {0}; float batch_distances[2] = {0}; if (paimon_vindex_reader_search_batch( @@ -269,6 +288,14 @@ static void run_roundtrip( } assert_id_in_cluster(batch_ids[0], 0); assert_id_in_cluster(batch_ids[1], 1); + if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { + if (paimon_vindex_reader_search_batch_with_query_bits( + reader, queries, 2, 1, 4, 16, 8, batch_ids, batch_distances, 2) != 0) { + fail_ffi("reader search batch with query bits failed"); + } + assert_id_in_cluster(batch_ids[0], 0); + assert_id_in_cluster(batch_ids[1], 1); + } paimon_vindex_reader_free(reader); free(buf.data); @@ -353,7 +380,7 @@ static void test_input_read_callback_error_propagates(void) { static void test_supported_index_roundtrips(void) { const char *flat_keys[] = {"index.type", "dimension", "nlist", "metric"}; - const char *flat_values[] = {"ivf_flat", "2", "4", "l2"}; + const char *flat_values[] = {"ivf_flat", "8", "4", "l2"}; run_roundtrip( "ivf_flat_roundtrip", flat_keys, @@ -364,18 +391,29 @@ static void test_supported_index_roundtrips(void) { 0); const char *pq_keys[] = {"index.type", "dimension", "nlist", "metric", "pq.m"}; - const char *pq_values[] = {"ivf_pq", "2", "4", "l2", "1"}; + const char *pq_values[] = {"ivf_pq", "8", "4", "l2", "4"}; run_roundtrip( "ivf_pq_roundtrip", pq_keys, pq_values, 5, PAIMON_VINDEX_INDEX_TYPE_IVF_PQ, - 1, + 4, + 0); + + const char *rq_keys[] = {"index.type", "dimension", "nlist", "metric"}; + const char *rq_values[] = {"ivf_rq", "8", "4", "l2"}; + run_roundtrip( + "ivf_rq_roundtrip", + rq_keys, + rq_values, + 4, + PAIMON_VINDEX_INDEX_TYPE_IVF_RQ, + 0, 0); const char *hnsw_flat_keys[] = {"index.type", "dimension", "nlist", "metric", "hnsw.m"}; - const char *hnsw_flat_values[] = {"ivf_hnsw_flat", "2", "4", "l2", "4"}; + const char *hnsw_flat_values[] = {"ivf_hnsw_flat", "8", "4", "l2", "4"}; run_roundtrip( "ivf_hnsw_flat_roundtrip", hnsw_flat_keys, @@ -386,7 +424,7 @@ static void test_supported_index_roundtrips(void) { 4); const char *hnsw_sq_keys[] = {"index.type", "dimension", "nlist", "metric", "hnsw.m"}; - const char *hnsw_sq_values[] = {"ivf_hnsw_sq", "2", "4", "l2", "4"}; + const char *hnsw_sq_values[] = {"ivf_hnsw_sq", "8", "4", "l2", "4"}; run_roundtrip( "ivf_hnsw_sq_roundtrip", hnsw_sq_keys, diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs index 9b6001d..4f060e8 100644 --- a/core/benches/ann_bench.rs +++ b/core/benches/ann_bench.rs @@ -27,6 +27,11 @@ use paimon_vindex_core::ivfhnswsq_io::{ search_batch_ivfhnswsq_reader, write_ivfhnswsq_index, IVFHNSWSQIndexReader, }; use paimon_vindex_core::ivfpq::{search_batch_reader, IVFPQIndex}; +use paimon_vindex_core::ivfrq::IVFRQIndex; +use paimon_vindex_core::ivfrq_io::{ + search_batch_ivfrq_reader_with_query_bits, write_ivfrq_index, IVFRQIndexReader, +}; +use paimon_vindex_core::rq::is_supported_query_bits; use std::env; use std::fs; use std::io; @@ -41,6 +46,7 @@ fn main() -> Result<(), Box> { println!("{}", CsvRow::header()); run_paimon_ivfpq(&cfg, &dataset, &ids, &workspace)?; + run_paimon_ivfrq(&cfg, &dataset, &ids, &workspace)?; run_paimon_ivfhnswflat(&cfg, &dataset, &ids, &workspace)?; run_paimon_ivfhnswsq(&cfg, &dataset, &ids, &workspace)?; @@ -59,6 +65,7 @@ struct Config { hnsw_m: usize, hnsw_ef_construction: usize, hnsw_ef_search: usize, + rq_query_bits: usize, clusters: usize, seed: u64, output_dir: PathBuf, @@ -76,6 +83,7 @@ impl Config { let hnsw_m = read_env("ANN_HNSW_M", 20)?; let hnsw_ef_construction = read_env("ANN_HNSW_EF_CONSTRUCTION", 150)?; let hnsw_ef_search = read_env("ANN_HNSW_EF_SEARCH", 80)?; + let rq_query_bits = read_env("ANN_RQ_QUERY_BITS", 0)?; let clusters = read_env("ANN_CLUSTERS", 32)?; let seed = read_env("ANN_SEED", 42)?; let output_dir = env::var("ANN_OUTPUT_DIR") @@ -96,6 +104,16 @@ impl Config { if d % pq_m != 0 { return Err(format!("ANN_D ({}) must be divisible by ANN_PQ_M ({})", d, pq_m).into()); } + if d % 8 != 0 { + return Err(format!("ANN_D ({}) must be divisible by 8 for IVF_RQ", d).into()); + } + if !is_supported_query_bits(rq_query_bits) { + return Err(format!( + "ANN_RQ_QUERY_BITS ({}) must be one of 0, 4, or 8", + rq_query_bits + ) + .into()); + } Ok(Self { n, @@ -108,6 +126,7 @@ impl Config { hnsw_m, hnsw_ef_construction, hnsw_ef_search, + rq_query_bits, clusters, seed, output_dir, @@ -269,6 +288,72 @@ fn run_paimon_ivfpq( Ok(()) } +fn run_paimon_ivfrq( + cfg: &Config, + dataset: &Dataset, + ids: &[i64], + workspace: &Path, +) -> io::Result<()> { + let path = workspace.join("paimon_ivfrq.index"); + let start = Instant::now(); + let mut index = IVFRQIndex::new(cfg.d, cfg.nlist, MetricType::L2); + index.train(&dataset.data, cfg.n); + index.add(&dataset.data, ids, cfg.n); + write_to_file(&path, |writer| write_ivfrq_index(&index, writer))?; + let build = start.elapsed(); + drop(index); + + let start = Instant::now(); + let file = fs::File::open(&path)?; + let mut reader = IVFRQIndexReader::open(file)?; + reader.ensure_loaded()?; + let read = start.elapsed(); + + let first_query = time_first_query(|| { + reader.search_with_query_bits( + &dataset.queries[..cfg.d], + cfg.k, + cfg.nprobe, + cfg.rq_query_bits, + ) + })?; + let search = time_search(|| { + search_batch_ivfrq_reader_with_query_bits( + &mut reader, + &dataset.queries, + cfg.nq, + cfg.k, + cfg.nprobe, + cfg.rq_query_bits, + ) + .map(|_| ()) + })?; + + print_row( + cfg, + "paimon", + "IVF_RQ", + None, + build, + read, + first_query, + search, + path.metadata()?.len(), + "index_bytes", + rq_note(cfg.rq_query_bits), + ); + Ok(()) +} + +fn rq_note(query_bits: usize) -> &'static str { + match query_bits { + 0 => "1bit query_bits=0", + 4 => "1bit query_bits=4", + 8 => "1bit query_bits=8", + _ => "1bit query_bits=invalid", + } +} + fn run_paimon_ivfhnswflat( cfg: &Config, dataset: &Dataset, diff --git a/core/src/index.rs b/core/src/index.rs index 25ea689..8dac92a 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -37,6 +37,12 @@ use crate::ivfpq::{ search_batch_reader, search_batch_reader_roaring_filter, search_with_reader, search_with_reader_roaring_filter, IVFPQIndex, }; +use crate::ivfrq::IVFRQIndex; +use crate::ivfrq_io::{ + search_batch_ivfrq_reader_roaring_filter_with_query_bits, + search_batch_ivfrq_reader_with_query_bits, write_ivfrq_index, IVFRQIndexReader, IVF_RQ_MAGIC, +}; +use crate::rq::{is_supported_query_bits, DEFAULT_RQ_QUERY_BITS}; use std::collections::{HashMap, HashSet}; use std::io; @@ -47,6 +53,7 @@ pub enum IndexType { IvfPq = 1, IvfHnswFlat = 2, IvfHnswSq = 3, + IvfRq = 4, } impl IndexType { @@ -56,6 +63,7 @@ impl IndexType { 1 => Some(Self::IvfPq), 2 => Some(Self::IvfHnswFlat), 3 => Some(Self::IvfHnswSq), + 4 => Some(Self::IvfRq), _ => None, } } @@ -66,6 +74,7 @@ impl IndexType { Self::IvfPq => "ivf_pq", Self::IvfHnswFlat => "ivf_hnsw_flat", Self::IvfHnswSq => "ivf_hnsw_sq", + Self::IvfRq => "ivf_rq", } } } @@ -84,6 +93,11 @@ pub enum VectorIndexConfig { metric: MetricType, use_opq: bool, }, + IvfRq { + dimension: usize, + nlist: usize, + metric: MetricType, + }, IvfHnswFlat { dimension: usize, nlist: usize, @@ -125,6 +139,11 @@ impl VectorIndexConfig { None => false, }, }, + IndexType::IvfRq => Self::IvfRq { + dimension: dimension?, + nlist: nlist?, + metric, + }, IndexType::IvfHnswFlat => Self::IvfHnswFlat { dimension: dimension?, nlist: nlist?, @@ -148,6 +167,7 @@ impl VectorIndexConfig { match self { Self::IvfFlat { .. } => IndexType::IvfFlat, Self::IvfPq { .. } => IndexType::IvfPq, + Self::IvfRq { .. } => IndexType::IvfRq, Self::IvfHnswFlat { .. } => IndexType::IvfHnswFlat, Self::IvfHnswSq { .. } => IndexType::IvfHnswSq, } @@ -157,6 +177,7 @@ impl VectorIndexConfig { match self { Self::IvfFlat { dimension, .. } | Self::IvfPq { dimension, .. } + | Self::IvfRq { dimension, .. } | Self::IvfHnswFlat { dimension, .. } | Self::IvfHnswSq { dimension, .. } => *dimension, } @@ -166,6 +187,7 @@ impl VectorIndexConfig { match self { Self::IvfFlat { nlist, .. } | Self::IvfPq { nlist, .. } + | Self::IvfRq { nlist, .. } | Self::IvfHnswFlat { nlist, .. } | Self::IvfHnswSq { nlist, .. } => *nlist, } @@ -250,10 +272,11 @@ fn parse_index_type_option(value: &str) -> io::Result { match value.trim() { "ivf_flat" => Ok(IndexType::IvfFlat), "ivf_pq" => Ok(IndexType::IvfPq), + "ivf_rq" => Ok(IndexType::IvfRq), "ivf_hnsw_flat" => Ok(IndexType::IvfHnswFlat), "ivf_hnsw_sq" => Ok(IndexType::IvfHnswSq), _ => Err(invalid_input(format!( - "unknown index.type '{}'; expected ivf_flat, ivf_pq, ivf_hnsw_flat, or ivf_hnsw_sq", + "unknown index.type '{}'; expected ivf_flat, ivf_pq, ivf_rq, ivf_hnsw_flat, or ivf_hnsw_sq", value ))), } @@ -294,6 +317,7 @@ pub struct VectorSearchParams { pub top_k: usize, pub nprobe: usize, pub ef_search: usize, + pub query_bits: usize, } impl VectorSearchParams { @@ -302,6 +326,7 @@ impl VectorSearchParams { top_k, nprobe, ef_search: 0, + query_bits: DEFAULT_RQ_QUERY_BITS, } } @@ -310,6 +335,30 @@ impl VectorSearchParams { top_k, nprobe, ef_search, + query_bits: DEFAULT_RQ_QUERY_BITS, + } + } + + pub fn with_query_bits(top_k: usize, nprobe: usize, query_bits: usize) -> Self { + Self { + top_k, + nprobe, + ef_search: 0, + query_bits, + } + } + + pub fn with_ef_search_and_query_bits( + top_k: usize, + nprobe: usize, + ef_search: usize, + query_bits: usize, + ) -> Self { + Self { + top_k, + nprobe, + ef_search, + query_bits, } } @@ -405,6 +454,7 @@ impl VectorIndexTraining { pub enum VectorIndexWriter { IvfFlat(IVFFlatIndex), IvfPq(IVFPQIndex), + IvfRq(IVFRQIndex), IvfHnswFlat(IVFHNSWFlatIndex), IvfHnswSq(IVFHNSWSQIndex), } @@ -429,6 +479,11 @@ impl VectorIndexWriter { metric, use_opq, } => Self::IvfPq(IVFPQIndex::new(dimension, nlist, m, metric, use_opq)), + VectorIndexConfig::IvfRq { + dimension, + nlist, + metric, + } => Self::IvfRq(IVFRQIndex::new(dimension, nlist, metric)), VectorIndexConfig::IvfHnswFlat { dimension, nlist, @@ -458,6 +513,7 @@ impl VectorIndexWriter { match self { Self::IvfFlat(_) => IndexType::IvfFlat, Self::IvfPq(_) => IndexType::IvfPq, + Self::IvfRq(_) => IndexType::IvfRq, Self::IvfHnswFlat(_) => IndexType::IvfHnswFlat, Self::IvfHnswSq(_) => IndexType::IvfHnswSq, } @@ -467,6 +523,7 @@ impl VectorIndexWriter { match self { Self::IvfFlat(index) => index.d, Self::IvfPq(index) => index.d, + Self::IvfRq(index) => index.d, Self::IvfHnswFlat(index) => index.flat.d, Self::IvfHnswSq(index) => index.d, } @@ -477,6 +534,7 @@ impl VectorIndexWriter { match self { Self::IvfFlat(index) => index.train(data, n), Self::IvfPq(index) => index.train(data, n), + Self::IvfRq(index) => index.train(data, n), Self::IvfHnswFlat(index) => index.train(data, n), Self::IvfHnswSq(index) => index.train(data, n), } @@ -494,6 +552,7 @@ impl VectorIndexWriter { match self { Self::IvfFlat(index) => index.add(data, ids, n), Self::IvfPq(index) => index.add(data, ids, n), + Self::IvfRq(index) => index.add(data, ids, n), Self::IvfHnswFlat(index) => index.add(data, ids, n), Self::IvfHnswSq(index) => index.add(data, ids, n), } @@ -504,6 +563,7 @@ impl VectorIndexWriter { match self { Self::IvfFlat(index) => write_ivfflat_index(index, out), Self::IvfPq(index) => write_index(index, out), + Self::IvfRq(index) => write_ivfrq_index(index, out), Self::IvfHnswFlat(index) => { index.build_graphs()?; write_ivfhnswflat_index(index, out) @@ -519,6 +579,7 @@ impl VectorIndexWriter { pub enum VectorIndexReader { IvfFlat(IVFFlatIndexReader), IvfPq(IVFPQIndexReader), + IvfRq(IVFRQIndexReader), IvfHnswFlat(IVFHNSWFlatIndexReader), IvfHnswSq(IVFHNSWSQIndexReader), } @@ -532,6 +593,7 @@ impl VectorIndexReader { match magic { IVFFLAT_MAGIC => Ok(Self::IvfFlat(IVFFlatIndexReader::open(reader)?)), MAGIC => Ok(Self::IvfPq(IVFPQIndexReader::open(reader)?)), + IVF_RQ_MAGIC => Ok(Self::IvfRq(IVFRQIndexReader::open(reader)?)), IVF_HNSW_FLAT_MAGIC => Ok(Self::IvfHnswFlat(IVFHNSWFlatIndexReader::open(reader)?)), IVF_HNSW_SQ_MAGIC => Ok(Self::IvfHnswSq(IVFHNSWSQIndexReader::open(reader)?)), _ => Err(io::Error::new( @@ -545,6 +607,7 @@ impl VectorIndexReader { match self { Self::IvfFlat(_) => IndexType::IvfFlat, Self::IvfPq(_) => IndexType::IvfPq, + Self::IvfRq(_) => IndexType::IvfRq, Self::IvfHnswFlat(_) => IndexType::IvfHnswFlat, Self::IvfHnswSq(_) => IndexType::IvfHnswSq, } @@ -570,6 +633,15 @@ impl VectorIndexReader { pq_m: Some(reader.m), hnsw: None, }, + Self::IvfRq(reader) => VectorIndexMetadata { + index_type: IndexType::IvfRq, + dimension: reader.d, + nlist: reader.nlist, + metric: reader.metric, + total_vectors: reader.total_vectors, + pq_m: None, + hnsw: None, + }, Self::IvfHnswFlat(reader) => VectorIndexMetadata { index_type: IndexType::IvfHnswFlat, dimension: reader.d, @@ -603,6 +675,7 @@ impl VectorIndexReader { match self { Self::IvfFlat(reader) => reader.ensure_loaded(), Self::IvfPq(reader) => reader.optimize_for_search(), + Self::IvfRq(reader) => reader.ensure_loaded(), Self::IvfHnswFlat(reader) => reader.ensure_loaded(), // IVF_HNSW_SQ warms SQ scan/fallback structures used by filtered // searches; normal unfiltered search primarily uses the HNSW graph. @@ -616,9 +689,13 @@ impl VectorIndexReader { params: VectorSearchParams, ) -> io::Result<(Vec, Vec)> { validate_query(query, self.dimension())?; + validate_query_bits_for_index(self.index_type(), params.query_bits)?; match self { Self::IvfFlat(reader) => reader.search(query, params.top_k, params.nprobe), Self::IvfPq(reader) => search_with_reader(reader, query, params.top_k, params.nprobe), + Self::IvfRq(reader) => { + reader.search_with_query_bits(query, params.top_k, params.nprobe, params.query_bits) + } Self::IvfHnswFlat(reader) => { reader.search(query, params.top_k, params.nprobe, params.hnsw_ef_search()) } @@ -635,6 +712,7 @@ impl VectorIndexReader { roaring_filter_bytes: &[u8], ) -> io::Result<(Vec, Vec)> { validate_query(query, self.dimension())?; + validate_query_bits_for_index(self.index_type(), params.query_bits)?; match self { Self::IvfFlat(reader) => reader.search_with_roaring_filter( query, @@ -649,6 +727,13 @@ impl VectorIndexReader { params.nprobe, roaring_filter_bytes, ), + Self::IvfRq(reader) => reader.search_with_roaring_filter_and_query_bits( + query, + params.top_k, + params.nprobe, + roaring_filter_bytes, + params.query_bits, + ), Self::IvfHnswFlat(reader) => reader.search_with_roaring_filter( query, params.top_k, @@ -673,6 +758,7 @@ impl VectorIndexReader { params: VectorSearchParams, ) -> io::Result<(Vec, Vec)> { validate_queries(queries, query_count, self.dimension())?; + validate_query_bits_for_index(self.index_type(), params.query_bits)?; match self { Self::IvfFlat(reader) => search_batch_ivfflat_reader( reader, @@ -684,6 +770,14 @@ impl VectorIndexReader { Self::IvfPq(reader) => { search_batch_reader(reader, queries, query_count, params.top_k, params.nprobe) } + Self::IvfRq(reader) => search_batch_ivfrq_reader_with_query_bits( + reader, + queries, + query_count, + params.top_k, + params.nprobe, + params.query_bits, + ), Self::IvfHnswFlat(reader) => search_batch_ivfhnswflat_reader( reader, queries, @@ -711,6 +805,7 @@ impl VectorIndexReader { roaring_filter_bytes: &[u8], ) -> io::Result<(Vec, Vec)> { validate_queries(queries, query_count, self.dimension())?; + validate_query_bits_for_index(self.index_type(), params.query_bits)?; match self { Self::IvfFlat(reader) => search_batch_ivfflat_reader_roaring_filter( reader, @@ -728,6 +823,15 @@ impl VectorIndexReader { params.nprobe, roaring_filter_bytes, ), + Self::IvfRq(reader) => search_batch_ivfrq_reader_roaring_filter_with_query_bits( + reader, + queries, + query_count, + params.top_k, + params.nprobe, + roaring_filter_bytes, + params.query_bits, + ), Self::IvfHnswFlat(reader) => search_batch_ivfhnswflat_reader_roaring_filter( reader, queries, @@ -763,6 +867,14 @@ fn validate_config(config: &VectorIndexConfig) -> io::Result<()> { )); } } + VectorIndexConfig::IvfRq { dimension, .. } => { + if !dimension.is_multiple_of(8) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("dimension {} must be divisible by 8 for IVF_RQ", dimension), + )); + } + } VectorIndexConfig::IvfHnswFlat { hnsw, .. } | VectorIndexConfig::IvfHnswSq { hnsw, .. } => { validate_hnsw_params(*hnsw)? } @@ -777,6 +889,25 @@ fn validate_hnsw_params(params: HnswBuildParams) -> io::Result<()> { validate_positive(params.max_level, "hnsw max_level") } +fn validate_query_bits_for_index(index_type: IndexType, query_bits: usize) -> io::Result<()> { + if query_bits == DEFAULT_RQ_QUERY_BITS { + return Ok(()); + } + if index_type != IndexType::IvfRq { + return Err(invalid_input(format!( + "query_bits is only supported by IVF_RQ, but index type is {}", + index_type.as_str() + ))); + } + if !is_supported_query_bits(query_bits) { + return Err(invalid_input(format!( + "invalid IVF_RQ query_bits {}; expected 0, 4, or 8", + query_bits + ))); + } + Ok(()) +} + fn validate_positive(value: usize, name: &str) -> io::Result<()> { if value == 0 { Err(invalid_input(format!("{} must be greater than 0", name))) @@ -969,6 +1100,11 @@ mod tests { metric: MetricType::L2, use_opq: false, }); + roundtrip(VectorIndexConfig::IvfRq { + dimension: 8, + nlist: 4, + metric: MetricType::L2, + }); roundtrip(VectorIndexConfig::IvfHnswFlat { dimension: 8, nlist: 4, @@ -998,6 +1134,11 @@ mod tests { metric: MetricType::L2, use_opq: false, }, + VectorIndexConfig::IvfRq { + dimension: 8, + nlist: 4, + metric: MetricType::L2, + }, VectorIndexConfig::IvfHnswFlat { dimension: 8, nlist: 4, @@ -1059,6 +1200,19 @@ mod tests { assert!(err.to_string().contains("must be divisible")); } + #[test] + fn unified_config_rejects_invalid_rq_dimension() { + let err = match VectorIndexTrainer::new(VectorIndexConfig::IvfRq { + dimension: 10, + nlist: 4, + metric: MetricType::L2, + }) { + Ok(_) => panic!("invalid RQ config should be rejected"), + Err(err) => err, + }; + assert!(err.to_string().contains("divisible by 8")); + } + fn options(values: &[(&str, &str)]) -> HashMap { values .iter() @@ -1096,6 +1250,26 @@ mod tests { _ => panic!("expected IVF PQ config"), } + match VectorIndexConfig::from_options(&options(&[ + ("index.type", "ivf_rq"), + ("dimension", "8"), + ("nlist", "4"), + ("metric", "cosine"), + ])) + .unwrap() + { + VectorIndexConfig::IvfRq { + dimension, + nlist, + metric, + } => { + assert_eq!(dimension, 8); + assert_eq!(nlist, 4); + assert_eq!(metric, MetricType::Cosine); + } + _ => panic!("expected IVF RQ config"), + } + match VectorIndexConfig::from_options(&options(&[ ("index.type", "ivf_hnsw_sq"), ("dimension", "8"), diff --git a/core/src/ivfrq.rs b/core/src/ivfrq.rs new file mode 100644 index 0000000..0051498 --- /dev/null +++ b/core/src/ivfrq.rs @@ -0,0 +1,314 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::distance::{fvec_madd, fvec_norm_l2sqr, preprocess_vectors, MetricType}; +use crate::ivfpq::RowIdFilter; +use crate::kmeans::{self, KMeansConfig}; +use crate::rq::{ + RQCodeFactors, RQRotation, RaBitQuantizer, DEFAULT_RQ_QUERY_BITS, DEFAULT_RQ_ROTATION_ROUNDS, + DEFAULT_RQ_ROTATION_SEED, RQ_BYTE_LUT_MIN_LIST_SIZE, +}; +use crate::topk::TopKHeap; + +pub struct IVFRQIndex { + pub d: usize, + pub nlist: usize, + pub metric: MetricType, + pub quantizer_centroids: Vec, + pub rotation_seed: u64, + pub rotation_rounds: u32, + pub ids: Vec>, + pub codes: Vec>, + pub factors: Vec>, + quantizer: RaBitQuantizer, + rotation: RQRotation, +} + +impl IVFRQIndex { + pub fn new(d: usize, nlist: usize, metric: MetricType) -> Self { + Self::with_rotation( + d, + nlist, + metric, + DEFAULT_RQ_ROTATION_SEED, + DEFAULT_RQ_ROTATION_ROUNDS, + ) + } + + pub fn with_rotation( + d: usize, + nlist: usize, + metric: MetricType, + rotation_seed: u64, + rotation_rounds: u32, + ) -> Self { + Self { + d, + nlist, + metric, + quantizer_centroids: Vec::new(), + rotation_seed, + rotation_rounds, + ids: vec![Vec::new(); nlist], + codes: vec![Vec::new(); nlist], + factors: vec![Vec::new(); nlist], + quantizer: RaBitQuantizer::new(d), + rotation: RQRotation::new(d, rotation_seed, rotation_rounds), + } + } + + pub fn train(&mut self, data: &[f32], n: usize) { + let processed = self.preprocess_vectors(data, n); + self.quantizer_centroids = + kmeans::kmeans_train(&KMeansConfig::default(), &processed, n, self.d, self.nlist); + } + + pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) { + let processed = self.preprocess_vectors(data, n); + let list_ids = kmeans::find_nearest_batch( + &processed, + n, + &self.quantizer_centroids, + self.nlist, + self.d, + ); + let code_size = self.code_size(); + let mut residual = vec![0.0f32; self.d]; + let mut code = vec![0u8; code_size]; + + for i in 0..n { + let list_id = list_ids[i]; + let vector = &processed[i * self.d..(i + 1) * self.d]; + self.write_rotated_residual(vector, list_id, &mut residual); + let factors = self + .quantizer + .encode(&residual, fvec_norm_l2sqr(vector), &mut code); + self.ids[list_id].push(ids[i]); + self.codes[list_id].extend_from_slice(&code); + self.factors[list_id].push(factors); + } + } + + pub fn total_vectors(&self) -> usize { + self.ids.iter().map(Vec::len).sum() + } + + pub fn code_size(&self) -> usize { + self.quantizer.code_size() + } + + pub fn search( + &self, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + result_distances: &mut [f32], + result_labels: &mut [i64], + ) { + self.search_with_filter( + queries, + nq, + k, + nprobe, + None, + DEFAULT_RQ_QUERY_BITS, + result_distances, + result_labels, + ); + } + + pub fn search_with_query_bits( + &self, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + query_bits: usize, + result_distances: &mut [f32], + result_labels: &mut [i64], + ) { + self.search_with_filter( + queries, + nq, + k, + nprobe, + None, + query_bits, + result_distances, + result_labels, + ); + } + + pub fn search_with_filter( + &self, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + query_bits: usize, + result_distances: &mut [f32], + result_labels: &mut [i64], + ) { + let processed_queries = self.preprocess_vectors(queries, nq); + let (all_probe_indices, _) = kmeans::find_topk_batch( + &processed_queries, + nq, + &self.quantizer_centroids, + self.nlist, + self.d, + nprobe, + ); + + for qi in 0..nq { + let query = &processed_queries[qi * self.d..(qi + 1) * self.d]; + let mut heap = TopKHeap::new(k); + for &list_id in &all_probe_indices[qi] { + self.scan_list(query, list_id, filter, query_bits, &mut heap); + } + + let sorted = heap.into_sorted(); + let out_base = qi * k; + for (i, &(dist, id)) in sorted.iter().enumerate() { + result_distances[out_base + i] = dist; + result_labels[out_base + i] = id; + } + for i in sorted.len()..k { + result_distances[out_base + i] = f32::MAX; + result_labels[out_base + i] = -1; + } + } + } + + pub(crate) fn preprocess_vectors(&self, data: &[f32], n: usize) -> Vec { + preprocess_vectors(data, n, self.d, self.metric) + } + + pub(crate) fn list_centroid(&self, list_id: usize) -> &[f32] { + &self.quantizer_centroids[list_id * self.d..(list_id + 1) * self.d] + } + + pub(crate) fn rotated_query_residual(&self, query: &[f32], list_id: usize) -> Vec { + let mut residual = vec![0.0f32; self.d]; + self.write_rotated_residual(query, list_id, &mut residual); + residual + } + + fn scan_list( + &self, + query: &[f32], + list_id: usize, + filter: Option<&dyn RowIdFilter>, + query_bits: usize, + heap: &mut TopKHeap, + ) { + let rotated_query_residual = self.rotated_query_residual(query, list_id); + let distance_context = self.quantizer.prepare_distance_context_with_query_bits( + rotated_query_residual, + query, + self.ids[list_id].len() >= RQ_BYTE_LUT_MIN_LIST_SIZE, + query_bits, + ); + let code_size = self.code_size(); + for (local_idx, &id) in self.ids[list_id].iter().enumerate() { + if filter.map(|f| !f.contains(id)).unwrap_or(false) { + continue; + } + let code = &self.codes[list_id][local_idx * code_size..(local_idx + 1) * code_size]; + let dist = self.quantizer.distance_to_code_prepared( + &distance_context, + code, + self.factors[list_id][local_idx], + self.metric, + ); + heap.push(dist, id); + } + } + + fn write_rotated_residual(&self, vector: &[f32], list_id: usize, out: &mut [f32]) { + fvec_madd(vector, self.list_centroid(list_id), -1.0, out); + self.rotation.apply(out); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ivfrq_recalls_query_vector() { + let d = 8; + let nlist = 4; + let n = 128; + let data: Vec = (0..n) + .flat_map(|i| { + let cluster = (i % nlist) as f32 * 100.0; + [cluster + i as f32 * 0.01, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + }) + .collect(); + let ids: Vec = (1000..1000 + n as i64).collect(); + + let mut index = IVFRQIndex::new(d, nlist, MetricType::L2); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut distances = vec![0.0; 5]; + let mut labels = vec![0; 5]; + index.search( + &data[7 * d..8 * d], + 1, + 5, + nlist, + &mut distances, + &mut labels, + ); + + assert_eq!(labels[0], ids[7]); + assert!(distances[0] <= 1e-4); + } + + #[test] + fn ivfrq_inner_product_recalls_query_vector() { + let d = 8; + let nlist = 1; + let n = 8; + let mut data = vec![0.0f32; n * d]; + for i in 0..n { + data[i * d + i] = 1.0; + } + let ids: Vec = (1000..1000 + n as i64).collect(); + + let mut index = IVFRQIndex::new(d, nlist, MetricType::InnerProduct); + index.train(&data, n); + index.add(&data, &ids, n); + + let query_id = 7; + let mut distances = vec![0.0; 5]; + let mut labels = vec![0; 5]; + index.search( + &data[query_id * d..(query_id + 1) * d], + 1, + 5, + nlist, + &mut distances, + &mut labels, + ); + + assert_eq!(labels[0], ids[query_id]); + } +} diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs new file mode 100644 index 0000000..925d434 --- /dev/null +++ b/core/src/ivfrq_io.rs @@ -0,0 +1,1115 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::distance::{fvec_madd, preprocess_vectors, MetricType}; +use crate::index_io_util::{ + decode_delta_varint_ids, encode_delta_varint_ids, validate_search_inputs, +}; +use crate::io::{PreadCursor, ReadRequest, SeekRead, SeekWrite}; +use crate::ivfpq::RowIdFilter; +use crate::ivfrq::IVFRQIndex; +use crate::kmeans; +use crate::rq::{ + is_supported_query_bits, RQCodeFactors, RQDistanceContext, RQRotation, RaBitQuantizer, + DEFAULT_RQ_QUERY_BITS, RQ_BYTE_LUT_MIN_LIST_SIZE, +}; +use crate::topk::TopKHeap; +use roaring::RoaringTreemap; +use std::io; + +pub const IVF_RQ_MAGIC: u32 = 0x49565251; // "IVRQ" +pub const IVF_RQ_VERSION: u32 = 1; +pub const IVF_RQ_HEADER_SIZE: usize = 64; + +const FLAG_DELTA_IDS: u32 = 1 << 0; +const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS; +const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS; +const FACTOR_BYTES: usize = 12; + +pub const IVF_RQ_NUM_BITS_ONE: u32 = 1; +pub const IVF_RQ_ROTATION_TYPE_KAC: u32 = 1; +pub const IVF_RQ_FACTOR_LAYOUT_RABITQ_V1: u32 = 1; + +const FORMAT_FLAG_EX_CODES_PRESENT: u32 = 1 << 0; +const FORMAT_FLAG_ERROR_FACTOR_PRESENT: u32 = 1 << 1; +const SUPPORTED_FORMAT_FLAGS: u32 = FORMAT_FLAG_EX_CODES_PRESENT | FORMAT_FLAG_ERROR_FACTOR_PRESENT; +const CURRENT_FORMAT_FLAGS: u32 = 0; + +pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Result<()> { + validate_index_shape(index)?; + let d = index.d; + let nlist = index.nlist; + let code_size = index.code_size(); + let total_vectors = index.ids.iter().try_fold(0i64, |sum, ids| { + let count = usize_to_i64(ids.len(), "total vector count")?; + sum.checked_add(count).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "total vector count exceeds i64 length limit", + ) + }) + })?; + + let mut sorted_lists: Vec<(Vec, Vec, Vec, Vec)> = + Vec::with_capacity(nlist); + for list_id in 0..nlist { + let count = index.ids[list_id].len(); + if count == 0 { + sorted_lists.push((Vec::new(), Vec::new(), Vec::new(), Vec::new())); + continue; + } + + let mut order: Vec = (0..count).collect(); + order.sort_by_key(|&idx| index.ids[list_id][idx]); + + let sorted_ids: Vec = order.iter().map(|&idx| index.ids[list_id][idx]).collect(); + let mut sorted_codes = Vec::with_capacity(count * code_size); + let mut sorted_factors = Vec::with_capacity(count); + for idx in order { + sorted_codes + .extend_from_slice(&index.codes[list_id][idx * code_size..(idx + 1) * code_size]); + sorted_factors.push(index.factors[list_id][idx]); + } + let (_, id_bytes) = encode_delta_varint_ids(&sorted_ids); + sorted_lists.push((sorted_ids, id_bytes, sorted_codes, sorted_factors)); + } + + write_u32_le(out, IVF_RQ_MAGIC)?; + write_u32_le(out, IVF_RQ_VERSION)?; + write_i32_le(out, usize_to_i32(d, "dimension")?)?; + write_i32_le(out, usize_to_i32(nlist, "nlist")?)?; + write_u32_le(out, index.metric as u32)?; + write_u32_le(out, REQUIRED_FLAGS)?; + write_i64_le(out, total_vectors)?; + write_u64_le(out, index.rotation_seed)?; + write_u32_le(out, index.rotation_rounds)?; + write_i32_le(out, usize_to_i32(code_size, "code size")?)?; + write_u32_le(out, IVF_RQ_NUM_BITS_ONE)?; + write_u32_le(out, IVF_RQ_ROTATION_TYPE_KAC)?; + write_u32_le(out, IVF_RQ_FACTOR_LAYOUT_RABITQ_V1)?; + write_u32_le(out, CURRENT_FORMAT_FLAGS)?; + + write_f32_slice(out, &index.quantizer_centroids)?; + + let offset_table_size = nlist.checked_mul(16).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "IVF-RQ offset table size overflow", + ) + })?; + let data_start = out + .pos() + .checked_add(offset_table_size as u64) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "IVF-RQ data start offset overflow", + ) + })?; + let mut list_offsets = vec![0i64; nlist]; + let mut list_counts = vec![0i32; nlist]; + let mut list_id_bytes_lens = vec![0i32; nlist]; + let mut current_offset = data_start; + + for list_id in 0..nlist { + list_offsets[list_id] = u64_to_i64(current_offset, "list offset")?; + let count = sorted_lists[list_id].0.len(); + list_counts[list_id] = usize_to_i32(count, "list count")?; + if count > 0 { + let id_bytes_len = sorted_lists[list_id].1.len(); + list_id_bytes_lens[list_id] = usize_to_i32(id_bytes_len, "delta ID section")?; + let code_bytes = checked_list_bytes(count, code_size)?; + let factor_bytes = checked_list_bytes(count, FACTOR_BYTES)?; + let list_bytes = 12usize + .checked_add(id_bytes_len) + .and_then(|len| len.checked_add(code_bytes)) + .and_then(|len| len.checked_add(factor_bytes)) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "IVF-RQ list size overflow") + })?; + current_offset = current_offset + .checked_add(list_bytes as u64) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "IVF-RQ offset overflow") + })?; + } + } + + for list_id in 0..nlist { + write_i64_le(out, list_offsets[list_id])?; + write_i32_le(out, list_counts[list_id])?; + write_i32_le(out, list_id_bytes_lens[list_id])?; + } + + for (sorted_ids, id_bytes, sorted_codes, sorted_factors) in sorted_lists { + if sorted_ids.is_empty() { + continue; + } + write_i64_le(out, sorted_ids[0])?; + write_i32_le(out, id_bytes.len() as i32)?; + out.write_all(&id_bytes)?; + out.write_all(&sorted_codes)?; + write_factors(out, &sorted_factors)?; + } + + Ok(()) +} + +pub struct IVFRQIndexReader { + reader: R, + pub d: usize, + pub nlist: usize, + pub metric: MetricType, + pub total_vectors: i64, + pub rotation_seed: u64, + pub rotation_rounds: u32, + pub code_size: usize, + pub num_bits: u32, + pub rotation_type: u32, + pub factor_layout: u32, + pub format_flags: u32, + pub quantizer_centroids: Vec, + pub list_offsets: Vec, + pub list_counts: Vec, + pub list_id_bytes_lens: Vec, + quantizer: RaBitQuantizer, + rotation: RQRotation, + delta_ids: bool, + loaded: bool, +} + +impl IVFRQIndexReader { + pub fn open(mut reader: R) -> io::Result { + let mut cursor = PreadCursor::new(&mut reader, 0); + let magic = read_u32_le(&mut cursor)?; + if magic != IVF_RQ_MAGIC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid IVF-RQ magic: 0x{:08X}", magic), + )); + } + let version = read_u32_le(&mut cursor)?; + if version != IVF_RQ_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported IVF-RQ version: {}", version), + )); + } + let d = validate_positive_i32(read_i32_le(&mut cursor)?, "d")? as usize; + let nlist = validate_positive_i32(read_i32_le(&mut cursor)?, "nlist")? as usize; + let metric_code = read_u32_le(&mut cursor)?; + let metric = MetricType::from_code(metric_code).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Unknown metric type: {}", metric_code), + ) + })?; + let flags = read_u32_le(&mut cursor)?; + let total_vectors = read_i64_le(&mut cursor)?; + if total_vectors < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("negative total vector count: {}", total_vectors), + )); + } + let rotation_seed = read_u64_le(&mut cursor)?; + let rotation_rounds = read_u32_le(&mut cursor)?; + let code_size = validate_positive_i32(read_i32_le(&mut cursor)?, "code_size")? as usize; + let num_bits = read_u32_le(&mut cursor)?; + let rotation_type = read_u32_le(&mut cursor)?; + let factor_layout = read_u32_le(&mut cursor)?; + let format_flags = read_u32_le(&mut cursor)?; + let expected_code_size = d.div_ceil(8); + if code_size != expected_code_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "IVF-RQ code_size {} does not match dimension-derived size {}", + code_size, expected_code_size + ), + )); + } + validate_rq_format(num_bits, rotation_type, factor_layout, format_flags)?; + + let unknown_flags = flags & !SUPPORTED_FLAGS; + if unknown_flags != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported IVF-RQ flags: 0x{:08X}", unknown_flags), + )); + } + if flags & REQUIRED_FLAGS != REQUIRED_FLAGS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-RQ v1 requires delta IDs", + )); + } + + Ok(Self { + reader, + d, + nlist, + metric, + total_vectors, + rotation_seed, + rotation_rounds, + code_size, + num_bits, + rotation_type, + factor_layout, + format_flags, + quantizer_centroids: Vec::new(), + list_offsets: Vec::new(), + list_counts: Vec::new(), + list_id_bytes_lens: Vec::new(), + quantizer: RaBitQuantizer::new(d), + rotation: RQRotation::new(d, rotation_seed, rotation_rounds), + delta_ids: true, + loaded: false, + }) + } + + pub fn ensure_loaded(&mut self) -> io::Result<()> { + if self.loaded { + return Ok(()); + } + let mut cursor = PreadCursor::new(&mut self.reader, IVF_RQ_HEADER_SIZE as u64); + self.quantizer_centroids = + read_f32_vec(&mut cursor, checked_section_size(self.nlist, self.d)?)?; + self.list_offsets = vec![0; self.nlist]; + self.list_counts = vec![0; self.nlist]; + self.list_id_bytes_lens = vec![0; self.nlist]; + for list_id in 0..self.nlist { + self.list_offsets[list_id] = read_i64_le(&mut cursor)?; + let count = read_i32_le(&mut cursor)?; + if count < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("negative list count {} at list {}", count, list_id), + )); + } + self.list_counts[list_id] = count; + let id_bytes_len = read_i32_le(&mut cursor)?; + if id_bytes_len < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("negative id_bytes_len {} at list {}", id_bytes_len, list_id), + )); + } + self.list_id_bytes_lens[list_id] = id_bytes_len; + } + self.loaded = true; + Ok(()) + } + + pub fn read_inverted_list(&mut self, list_id: usize) -> io::Result { + self.ensure_loaded()?; + if list_id >= self.nlist { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("list_id {} out of range (nlist={})", list_id, self.nlist), + )); + } + let count = self.list_counts[list_id] as usize; + if count == 0 { + return Ok(RQReadList { + list_id, + ids: Vec::new(), + codes: Vec::new(), + factors: Vec::new(), + }); + } + if !self.delta_ids { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-RQ reader only supports delta IDs", + )); + } + let offset = checked_list_offset(self.list_offsets[list_id], list_id)?; + let id_bytes_len = self.list_id_bytes_lens[list_id] as usize; + let code_bytes = checked_list_bytes(count, self.code_size)?; + let factor_bytes = checked_list_bytes(count, FACTOR_BYTES)?; + let payload_len = 12usize + .checked_add(id_bytes_len) + .and_then(|len| len.checked_add(code_bytes)) + .and_then(|len| len.checked_add(factor_bytes)) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-RQ list payload overflow") + })?; + let mut payload = vec![0u8; payload_len]; + self.reader + .pread(&mut [ReadRequest::new(offset, &mut payload)])?; + let base_id = i64::from_le_bytes(payload[0..8].try_into().unwrap()); + let encoded_len = i32::from_le_bytes(payload[8..12].try_into().unwrap()); + if encoded_len < 0 || encoded_len as usize != id_bytes_len { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-RQ id_bytes_len mismatch", + )); + } + let ids = decode_delta_varint_ids(base_id, &payload[12..12 + id_bytes_len], count)?; + let code_start = 12 + id_bytes_len; + let factor_start = code_start + code_bytes; + Ok(RQReadList { + list_id, + ids, + codes: payload[code_start..factor_start].to_vec(), + factors: bytes_to_factors(&payload[factor_start..])?, + }) + } + + pub fn search( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + ) -> io::Result<(Vec, Vec)> { + self.search_with_query_bits(query, k, nprobe, DEFAULT_RQ_QUERY_BITS) + } + + pub fn search_with_query_bits( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + query_bits: usize, + ) -> io::Result<(Vec, Vec)> { + self.search_with_filter(query, k, nprobe, None, query_bits) + } + + pub fn search_with_filter( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + query_bits: usize, + ) -> io::Result<(Vec, Vec)> { + search_batch_ivfrq_reader_filter_with_query_bits( + self, query, 1, k, nprobe, filter, query_bits, + ) + } + + pub fn search_with_roaring_filter( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], + ) -> io::Result<(Vec, Vec)> { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.search_with_filter(query, k, nprobe, Some(&filter), DEFAULT_RQ_QUERY_BITS) + } + + pub fn search_with_roaring_filter_and_query_bits( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], + query_bits: usize, + ) -> io::Result<(Vec, Vec)> { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.search_with_filter(query, k, nprobe, Some(&filter), query_bits) + } +} + +pub struct RQReadList { + pub list_id: usize, + pub ids: Vec, + pub codes: Vec, + pub factors: Vec, +} + +pub fn search_batch_ivfrq_reader( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, +) -> io::Result<(Vec, Vec)> { + search_batch_ivfrq_reader_with_query_bits(reader, queries, nq, k, nprobe, DEFAULT_RQ_QUERY_BITS) +} + +pub fn search_batch_ivfrq_reader_with_query_bits( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + query_bits: usize, +) -> io::Result<(Vec, Vec)> { + search_batch_ivfrq_reader_filter_with_query_bits( + reader, queries, nq, k, nprobe, None, query_bits, + ) +} + +pub fn search_batch_ivfrq_reader_filter( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, +) -> io::Result<(Vec, Vec)> { + search_batch_ivfrq_reader_filter_with_query_bits( + reader, + queries, + nq, + k, + nprobe, + filter, + DEFAULT_RQ_QUERY_BITS, + ) +} + +pub fn search_batch_ivfrq_reader_filter_with_query_bits( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + query_bits: usize, +) -> io::Result<(Vec, Vec)> { + reader.ensure_loaded()?; + validate_search_inputs(queries, nq, reader.d, k, nprobe)?; + validate_query_bits(query_bits)?; + + let processed = preprocess_vectors(queries, nq, reader.d, reader.metric); + let (all_probe_indices, _) = kmeans::find_topk_batch( + &processed, + nq, + &reader.quantizer_centroids, + reader.nlist, + reader.d, + nprobe, + ); + + let mut list_to_queries = vec![Vec::new(); reader.nlist]; + let mut unique_lists = Vec::new(); + for (qi, probe_indices) in all_probe_indices.iter().enumerate() { + for &list_id in probe_indices { + if list_to_queries[list_id].is_empty() { + unique_lists.push(list_id); + } + list_to_queries[list_id].push(qi); + } + } + + let mut heaps: Vec = (0..nq).map(|_| TopKHeap::new(k)).collect(); + for list_id in unique_lists { + let count = reader.list_counts[list_id] as usize; + if count == 0 { + continue; + } + let read_list = reader.read_inverted_list(list_id)?; + if read_list.list_id != list_id { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-RQ inverted list read returned wrong list id", + )); + } + let centroid = &reader.quantizer_centroids[list_id * reader.d..(list_id + 1) * reader.d]; + for &qi in &list_to_queries[list_id] { + let query = &processed[qi * reader.d..(qi + 1) * reader.d]; + let rotated_query_residual = + rotated_residual(query, centroid, reader.d, &reader.rotation); + let distance_context = reader.quantizer.prepare_distance_context_with_query_bits( + rotated_query_residual, + query, + count >= RQ_BYTE_LUT_MIN_LIST_SIZE, + query_bits, + ); + scan_read_list( + &read_list, + &reader.quantizer, + reader.code_size, + reader.metric, + &distance_context, + filter, + &mut heaps[qi], + ); + } + } + + let mut result_ids = vec![-1i64; nq * k]; + let mut result_dists = vec![f32::MAX; nq * k]; + for qi in 0..nq { + let sorted = std::mem::replace(&mut heaps[qi], TopKHeap::new(0)).into_sorted(); + let base = qi * k; + for (i, &(dist, id)) in sorted.iter().enumerate() { + result_ids[base + i] = id; + result_dists[base + i] = dist; + } + } + Ok((result_ids, result_dists)) +} + +pub fn search_batch_ivfrq_reader_roaring_filter( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], +) -> io::Result<(Vec, Vec)> { + search_batch_ivfrq_reader_roaring_filter_with_query_bits( + reader, + queries, + nq, + k, + nprobe, + roaring_filter_bytes, + DEFAULT_RQ_QUERY_BITS, + ) +} + +pub fn search_batch_ivfrq_reader_roaring_filter_with_query_bits( + reader: &mut IVFRQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], + query_bits: usize, +) -> io::Result<(Vec, Vec)> { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + search_batch_ivfrq_reader_filter_with_query_bits( + reader, + queries, + nq, + k, + nprobe, + Some(&filter), + query_bits, + ) +} + +fn scan_read_list( + read_list: &RQReadList, + quantizer: &RaBitQuantizer, + code_size: usize, + metric: MetricType, + distance_context: &RQDistanceContext, + filter: Option<&dyn RowIdFilter>, + heap: &mut TopKHeap, +) { + for (local_idx, &id) in read_list.ids.iter().enumerate() { + if filter.map(|f| !f.contains(id)).unwrap_or(false) { + continue; + } + let code = &read_list.codes[local_idx * code_size..(local_idx + 1) * code_size]; + let dist = quantizer.distance_to_code_prepared( + distance_context, + code, + read_list.factors[local_idx], + metric, + ); + heap.push(dist, id); + } +} + +fn rotated_residual(query: &[f32], centroid: &[f32], d: usize, rotation: &RQRotation) -> Vec { + let mut residual = vec![0.0f32; d]; + fvec_madd(query, centroid, -1.0, &mut residual); + rotation.apply(&mut residual); + residual +} + +fn validate_query_bits(query_bits: usize) -> io::Result<()> { + if is_supported_query_bits(query_bits) { + Ok(()) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "invalid IVF-RQ query_bits {}; expected 0, 4, or 8", + query_bits + ), + )) + } +} + +fn validate_rq_format( + num_bits: u32, + rotation_type: u32, + factor_layout: u32, + format_flags: u32, +) -> io::Result<()> { + if !matches!(num_bits, 1 | 2 | 4 | 8) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "unsupported IVF-RQ num_bits {}; expected 1, 2, 4, or 8", + num_bits + ), + )); + } + if rotation_type != IVF_RQ_ROTATION_TYPE_KAC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "unsupported IVF-RQ rotation_type {}; expected {}", + rotation_type, IVF_RQ_ROTATION_TYPE_KAC + ), + )); + } + if factor_layout != IVF_RQ_FACTOR_LAYOUT_RABITQ_V1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "unsupported IVF-RQ factor_layout {}; expected {}", + factor_layout, IVF_RQ_FACTOR_LAYOUT_RABITQ_V1 + ), + )); + } + let unknown_format_flags = format_flags & !SUPPORTED_FORMAT_FLAGS; + if unknown_format_flags != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "unsupported IVF-RQ format_flags: 0x{:08X}", + unknown_format_flags + ), + )); + } + if num_bits != IVF_RQ_NUM_BITS_ONE || format_flags != CURRENT_FORMAT_FLAGS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "IVF-RQ reader currently supports only num_bits=1 without optional sections; got num_bits={}, format_flags=0x{:08X}", + num_bits, format_flags + ), + )); + } + Ok(()) +} + +fn validate_index_shape(index: &IVFRQIndex) -> io::Result<()> { + if index.d == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "dimension must be greater than 0", + )); + } + if index.nlist == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "nlist must be greater than 0", + )); + } + if index.ids.len() != index.nlist + || index.codes.len() != index.nlist + || index.factors.len() != index.nlist + { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "IVF-RQ list storage does not match nlist", + )); + } + let centroid_len = checked_section_size(index.nlist, index.d)?; + if index.quantizer_centroids.len() != centroid_len { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "centroid length {} does not match nlist*d {}", + index.quantizer_centroids.len(), + centroid_len + ), + )); + } + let code_size = index.code_size(); + for list_id in 0..index.nlist { + let count = index.ids[list_id].len(); + let expected_code_len = checked_list_bytes(count, code_size)?; + if index.codes[list_id].len() != expected_code_len { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "list {} code length {} does not match ids*code_size {}", + list_id, + index.codes[list_id].len(), + expected_code_len + ), + )); + } + if index.factors[list_id].len() != count { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "list {} factor count {} does not match ids {}", + list_id, + index.factors[list_id].len(), + count + ), + )); + } + } + Ok(()) +} + +fn write_factors(out: &mut dyn SeekWrite, factors: &[RQCodeFactors]) -> io::Result<()> { + let mut bytes = Vec::with_capacity(factors.len() * FACTOR_BYTES); + for factor in factors { + bytes.extend_from_slice(&factor.residual_norm_sqr.to_le_bytes()); + bytes.extend_from_slice(&factor.vector_norm_sqr.to_le_bytes()); + bytes.extend_from_slice(&factor.dp_multiplier.to_le_bytes()); + } + out.write_all(&bytes) +} + +fn bytes_to_factors(bytes: &[u8]) -> io::Result> { + if !bytes.len().is_multiple_of(FACTOR_BYTES) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-RQ factor section is not 12-byte aligned", + )); + } + Ok(bytes + .chunks_exact(FACTOR_BYTES) + .map(|chunk| RQCodeFactors { + residual_norm_sqr: f32::from_le_bytes(chunk[0..4].try_into().unwrap()), + vector_norm_sqr: f32::from_le_bytes(chunk[4..8].try_into().unwrap()), + dp_multiplier: f32::from_le_bytes(chunk[8..12].try_into().unwrap()), + }) + .collect()) +} + +fn write_u32_le(out: &mut dyn SeekWrite, v: u32) -> io::Result<()> { + out.write_all(&v.to_le_bytes()) +} + +fn write_i32_le(out: &mut dyn SeekWrite, v: i32) -> io::Result<()> { + out.write_all(&v.to_le_bytes()) +} + +fn write_i64_le(out: &mut dyn SeekWrite, v: i64) -> io::Result<()> { + out.write_all(&v.to_le_bytes()) +} + +fn write_u64_le(out: &mut dyn SeekWrite, v: u64) -> io::Result<()> { + out.write_all(&v.to_le_bytes()) +} + +fn write_f32_slice(out: &mut dyn SeekWrite, data: &[f32]) -> io::Result<()> { + let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); + out.write_all(&bytes) +} + +fn read_u32_le(reader: &mut PreadCursor<'_, R>) -> io::Result { + let mut buf = [0u8; 4]; + reader.read_exact(&mut buf)?; + Ok(u32::from_le_bytes(buf)) +} + +fn read_i32_le(reader: &mut PreadCursor<'_, R>) -> io::Result { + let mut buf = [0u8; 4]; + reader.read_exact(&mut buf)?; + Ok(i32::from_le_bytes(buf)) +} + +fn read_i64_le(reader: &mut PreadCursor<'_, R>) -> io::Result { + let mut buf = [0u8; 8]; + reader.read_exact(&mut buf)?; + Ok(i64::from_le_bytes(buf)) +} + +fn read_u64_le(reader: &mut PreadCursor<'_, R>) -> io::Result { + let mut buf = [0u8; 8]; + reader.read_exact(&mut buf)?; + Ok(u64::from_le_bytes(buf)) +} + +fn validate_positive_i32(val: i32, field: &str) -> io::Result { + if val <= 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid header field {}: {} (must be positive)", field, val), + )); + } + Ok(val) +} + +fn usize_to_i32(value: usize, field: &str) -> io::Result { + if value > i32::MAX as usize { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("{} exceeds i32 length limit: {}", field, value), + )); + } + Ok(value as i32) +} + +fn usize_to_i64(value: usize, field: &str) -> io::Result { + if value > i64::MAX as usize { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("{} exceeds i64 length limit: {}", field, value), + )); + } + Ok(value as i64) +} + +fn u64_to_i64(value: u64, field: &str) -> io::Result { + if value > i64::MAX as u64 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("{} exceeds i64 offset limit: {}", field, value), + )); + } + Ok(value as i64) +} + +const MAX_SECTION_ELEMENTS: usize = 1 << 30; + +fn checked_section_size(a: usize, b: usize) -> io::Result { + let result = a.checked_mul(b).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "section size overflow in IVF-RQ header", + ) + })?; + if result > MAX_SECTION_ELEMENTS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "section size {} exceeds maximum {}", + result, MAX_SECTION_ELEMENTS + ), + )); + } + Ok(result) +} + +fn checked_list_offset(offset: i64, list_id: usize) -> io::Result { + if offset < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("negative list offset {} at list {}", offset, list_id), + )); + } + Ok(offset as u64) +} + +fn checked_list_bytes(count: usize, bytes_per_entry: usize) -> io::Result { + count + .checked_mul(bytes_per_entry) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "IVF-RQ list byte size overflow")) +} + +fn read_f32_vec( + reader: &mut PreadCursor<'_, R>, + count: usize, +) -> io::Result> { + let mut buf = vec![0u8; count * 4]; + reader.read_exact(&mut buf)?; + bytes_to_f32_vec(&buf) +} + +fn bytes_to_f32_vec(bytes: &[u8]) -> io::Result> { + if !bytes.len().is_multiple_of(4) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "f32 byte section is not 4-byte aligned", + )); + } + Ok(bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect()) +} + +fn decode_roaring_filter(bytes: &[u8]) -> io::Result { + RoaringTreemap::deserialize_from(bytes).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("invalid RoaringTreemap filter: {}", e), + ) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::io::PosWriter; + use std::io::Cursor; + + #[test] + fn ivfrq_write_read_search_roundtrip() { + let d = 8; + let nlist = 4; + let n = 128; + let data: Vec = (0..n) + .flat_map(|i| { + let cluster = (i % nlist) as f32 * 100.0; + [cluster + i as f32 * 0.01, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + }) + .collect(); + let ids: Vec = (1000..1000 + n as i64).collect(); + let mut index = IVFRQIndex::new(d, nlist, MetricType::L2); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_ivfrq_index(&index, &mut writer).unwrap(); + + let mut reader = IVFRQIndexReader::open(Cursor::new(buf)).unwrap(); + let (labels, distances) = reader.search(&data[7 * d..8 * d], 5, nlist).unwrap(); + + assert_eq!(labels[0], ids[7]); + assert!(distances[0] <= 1e-4); + } + + #[test] + fn ivfrq_reader_search_with_query_bits() { + let d = 16; + let nlist = 4; + let n = 128; + let data: Vec = (0..n) + .flat_map(|i| { + let cluster = (i % nlist) as f32 * 25.0; + (0..d).map(move |j| cluster + i as f32 * 0.02 + j as f32 * 0.125) + }) + .collect(); + let ids: Vec = (1000..1000 + n as i64).collect(); + let mut index = IVFRQIndex::new(d, nlist, MetricType::L2); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_ivfrq_index(&index, &mut writer).unwrap(); + + let query = &data[37 * d..38 * d]; + for query_bits in [0, 4, 8] { + let mut reader = IVFRQIndexReader::open(Cursor::new(buf.clone())).unwrap(); + let (labels, distances) = reader + .search_with_query_bits(query, 5, nlist, query_bits) + .unwrap(); + assert_eq!(labels[0], ids[37], "query_bits={}", query_bits); + if query_bits == 0 { + assert!(distances[0] <= 1e-3); + } else { + assert!(distances[0].is_finite()); + } + } + } + + #[test] + fn ivfrq_reader_rejects_invalid_query_bits() { + let d = 8; + let nlist = 1; + let n = 8; + let data: Vec = (0..n) + .flat_map(|i| (0..d).map(move |j| i as f32 + j as f32 * 0.25)) + .collect(); + let ids: Vec = (0..n as i64).collect(); + let mut index = IVFRQIndex::new(d, nlist, MetricType::L2); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_ivfrq_index(&index, &mut writer).unwrap(); + + let mut reader = IVFRQIndexReader::open(Cursor::new(buf)).unwrap(); + let err = reader + .search_with_query_bits(&data[0..d], 1, nlist, 7) + .unwrap_err(); + assert!(err.to_string().contains("query_bits")); + } + + #[test] + fn ivfrq_header_records_current_rq_format_fields() { + let index = tiny_ivfrq_index(); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_ivfrq_index(&index, &mut writer).unwrap(); + + assert_eq!( + u32::from_le_bytes(buf[48..52].try_into().unwrap()), + IVF_RQ_NUM_BITS_ONE + ); + assert_eq!( + u32::from_le_bytes(buf[52..56].try_into().unwrap()), + IVF_RQ_ROTATION_TYPE_KAC + ); + assert_eq!( + u32::from_le_bytes(buf[56..60].try_into().unwrap()), + IVF_RQ_FACTOR_LAYOUT_RABITQ_V1 + ); + assert_eq!(u32::from_le_bytes(buf[60..64].try_into().unwrap()), 0); + + let reader = IVFRQIndexReader::open(Cursor::new(buf)).unwrap(); + assert_eq!(reader.num_bits, IVF_RQ_NUM_BITS_ONE); + assert_eq!(reader.rotation_type, IVF_RQ_ROTATION_TYPE_KAC); + assert_eq!(reader.factor_layout, IVF_RQ_FACTOR_LAYOUT_RABITQ_V1); + assert_eq!(reader.format_flags, 0); + } + + #[test] + fn ivfrq_reader_rejects_reserved_future_format_fields() { + let index = tiny_ivfrq_index(); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_ivfrq_index(&index, &mut writer).unwrap(); + + let mut future_num_bits = buf.clone(); + future_num_bits[48..52].copy_from_slice(&2u32.to_le_bytes()); + let err = open_err(future_num_bits); + assert!(err.to_string().contains("num_bits=1")); + + let mut unknown_rotation = buf.clone(); + unknown_rotation[52..56].copy_from_slice(&99u32.to_le_bytes()); + let err = open_err(unknown_rotation); + assert!(err.to_string().contains("rotation_type")); + + let mut unknown_layout = buf.clone(); + unknown_layout[56..60].copy_from_slice(&99u32.to_le_bytes()); + let err = open_err(unknown_layout); + assert!(err.to_string().contains("factor_layout")); + + let mut optional_sections = buf; + optional_sections[60..64].copy_from_slice(&FORMAT_FLAG_EX_CODES_PRESENT.to_le_bytes()); + let err = open_err(optional_sections); + assert!(err.to_string().contains("optional sections")); + } + + fn open_err(buf: Vec) -> io::Error { + match IVFRQIndexReader::open(Cursor::new(buf)) { + Ok(_) => panic!("expected IVF-RQ open to fail"), + Err(err) => err, + } + } + + fn tiny_ivfrq_index() -> IVFRQIndex { + let d = 8; + let nlist = 2; + let data: Vec = vec![ + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, + ]; + let ids = vec![7, 42]; + let mut index = IVFRQIndex::new(d, nlist, MetricType::L2); + index.train(&data, 2); + index.add(&data, &ids, 2); + index + } +} diff --git a/core/src/lib.rs b/core/src/lib.rs index 2ef54db..d332b8a 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -33,9 +33,12 @@ pub mod ivfhnswflat_io; pub mod ivfhnswsq; pub mod ivfhnswsq_io; pub mod ivfpq; +pub mod ivfrq; +pub mod ivfrq_io; pub mod kmeans; pub mod opq; pub mod pq; +pub mod rq; pub mod shuffler; pub mod sq; pub mod topk; diff --git a/core/src/rq.rs b/core/src/rq.rs new file mode 100644 index 0000000..2843ccd --- /dev/null +++ b/core/src/rq.rs @@ -0,0 +1,682 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::distance::{fvec_norm_l2sqr, MetricType}; + +pub const DEFAULT_RQ_ROTATION_SEED: u64 = 0x9E37_79B9_7F4A_7C15; +pub const DEFAULT_RQ_ROTATION_ROUNDS: u32 = 3; +pub const RQ_BYTE_LUT_MIN_LIST_SIZE: usize = 64; +pub const DEFAULT_RQ_QUERY_BITS: usize = 0; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RQCodeFactors { + pub residual_norm_sqr: f32, + pub vector_norm_sqr: f32, + pub dp_multiplier: f32, +} + +impl RQCodeFactors { + pub fn zero() -> Self { + Self { + residual_norm_sqr: 0.0, + vector_norm_sqr: 0.0, + dp_multiplier: 0.0, + } + } +} + +#[derive(Debug, Clone)] +pub struct RQRotation { + d: usize, + seed: u64, + rounds: u32, + ops: Vec, +} + +#[derive(Debug, Clone, Copy)] +struct KacOp { + i: usize, + j: usize, + cos: f32, + sin: f32, +} + +impl RQRotation { + pub fn new(d: usize, seed: u64, rounds: u32) -> Self { + let mut rng = SplitMix64::new(seed ^ (d as u64).rotate_left(17)); + let mut ops = Vec::new(); + if d >= 2 { + for _ in 0..rounds { + let mut order: Vec = (0..d).collect(); + for i in (1..d).rev() { + let j = rng.next_usize(i + 1); + order.swap(i, j); + } + for pair in order.chunks_exact(2) { + let angle = (rng.next_f32() * 2.0 - 1.0) * std::f32::consts::PI; + let (sin, cos) = angle.sin_cos(); + ops.push(KacOp { + i: pair[0], + j: pair[1], + cos, + sin, + }); + } + } + } + Self { + d, + seed, + rounds, + ops, + } + } + + pub fn seed(&self) -> u64 { + self.seed + } + + pub fn rounds(&self) -> u32 { + self.rounds + } + + pub fn apply(&self, values: &mut [f32]) { + debug_assert_eq!(values.len(), self.d); + for op in &self.ops { + let x = values[op.i]; + let y = values[op.j]; + values[op.i] = op.cos * x - op.sin * y; + values[op.j] = op.sin * x + op.cos * y; + } + } +} + +#[derive(Debug, Clone)] +pub struct RaBitQuantizer { + d: usize, + inv_sqrt_d: f32, +} + +#[derive(Debug, Clone)] +pub struct RQDistanceContext { + d: usize, + code_size: usize, + rotated_query_residual: Vec, + query_residual_norm_sqr: f32, + query_norm_sqr: f32, + byte_signed_sums: Option>, + quantized_query: Option, +} + +#[derive(Debug, Clone)] +struct RQQuantizedQuery { + scale: f32, + sign_bits: Vec, + magnitude_bit_planes: Vec>, +} + +impl RaBitQuantizer { + pub fn new(d: usize) -> Self { + let inv_sqrt_d = if d == 0 { 1.0 } else { 1.0 / (d as f32).sqrt() }; + Self { d, inv_sqrt_d } + } + + pub fn code_size(&self) -> usize { + self.d.div_ceil(8) + } + + pub fn encode( + &self, + rotated_residual: &[f32], + vector_norm_sqr: f32, + code: &mut [u8], + ) -> RQCodeFactors { + debug_assert_eq!(rotated_residual.len(), self.d); + debug_assert!(code.len() >= self.code_size()); + code[..self.code_size()].fill(0); + + let (residual_norm_sqr, abs_sum) = fvec_norm_l2sqr_abs_sum(rotated_residual); + for (byte_idx, chunk) in rotated_residual.chunks(8).enumerate() { + let mut byte = 0u8; + for (bit, &value) in chunk.iter().enumerate() { + if value > 0.0 { + byte |= 1u8 << bit; + } + } + code[byte_idx] = byte; + } + + let dp_multiplier = if abs_sum > f32::EPSILON { + residual_norm_sqr / (abs_sum * self.inv_sqrt_d) + } else { + 0.0 + }; + RQCodeFactors { + residual_norm_sqr, + vector_norm_sqr, + dp_multiplier, + } + } + + pub fn distance_to_code( + &self, + rotated_query_residual: &[f32], + query: &[f32], + code: &[u8], + factors: RQCodeFactors, + metric: MetricType, + ) -> f32 { + debug_assert_eq!(rotated_query_residual.len(), self.d); + debug_assert_eq!(query.len(), self.d); + + let context = self.prepare_distance_context(rotated_query_residual.to_vec(), query, false); + self.distance_to_code_prepared(&context, code, factors, metric) + } + + pub fn prepare_distance_context( + &self, + rotated_query_residual: Vec, + query: &[f32], + build_byte_lut: bool, + ) -> RQDistanceContext { + self.prepare_distance_context_with_query_bits( + rotated_query_residual, + query, + build_byte_lut, + DEFAULT_RQ_QUERY_BITS, + ) + } + + pub fn prepare_distance_context_with_query_bits( + &self, + rotated_query_residual: Vec, + query: &[f32], + build_byte_lut: bool, + query_bits: usize, + ) -> RQDistanceContext { + debug_assert_eq!(rotated_query_residual.len(), self.d); + debug_assert_eq!(query.len(), self.d); + assert!( + is_supported_query_bits(query_bits), + "unsupported IVF-RQ query_bits {}; expected 0, 4, or 8", + query_bits + ); + + let query_residual_norm_sqr = fvec_norm_l2sqr(&rotated_query_residual); + let query_norm_sqr = fvec_norm_l2sqr(query); + let quantized_query = if query_bits == DEFAULT_RQ_QUERY_BITS { + None + } else { + Some(self.quantize_query(&rotated_query_residual, query_bits)) + }; + let byte_signed_sums = if quantized_query.is_none() && build_byte_lut { + Some(self.build_byte_signed_sums(&rotated_query_residual)) + } else { + None + }; + + RQDistanceContext { + d: self.d, + code_size: self.code_size(), + rotated_query_residual, + query_residual_norm_sqr, + query_norm_sqr, + byte_signed_sums, + quantized_query, + } + } + + pub fn distance_to_code_prepared( + &self, + context: &RQDistanceContext, + code: &[u8], + factors: RQCodeFactors, + metric: MetricType, + ) -> f32 { + debug_assert_eq!(context.d, self.d); + debug_assert!(code.len() >= context.code_size); + + let signed_query_sum = self.signed_query_sum(context, code); + let approx_ip = factors.dp_multiplier * signed_query_sum * self.inv_sqrt_d; + let approx_l2 = (factors.residual_norm_sqr + context.query_residual_norm_sqr + - 2.0 * approx_ip) + .max(0.0); + + match metric { + MetricType::L2 => approx_l2, + MetricType::Cosine => 0.5 * approx_l2, + MetricType::InnerProduct => { + let base = factors.residual_norm_sqr - factors.vector_norm_sqr; + let pre_dist = base + context.query_residual_norm_sqr - 2.0 * approx_ip; + 0.5 * (pre_dist - context.query_norm_sqr) + } + } + } + + fn signed_query_sum(&self, context: &RQDistanceContext, code: &[u8]) -> f32 { + if let Some(quantized_query) = &context.quantized_query { + return quantized_query.signed_query_sum(code, context.code_size); + } + + if let Some(byte_signed_sums) = &context.byte_signed_sums { + let mut sum = 0.0f32; + for byte_idx in 0..context.code_size { + sum += byte_signed_sums[byte_idx * 256 + code[byte_idx] as usize]; + } + return sum; + } + + let mut sum = 0.0f32; + for (dim, &value) in context.rotated_query_residual.iter().enumerate() { + sum += if get_bit(code, dim) { value } else { -value }; + } + sum + } + + fn quantize_query( + &self, + rotated_query_residual: &[f32], + query_bits: usize, + ) -> RQQuantizedQuery { + let magnitude_bits = query_bits - 1; + let max_level = (1usize << magnitude_bits) - 1; + let code_size = self.code_size(); + let max_abs = rotated_query_residual + .iter() + .map(|value| value.abs()) + .fold(0.0f32, f32::max); + let scale = if max_abs > f32::EPSILON { + max_abs / max_level as f32 + } else { + 0.0 + }; + let mut sign_bits = vec![0u8; code_size]; + let mut magnitude_bit_planes = vec![vec![0u8; code_size]; magnitude_bits]; + + if scale == 0.0 { + return RQQuantizedQuery { + scale, + sign_bits, + magnitude_bit_planes, + }; + } + + for (dim, &value) in rotated_query_residual.iter().enumerate() { + if value >= 0.0 { + sign_bits[dim / 8] |= 1u8 << (dim % 8); + } + let level = (value.abs() / scale).round().clamp(0.0, max_level as f32) as usize; + for (bit, plane) in magnitude_bit_planes.iter_mut().enumerate() { + if (level >> bit) & 1 != 0 { + plane[dim / 8] |= 1u8 << (dim % 8); + } + } + } + + RQQuantizedQuery { + scale, + sign_bits, + magnitude_bit_planes, + } + } + + fn build_byte_signed_sums(&self, rotated_query_residual: &[f32]) -> Vec { + let code_size = self.code_size(); + let mut byte_signed_sums = vec![0.0f32; code_size * 256]; + for byte_idx in 0..code_size { + let dim_base = byte_idx * 8; + let dim_end = (dim_base + 8).min(self.d); + for pattern in 0..256usize { + let mut sum = 0.0f32; + for dim in dim_base..dim_end { + let bit = (pattern >> (dim - dim_base)) & 1; + let value = rotated_query_residual[dim]; + sum += if bit != 0 { value } else { -value }; + } + byte_signed_sums[byte_idx * 256 + pattern] = sum; + } + } + byte_signed_sums + } +} + +impl RQQuantizedQuery { + fn signed_query_sum(&self, code: &[u8], code_size: usize) -> f32 { + if self.scale == 0.0 { + return 0.0; + } + + let mut signed_level_sum = 0i64; + for (bit, plane) in self.magnitude_bit_planes.iter().enumerate() { + let weight = 1i64 << bit; + let mut plane_sum = 0i64; + let mut offset = 0usize; + + while offset + 8 <= code_size { + let selected = u64::from_le_bytes(plane[offset..offset + 8].try_into().unwrap()); + if selected != 0 { + let code_bits = + u64::from_le_bytes(code[offset..offset + 8].try_into().unwrap()); + let sign_bits = + u64::from_le_bytes(self.sign_bits[offset..offset + 8].try_into().unwrap()); + let same_sign = !(code_bits ^ sign_bits) & selected; + plane_sum += 2 * same_sign.count_ones() as i64 - selected.count_ones() as i64; + } + offset += 8; + } + + while offset < code_size { + let selected = plane[offset]; + if selected != 0 { + let same_sign = !(code[offset] ^ self.sign_bits[offset]) & selected; + plane_sum += 2 * same_sign.count_ones() as i64 - selected.count_ones() as i64; + } + offset += 1; + } + + signed_level_sum += weight * plane_sum; + } + + self.scale * signed_level_sum as f32 + } +} + +#[inline] +pub fn is_supported_query_bits(query_bits: usize) -> bool { + matches!(query_bits, 0 | 4 | 8) +} + +fn get_bit(code: &[u8], dim: usize) -> bool { + code[dim / 8] & (1u8 << (dim % 8)) != 0 +} + +#[cfg(target_arch = "x86_64")] +#[inline] +fn fvec_norm_l2sqr_abs_sum(values: &[f32]) -> (f32, f32) { + if is_x86_feature_detected!("avx2") && values.len() >= 8 { + unsafe { fvec_norm_l2sqr_abs_sum_avx2(values) } + } else { + fvec_norm_l2sqr_abs_sum_scalar(values) + } +} + +#[cfg(target_arch = "aarch64")] +#[inline] +fn fvec_norm_l2sqr_abs_sum(values: &[f32]) -> (f32, f32) { + unsafe { fvec_norm_l2sqr_abs_sum_neon(values) } +} + +#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] +#[inline] +fn fvec_norm_l2sqr_abs_sum(values: &[f32]) -> (f32, f32) { + fvec_norm_l2sqr_abs_sum_scalar(values) +} + +#[cfg(any( + target_arch = "x86_64", + not(any(target_arch = "x86_64", target_arch = "aarch64")) +))] +#[inline] +fn fvec_norm_l2sqr_abs_sum_scalar(values: &[f32]) -> (f32, f32) { + let mut norm = 0.0f32; + let mut abs_sum = 0.0f32; + for &value in values { + norm += value * value; + abs_sum += value.abs(); + } + (norm, abs_sum) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn fvec_norm_l2sqr_abs_sum_avx2(values: &[f32]) -> (f32, f32) { + use std::arch::x86_64::*; + + let n = values.len(); + let abs_mask = _mm256_castsi256_ps(_mm256_set1_epi32(0x7fff_ffff)); + let mut norm_sum = _mm256_setzero_ps(); + let mut abs_sum_vec = _mm256_setzero_ps(); + let mut i = 0; + while i + 8 <= n { + let value = unsafe { _mm256_loadu_ps(values.as_ptr().add(i)) }; + norm_sum = _mm256_add_ps(norm_sum, _mm256_mul_ps(value, value)); + abs_sum_vec = _mm256_add_ps(abs_sum_vec, _mm256_and_ps(value, abs_mask)); + i += 8; + } + + let norm_hi = _mm256_extractf128_ps::<1>(norm_sum); + let norm_lo = _mm256_castps256_ps128(norm_sum); + let norm_128 = _mm_add_ps(norm_lo, norm_hi); + let norm_64 = _mm_add_ps(norm_128, _mm_movehl_ps(norm_128, norm_128)); + let norm_32 = _mm_add_ss(norm_64, _mm_shuffle_ps::<1>(norm_64, norm_64)); + let mut norm = _mm_cvtss_f32(norm_32); + + let abs_hi = _mm256_extractf128_ps::<1>(abs_sum_vec); + let abs_lo = _mm256_castps256_ps128(abs_sum_vec); + let abs_128 = _mm_add_ps(abs_lo, abs_hi); + let abs_64 = _mm_add_ps(abs_128, _mm_movehl_ps(abs_128, abs_128)); + let abs_32 = _mm_add_ss(abs_64, _mm_shuffle_ps::<1>(abs_64, abs_64)); + let mut abs_sum = _mm_cvtss_f32(abs_32); + + while i < n { + let value = unsafe { *values.get_unchecked(i) }; + norm += value * value; + abs_sum += value.abs(); + i += 1; + } + (norm, abs_sum) +} + +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +unsafe fn fvec_norm_l2sqr_abs_sum_neon(values: &[f32]) -> (f32, f32) { + use std::arch::aarch64::*; + + let n = values.len(); + let mut norm_sum = vdupq_n_f32(0.0); + let mut abs_sum_vec = vdupq_n_f32(0.0); + let mut i = 0; + while i + 4 <= n { + let value = unsafe { vld1q_f32(values.as_ptr().add(i)) }; + norm_sum = vmlaq_f32(norm_sum, value, value); + abs_sum_vec = vaddq_f32(abs_sum_vec, vabsq_f32(value)); + i += 4; + } + + let mut norm = vaddvq_f32(norm_sum); + let mut abs_sum = vaddvq_f32(abs_sum_vec); + while i < n { + let value = unsafe { *values.get_unchecked(i) }; + norm += value * value; + abs_sum += value.abs(); + i += 1; + } + (norm, abs_sum) +} + +#[derive(Debug, Clone, Copy)] +struct SplitMix64 { + state: u64, +} + +impl SplitMix64 { + fn new(seed: u64) -> Self { + Self { state: seed } + } + + fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.state; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn next_usize(&mut self, upper: usize) -> usize { + (self.next_u64() % upper as u64) as usize + } + + fn next_f32(&mut self) -> f32 { + let mantissa = (self.next_u64() >> 40) as u32; + mantissa as f32 / ((1u32 << 24) as f32) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rabit_quantizer_estimates_self_distance_as_zero() { + let d = 8; + let rotation = RQRotation::new(d, DEFAULT_RQ_ROTATION_SEED, DEFAULT_RQ_ROTATION_ROUNDS); + let quantizer = RaBitQuantizer::new(d); + let centroid = vec![1.0; d]; + let vector = vec![2.0, 1.5, 0.75, 3.0, 1.25, -1.0, 4.0, 2.5]; + let mut residual: Vec = vector + .iter() + .zip(centroid.iter()) + .map(|(&x, &c)| x - c) + .collect(); + rotation.apply(&mut residual); + + let mut code = vec![0u8; quantizer.code_size()]; + let factors = quantizer.encode(&residual, fvec_norm_l2sqr(&vector), &mut code); + let dist = quantizer.distance_to_code(&residual, &vector, &code, factors, MetricType::L2); + + assert!( + dist <= 1e-5, + "self distance should be close to zero: {dist}" + ); + } + + #[test] + fn distance_context_byte_lut_matches_scalar_path() { + let d = 16; + let quantizer = RaBitQuantizer::new(d); + let rotated_residual: Vec = (0..d).map(|i| i as f32 * 0.25 - 1.5).collect(); + let query: Vec = (0..d).map(|i| (i as f32 + 1.0) * 0.125).collect(); + let rotated_query_residual: Vec = (0..d).map(|i| (i as f32 - 3.0) * 0.2).collect(); + + let mut code = vec![0u8; quantizer.code_size()]; + let factors = quantizer.encode(&rotated_residual, fvec_norm_l2sqr(&query), &mut code); + code[0] = 0b1010_0101; + code[1] = 0b0101_1010; + + let scalar_context = + quantizer.prepare_distance_context(rotated_query_residual.clone(), &query, false); + let lut_context = quantizer.prepare_distance_context(rotated_query_residual, &query, true); + + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + let scalar = + quantizer.distance_to_code_prepared(&scalar_context, &code, factors, metric); + let lut = quantizer.distance_to_code_prepared(&lut_context, &code, factors, metric); + assert!( + (scalar - lut).abs() < 1e-5, + "metric {:?}: scalar {} != lut {}", + metric, + scalar, + lut + ); + } + } + + #[test] + fn quantized_query_bit_planes_match_scalar_quantization() { + let d = 24; + let quantizer = RaBitQuantizer::new(d); + let rotated_query_residual: Vec = (0..d).map(|i| (i as f32 - 11.0) * 0.17).collect(); + let query: Vec = (0..d).map(|i| (i as f32 + 1.0) * 0.03125).collect(); + let mut code = vec![0u8; quantizer.code_size()]; + for (byte_idx, byte) in code.iter_mut().enumerate() { + *byte = if byte_idx % 2 == 0 { + 0b1010_1100 + } else { + 0b0101_0011 + }; + } + + for query_bits in [4, 8] { + let context = quantizer.prepare_distance_context_with_query_bits( + rotated_query_residual.clone(), + &query, + true, + query_bits, + ); + let quantized_query = context.quantized_query.as_ref().unwrap(); + let actual = quantized_query.signed_query_sum(&code, quantizer.code_size()); + let expected = + scalar_quantized_signed_query_sum(&rotated_query_residual, &code, query_bits); + + assert!( + (actual - expected).abs() < 1e-5, + "query_bits {}: {} != {}", + query_bits, + actual, + expected + ); + assert!(context.byte_signed_sums.is_none()); + } + } + + #[test] + fn norm_l2sqr_abs_sum_helper_matches_expected_values() { + let values = [-3.0f32, 4.0, 0.5, -0.25, 8.0, -2.0, 1.25, -6.0, 7.0]; + let (norm, abs_sum) = fvec_norm_l2sqr_abs_sum(&values); + + let expected_norm: f32 = values.iter().map(|value| value * value).sum(); + let expected_abs_sum: f32 = values.iter().map(|value| value.abs()).sum(); + assert!((norm - expected_norm).abs() < 1e-5); + assert!((abs_sum - expected_abs_sum).abs() < 1e-5); + } + + #[test] + fn rotation_preserves_l2_norm() { + let d = 16; + let rotation = RQRotation::new(d, 11, 3); + let mut vector: Vec = (0..d).map(|i| i as f32 - 3.0).collect(); + let before = fvec_norm_l2sqr(&vector); + rotation.apply(&mut vector); + let after = fvec_norm_l2sqr(&vector); + assert!((before - after).abs() < 1e-4); + } + + fn scalar_quantized_signed_query_sum( + rotated_query_residual: &[f32], + code: &[u8], + query_bits: usize, + ) -> f32 { + let magnitude_bits = query_bits - 1; + let max_level = (1usize << magnitude_bits) - 1; + let max_abs = rotated_query_residual + .iter() + .map(|value| value.abs()) + .fold(0.0f32, f32::max); + if max_abs <= f32::EPSILON { + return 0.0; + } + let scale = max_abs / max_level as f32; + let mut sum = 0i64; + for (dim, &value) in rotated_query_residual.iter().enumerate() { + let code_sign = if get_bit(code, dim) { 1i64 } else { -1i64 }; + let query_sign = if value >= 0.0 { 1i64 } else { -1i64 }; + let level = (value.abs() / scale).round().clamp(0.0, max_level as f32) as i64; + sum += code_sign * query_sign * level; + } + scale * sum as f32 + } +} diff --git a/core/tests/fixtures/ivf_rq_v1.hex b/core/tests/fixtures/ivf_rq_v1.hex new file mode 100644 index 0000000..df50d56 --- /dev/null +++ b/core/tests/fixtures/ivf_rq_v1.hex @@ -0,0 +1,8 @@ +51 52 56 49 01 00 00 00 08 00 00 00 02 00 00 00 00 00 00 00 01 00 00 00 03 00 00 00 00 00 00 00 +15 7c 4a 7f b9 79 37 9e 03 00 00 00 01 00 00 00 01 00 00 00 01 00 00 00 01 00 00 00 00 00 00 00 +00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 +00 00 20 41 00 00 20 41 00 00 20 41 00 00 20 41 00 00 20 41 00 00 20 41 00 00 20 41 00 00 20 41 +a0 00 00 00 00 00 00 00 02 00 00 00 02 00 00 00 c8 00 00 00 00 00 00 00 01 00 00 00 01 00 00 00 +07 00 00 00 00 00 00 00 02 00 00 00 00 23 00 ff 00 00 00 00 00 00 00 00 00 00 00 00 00 00 80 3f +00 00 80 3f 00 00 00 00 63 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 +00 00 diff --git a/core/tests/storage_format_fixtures.rs b/core/tests/storage_format_fixtures.rs index fea6da3..8f52ca5 100644 --- a/core/tests/storage_format_fixtures.rs +++ b/core/tests/storage_format_fixtures.rs @@ -28,6 +28,9 @@ use paimon_vindex_core::ivfhnswflat_io::write_ivfhnswflat_index; use paimon_vindex_core::ivfhnswsq::IVFHNSWSQIndex; use paimon_vindex_core::ivfhnswsq_io::write_ivfhnswsq_index; use paimon_vindex_core::ivfpq::IVFPQIndex; +use paimon_vindex_core::ivfrq::IVFRQIndex; +use paimon_vindex_core::ivfrq_io::write_ivfrq_index; +use paimon_vindex_core::rq::RQCodeFactors; use paimon_vindex_core::sq::ScalarQuantizer; use std::fmt::Write as _; use std::io::Cursor; @@ -176,6 +179,21 @@ fn fixture_cases() -> Vec { params: VectorSearchParams::new(2, 2), expected_first_id: 5, }, + FixtureCase { + name: "ivf_rq_v1", + fixture_hex: include_str!("fixtures/ivf_rq_v1.hex"), + build: build_ivf_rq_fixture, + index_type: IndexType::IvfRq, + dimension: 8, + nlist: 2, + metric: MetricType::L2, + total_vectors: 3, + pq_m: None, + hnsw: None, + query: vec![0.0; 8], + params: VectorSearchParams::new(2, 2), + expected_first_id: 7, + }, FixtureCase { name: "ivf_hnsw_flat_v1", fixture_hex: include_str!("fixtures/ivf_hnsw_flat_v1.hex"), @@ -251,6 +269,30 @@ fn build_ivf_pq_4bit_fixture() -> Vec { buf } +fn build_ivf_rq_fixture() -> Vec { + let mut index = IVFRQIndex::new(8, 2, MetricType::L2); + index.quantizer_centroids = vec![ + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, + ]; + index.ids = vec![vec![42, 7], vec![99]]; + index.codes = vec![vec![0xFF, 0x00], vec![0x00]]; + index.factors = vec![ + vec![ + RQCodeFactors { + residual_norm_sqr: 1.0, + vector_norm_sqr: 1.0, + dp_multiplier: 0.0, + }, + RQCodeFactors::zero(), + ], + vec![RQCodeFactors::zero()], + ]; + + let mut buf = Vec::new(); + write_ivfrq_index(&index, &mut PosWriter::new(&mut buf)).unwrap(); + buf +} + fn build_ivf_hnsw_flat_fixture() -> Vec { let mut index = IVFHNSWFlatIndex::new(2, 2, MetricType::L2, fixture_hnsw_params()); index.flat.quantizer_centroids = vec![0.0, 0.0, 10.0, 10.0]; diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp index ec1b5e3..5d097f5 100644 --- a/cpp/test_vindex.cpp +++ b/cpp/test_vindex.cpp @@ -45,7 +45,7 @@ struct MemBuffer { size_t pos = 0; }; -constexpr size_t kRoundtripDimension = 2; +constexpr size_t kRoundtripDimension = 8; constexpr size_t kRoundtripNlist = 4; constexpr size_t kRoundtripPerList = 128; constexpr size_t kRoundtripVectorCount = kRoundtripNlist * kRoundtripPerList; @@ -82,8 +82,11 @@ static std::vector roundtrip_data() { size_t cluster = i / kRoundtripPerList; size_t local = i % kRoundtripPerList; float center = static_cast(cluster) * 20.0f; - data[i * kRoundtripDimension] = center + static_cast(local % 16) * 0.001f; - data[i * kRoundtripDimension + 1] = center + static_cast(local / 16) * 0.001f; + for (size_t dim = 0; dim < kRoundtripDimension; dim++) { + data[i * kRoundtripDimension + dim] = + center + static_cast(dim) * 0.01f + + static_cast(local % 16) * 0.001f; + } } return data; } @@ -104,6 +107,14 @@ static void assert_id_in_cluster(int64_t id, size_t cluster) { ASSERT_TRUE(id < base + static_cast(kRoundtripPerList)); } +static std::vector query_for_center(float center) { + std::vector query(kRoundtripDimension); + for (size_t dim = 0; dim < kRoundtripDimension; dim++) { + query[dim] = center + static_cast(dim) * 0.01f; + } + return query; +} + static void run_roundtrip( const char* name, const std::vector>& options, @@ -113,12 +124,12 @@ static void run_roundtrip( std::vector data = roundtrip_data(); std::vector ids = roundtrip_ids(); paimon::vindex::Trainer trainer(options); - ASSERT_EQ(trainer.dimension(), 2); + ASSERT_EQ(trainer.dimension(), kRoundtripDimension); paimon::vindex::Training training = trainer.add_training_vectors(data.data(), kRoundtripVectorCount).finish_training(); paimon::vindex::Writer writer(std::move(training)); - ASSERT_EQ(writer.dimension(), 2); + ASSERT_EQ(writer.dimension(), kRoundtripDimension); writer.add_vectors(ids.data(), data.data(), kRoundtripVectorCount); MemBuffer buf; @@ -128,7 +139,7 @@ static void run_roundtrip( paimon::vindex::Reader reader(make_input(buf)); auto metadata = reader.metadata(); ASSERT_EQ(metadata.index_type, expected_index_type); - ASSERT_EQ(metadata.dimension, 2); + ASSERT_EQ(metadata.dimension, kRoundtripDimension); ASSERT_EQ(metadata.nlist, 4); ASSERT_EQ(metadata.metric, PAIMON_VINDEX_METRIC_L2); ASSERT_EQ(metadata.total_vectors, kRoundtripVectorCount); @@ -137,17 +148,31 @@ static void run_roundtrip( reader.optimize_for_search(); - const float query[] = {0.0f, 0.0f}; - auto result = reader.search(query, 2, 4, 16); + auto query = query_for_center(0.0f); + auto result = reader.search(query.data(), 2, 4, 16); ASSERT_EQ(result.ids.size(), 2); assert_id_in_cluster(result.ids[0], 0); ASSERT_TRUE(std::isfinite(result.distances[0])); + if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { + auto query_bits_result = reader.search(query.data(), 2, 4, 16, 4); + assert_id_in_cluster(query_bits_result.ids[0], 0); + ASSERT_TRUE(std::isfinite(query_bits_result.distances[0])); + } - const float queries[] = {0.0f, 0.0f, 20.0f, 20.0f}; - auto batch = reader.search_batch(queries, 2, 1, 4, 16); + auto query0 = query_for_center(0.0f); + auto query1 = query_for_center(20.0f); + std::vector queries; + queries.insert(queries.end(), query0.begin(), query0.end()); + queries.insert(queries.end(), query1.begin(), query1.end()); + auto batch = reader.search_batch(queries.data(), 2, 1, 4, 16); ASSERT_EQ(batch.ids.size(), 2); assert_id_in_cluster(batch.ids[0], 0); assert_id_in_cluster(batch.ids[1], 1); + if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { + auto query_bits_batch = reader.search_batch(queries.data(), 2, 1, 4, 16, 8); + assert_id_in_cluster(query_bits_batch.ids[0], 0); + assert_id_in_cluster(query_bits_batch.ids[1], 1); + } printf("PASS %s\n", name); } @@ -156,7 +181,7 @@ static void test_supported_index_roundtrips() { "ivf_flat_roundtrip", { {"index.type", "ivf_flat"}, - {"dimension", "2"}, + {"dimension", "8"}, {"nlist", "4"}, {"metric", "l2"}, }, @@ -168,20 +193,32 @@ static void test_supported_index_roundtrips() { "ivf_pq_roundtrip", { {"index.type", "ivf_pq"}, - {"dimension", "2"}, + {"dimension", "8"}, {"nlist", "4"}, {"metric", "l2"}, - {"pq.m", "1"}, + {"pq.m", "4"}, }, PAIMON_VINDEX_INDEX_TYPE_IVF_PQ, - 1, + 4, + 0); + + run_roundtrip( + "ivf_rq_roundtrip", + { + {"index.type", "ivf_rq"}, + {"dimension", "8"}, + {"nlist", "4"}, + {"metric", "l2"}, + }, + PAIMON_VINDEX_INDEX_TYPE_IVF_RQ, + 0, 0); run_roundtrip( "ivf_hnsw_flat_roundtrip", { {"index.type", "ivf_hnsw_flat"}, - {"dimension", "2"}, + {"dimension", "8"}, {"nlist", "4"}, {"metric", "l2"}, {"hnsw.m", "4"}, @@ -194,7 +231,7 @@ static void test_supported_index_roundtrips() { "ivf_hnsw_sq_roundtrip", { {"index.type", "ivf_hnsw_sq"}, - {"dimension", "2"}, + {"dimension", "8"}, {"nlist", "4"}, {"metric", "l2"}, {"hnsw.m", "4"}, diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 1e719e3..60893ed 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -35,6 +35,7 @@ pub const PAIMON_VINDEX_INDEX_TYPE_IVF_FLAT: u32 = 0; pub const PAIMON_VINDEX_INDEX_TYPE_IVF_PQ: u32 = 1; pub const PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_FLAT: u32 = 2; pub const PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_SQ: u32 = 3; +pub const PAIMON_VINDEX_INDEX_TYPE_IVF_RQ: u32 = 4; pub const PAIMON_VINDEX_METRIC_L2: u32 = 0; pub const PAIMON_VINDEX_METRIC_INNER_PRODUCT: u32 = 1; @@ -683,6 +684,31 @@ pub unsafe extern "C" fn paimon_vindex_reader_search( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_with_query_bits( + handle: *mut PaimonVindexReaderHandle, + query: *const f32, + top_k: usize, + nprobe: usize, + ef_search: usize, + query_bits: usize, + out_ids: *mut i64, + out_distances: *mut f32, + result_len: usize, +) -> c_int { + ffi_status(|| { + let handle = unsafe { reader_mut(handle) }?; + let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; + let params = + VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let (ids, distances) = handle + .inner + .search(query, params) + .map_err(|e| format!("search: {}", e))?; + copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( handle: *mut PaimonVindexReaderHandle, @@ -709,6 +735,34 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter_and_query_bits( + handle: *mut PaimonVindexReaderHandle, + query: *const f32, + top_k: usize, + nprobe: usize, + ef_search: usize, + query_bits: usize, + roaring_filter: *const u8, + roaring_filter_len: usize, + out_ids: *mut i64, + out_distances: *mut f32, + result_len: usize, +) -> c_int { + ffi_status(|| { + let handle = unsafe { reader_mut(handle) }?; + let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; + let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; + let params = + VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let (ids, distances) = handle + .inner + .search_with_roaring_filter(query, params, filter) + .map_err(|e| format!("search_with_roaring_filter: {}", e))?; + copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_batch( handle: *mut PaimonVindexReaderHandle, @@ -742,6 +796,41 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch( }) } +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_query_bits( + handle: *mut PaimonVindexReaderHandle, + queries: *const f32, + query_count: usize, + top_k: usize, + nprobe: usize, + ef_search: usize, + query_bits: usize, + out_ids: *mut i64, + out_distances: *mut f32, + result_len: usize, +) -> c_int { + ffi_status(|| { + let handle = unsafe { reader_mut(handle) }?; + let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; + let queries = unsafe { const_slice(queries, query_len, "queries") }?; + let expected_len = checked_len(query_count, top_k, "batch result")?; + let params = + VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let (ids, distances) = handle + .inner + .search_batch(queries, query_count, params) + .map_err(|e| format!("search_batch: {}", e))?; + copy_search_result( + &ids, + &distances, + out_ids, + out_distances, + result_len, + expected_len, + ) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter( handle: *mut PaimonVindexReaderHandle, @@ -777,3 +866,41 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter( ) }) } + +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( + handle: *mut PaimonVindexReaderHandle, + queries: *const f32, + query_count: usize, + top_k: usize, + nprobe: usize, + ef_search: usize, + query_bits: usize, + roaring_filter: *const u8, + roaring_filter_len: usize, + out_ids: *mut i64, + out_distances: *mut f32, + result_len: usize, +) -> c_int { + ffi_status(|| { + let handle = unsafe { reader_mut(handle) }?; + let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; + let queries = unsafe { const_slice(queries, query_len, "queries") }?; + let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; + let expected_len = checked_len(query_count, top_k, "batch result")?; + let params = + VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let (ids, distances) = handle + .inner + .search_batch_with_roaring_filter(queries, query_count, params, filter) + .map_err(|e| format!("search_batch_with_roaring_filter: {}", e))?; + copy_search_result( + &ids, + &distances, + out_ids, + out_distances, + result_len, + expected_len, + ) + }) +} diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp index 1216835..33e105c 100644 --- a/include/paimon_vindex.hpp +++ b/include/paimon_vindex.hpp @@ -354,16 +354,22 @@ class Reader { check(paimon_vindex_reader_optimize_for_search(handle_)); } - SearchResult search(const float* query, size_t top_k, size_t nprobe, size_t ef_search = 0) { + SearchResult search( + const float* query, + size_t top_k, + size_t nprobe, + size_t ef_search = 0, + size_t query_bits = 0) { SearchResult result; result.ids.resize(top_k); result.distances.resize(top_k); - check(paimon_vindex_reader_search( + check(paimon_vindex_reader_search_with_query_bits( handle_, query, top_k, nprobe, ef_search, + query_bits, result.ids.data(), result.distances.data(), top_k)); @@ -377,15 +383,27 @@ class Reader { size_t ef_search, const uint8_t* filter, size_t filter_len) { + return search_with_roaring_filter(query, top_k, nprobe, ef_search, 0, filter, filter_len); + } + + SearchResult search_with_roaring_filter( + const float* query, + size_t top_k, + size_t nprobe, + size_t ef_search, + size_t query_bits, + const uint8_t* filter, + size_t filter_len) { SearchResult result; result.ids.resize(top_k); result.distances.resize(top_k); - check(paimon_vindex_reader_search_with_roaring_filter( + check(paimon_vindex_reader_search_with_roaring_filter_and_query_bits( handle_, query, top_k, nprobe, ef_search, + query_bits, filter, filter_len, result.ids.data(), @@ -399,18 +417,20 @@ class Reader { size_t query_count, size_t top_k, size_t nprobe, - size_t ef_search = 0) { + size_t ef_search = 0, + size_t query_bits = 0) { const size_t result_len = query_count * top_k; SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch( + check(paimon_vindex_reader_search_batch_with_query_bits( handle_, queries, query_count, top_k, nprobe, ef_search, + query_bits, result.ids.data(), result.distances.data(), result_len)); @@ -425,17 +445,31 @@ class Reader { size_t ef_search, const uint8_t* filter, size_t filter_len) { + return search_batch_with_roaring_filter( + queries, query_count, top_k, nprobe, ef_search, 0, filter, filter_len); + } + + SearchResult search_batch_with_roaring_filter( + const float* queries, + size_t query_count, + size_t top_k, + size_t nprobe, + size_t ef_search, + size_t query_bits, + const uint8_t* filter, + size_t filter_len) { const size_t result_len = query_count * top_k; SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch_with_roaring_filter( + check(paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( handle_, queries, query_count, top_k, nprobe, ef_search, + query_bits, filter, filter_len, result.ids.data(), diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java index e7cf220..d4d4445 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java @@ -51,12 +51,33 @@ private VectorIndexNative() {} static native VectorSearchResult search(long ptr, float[] query, int k, int nprobe, int efSearch); + static native VectorSearchResult searchWithQueryBits( + long ptr, float[] query, int k, int nprobe, int efSearch, int queryBits); + static native VectorSearchResult searchWithRoaringFilter( long ptr, float[] query, int k, int nprobe, int efSearch, byte[] roaringFilter); + static native VectorSearchResult searchWithRoaringFilterAndQueryBits( + long ptr, + float[] query, + int k, + int nprobe, + int efSearch, + int queryBits, + byte[] roaringFilter); + static native VectorSearchBatchResult searchBatch( long ptr, float[] queries, int queryCount, int k, int nprobe, int efSearch); + static native VectorSearchBatchResult searchBatchWithQueryBits( + long ptr, + float[] queries, + int queryCount, + int k, + int nprobe, + int efSearch, + int queryBits); + static native VectorSearchBatchResult searchBatchWithRoaringFilter( long ptr, float[] queries, @@ -66,5 +87,15 @@ static native VectorSearchBatchResult searchBatchWithRoaringFilter( int efSearch, byte[] roaringFilter); + static native VectorSearchBatchResult searchBatchWithRoaringFilterAndQueryBits( + long ptr, + float[] queries, + int queryCount, + int k, + int nprobe, + int efSearch, + int queryBits, + byte[] roaringFilter); + static native void freeReader(long ptr); } diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java index 2494cfc..847aae5 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java @@ -82,11 +82,17 @@ public VectorSearchResult search(float[] query, int topK, int nprobe) { } public VectorSearchResult search(float[] query, int topK, int nprobe, int efSearch) { + return search(query, topK, nprobe, efSearch, 0); + } + + public VectorSearchResult search( + float[] query, int topK, int nprobe, int efSearch, int queryBits) { validateQuery(query); synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.search(requireOpen(), query, topK, nprobe, efSearch); + return VectorIndexNative.searchWithQueryBits( + requireOpen(), query, topK, nprobe, efSearch, queryBits); } finally { exitNativeHandle(); } @@ -99,6 +105,16 @@ public VectorSearchResult search(float[] query, int topK, int nprobe, byte[] roa public VectorSearchResult search( float[] query, int topK, int nprobe, int efSearch, byte[] roaringFilter) { + return search(query, topK, nprobe, efSearch, 0, roaringFilter); + } + + public VectorSearchResult search( + float[] query, + int topK, + int nprobe, + int efSearch, + int queryBits, + byte[] roaringFilter) { validateQuery(query); if (roaringFilter == null) { throw new NullPointerException("roaringFilter"); @@ -106,8 +122,8 @@ public VectorSearchResult search( synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchWithRoaringFilter( - requireOpen(), query, topK, nprobe, efSearch, roaringFilter); + return VectorIndexNative.searchWithRoaringFilterAndQueryBits( + requireOpen(), query, topK, nprobe, efSearch, queryBits, roaringFilter); } finally { exitNativeHandle(); } @@ -121,14 +137,19 @@ public VectorSearchBatchResult searchBatch( public VectorSearchBatchResult searchBatch( float[] queries, int queryCount, int topK, int nprobe, int efSearch) { + return searchBatch(queries, queryCount, topK, nprobe, efSearch, 0); + } + + public VectorSearchBatchResult searchBatch( + float[] queries, int queryCount, int topK, int nprobe, int efSearch, int queryBits) { if (queries == null) { throw new NullPointerException("queries"); } synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchBatch( - requireOpen(), queries, queryCount, topK, nprobe, efSearch); + return VectorIndexNative.searchBatchWithQueryBits( + requireOpen(), queries, queryCount, topK, nprobe, efSearch, queryBits); } finally { exitNativeHandle(); } @@ -147,6 +168,17 @@ public VectorSearchBatchResult searchBatch( int nprobe, int efSearch, byte[] roaringFilter) { + return searchBatch(queries, queryCount, topK, nprobe, efSearch, 0, roaringFilter); + } + + public VectorSearchBatchResult searchBatch( + float[] queries, + int queryCount, + int topK, + int nprobe, + int efSearch, + int queryBits, + byte[] roaringFilter) { if (queries == null) { throw new NullPointerException("queries"); } @@ -156,8 +188,15 @@ public VectorSearchBatchResult searchBatch( synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchBatchWithRoaringFilter( - requireOpen(), queries, queryCount, topK, nprobe, efSearch, roaringFilter); + return VectorIndexNative.searchBatchWithRoaringFilterAndQueryBits( + requireOpen(), + queries, + queryCount, + topK, + nprobe, + efSearch, + queryBits, + roaringFilter); } finally { exitNativeHandle(); } diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java index 5d75836..fb3a406 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java @@ -232,13 +232,24 @@ private static void testReaderAndWriterApiCompile() { reader.optimizeForSearch(); reader.search(new float[] {0.0f, 1.0f}, 10, 4); reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32); + reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, 4); reader.search(new float[] {0.0f, 1.0f}, 10, 4, new byte[] {1, 2}); reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, new byte[] {1, 2}); + reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, 4, new byte[] {1, 2}); reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4); reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32); + reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32, 4); reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, new byte[] {1, 2}); reader.searchBatch( new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32, new byte[] {1, 2}); + reader.searchBatch( + new float[] {0.0f, 1.0f, 2.0f, 3.0f}, + 2, + 10, + 4, + 32, + 4, + new byte[] {1, 2}); VectorIndexTraining training = VectorIndexTrainer.train(options, new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2); diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java index 07db0ca..3ee2164 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java @@ -23,7 +23,7 @@ public class VectorIndexNativeValidationTest { - private static final int ROUNDTRIP_DIMENSION = 2; + private static final int ROUNDTRIP_DIMENSION = 8; private static final int ROUNDTRIP_NLIST = 4; private static final int ROUNDTRIP_PER_LIST = 128; private static final int ROUNDTRIP_VECTOR_COUNT = ROUNDTRIP_NLIST * ROUNDTRIP_PER_LIST; @@ -168,7 +168,9 @@ private static void testStagedTrainingRoundtrip() { runStagedTrainingRoundtrip( "ivf_flat", ivfFlatOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); runStagedTrainingRoundtrip( - "ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST, 1), 1, 0); + "ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST, 4), 4, 0); + runStagedTrainingRoundtrip( + "ivf_rq", ivfRqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); runStagedTrainingRoundtrip( "ivf_hnsw_flat", ivfHnswOptions("ivf_hnsw_flat", ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), @@ -301,7 +303,8 @@ public void run() { private static void testSupportedIndexRoundtrips() { runRoundtrip("ivf_flat", ivfFlatOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); - runRoundtrip("ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST, 1), 1, 0); + runRoundtrip("ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST, 4), 4, 0); + runRoundtrip("ivf_rq", ivfRqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); runRoundtrip( "ivf_hnsw_flat", ivfHnswOptions("ivf_hnsw_flat", ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), @@ -346,16 +349,38 @@ private static void assertRoundtrip( reader.optimizeForSearch(); - VectorSearchResult single = reader.search(new float[] {0.0f, 0.0f}, 2, 4, 16); + VectorSearchResult single = reader.search(queryForCenter(0.0f), 2, 4, 16); assertIdInCluster(single.ids()[0], 0); assertFinite(single.distances()[0], indexType + " single distance"); VectorSearchBatchResult batch = - reader.searchBatch(new float[] {0.0f, 0.0f, 20.0f, 20.0f}, 2, 1, 4, 16); + reader.searchBatch(batchQueries(), 2, 1, 4, 16); assertIdInCluster(batch.ids()[0], 0); assertIdInCluster(batch.ids()[1], 1); assertFinite(batch.distances()[0], indexType + " batch distance 0"); assertFinite(batch.distances()[1], indexType + " batch distance 1"); + + if ("ivf_rq".equals(indexType)) { + VectorSearchResult queryBitsSingle = + reader.search(queryForCenter(0.0f), 2, 4, 16, 4); + assertIdInCluster(queryBitsSingle.ids()[0], 0); + assertFinite(queryBitsSingle.distances()[0], "ivf_rq queryBits single distance"); + + VectorSearchBatchResult queryBitsBatch = + reader.searchBatch(batchQueries(), 2, 1, 4, 16, 8); + assertIdInCluster(queryBitsBatch.ids()[0], 0); + assertIdInCluster(queryBitsBatch.ids()[1], 1); + + assertThrowsMessage( + RuntimeException.class, + "query_bits", + new ThrowingRunnable() { + @Override + public void run() { + reader.search(queryForCenter(0.0f), 2, 4, 16, 7); + } + }); + } } finally { reader.close(); } @@ -436,6 +461,12 @@ private static Map ivfPqOptions(int dimension, int nlist, int m) return options; } + private static Map ivfRqOptions(int dimension, int nlist) { + Map options = ivfFlatOptions(dimension, nlist); + options.put("index.type", "ivf_rq"); + return options; + } + private static Map ivfHnswOptions(String indexType, int dimension, int nlist) { Map options = ivfFlatOptions(dimension, nlist); options.put("index.type", indexType); @@ -451,12 +482,31 @@ private static float[] roundtripData() { int cluster = i / ROUNDTRIP_PER_LIST; int local = i % ROUNDTRIP_PER_LIST; float center = cluster * 20.0f; - data[i * ROUNDTRIP_DIMENSION] = center + (local % 16) * 0.001f; - data[i * ROUNDTRIP_DIMENSION + 1] = center + (local / 16) * 0.001f; + for (int dim = 0; dim < ROUNDTRIP_DIMENSION; dim++) { + data[i * ROUNDTRIP_DIMENSION + dim] = + center + dim * 0.01f + (local % 16) * 0.001f; + } } return data; } + private static float[] queryForCenter(float center) { + float[] query = new float[ROUNDTRIP_DIMENSION]; + for (int dim = 0; dim < ROUNDTRIP_DIMENSION; dim++) { + query[dim] = center + dim * 0.01f; + } + return query; + } + + private static float[] batchQueries() { + float[] first = queryForCenter(0.0f); + float[] second = queryForCenter(20.0f); + float[] queries = new float[2 * ROUNDTRIP_DIMENSION]; + System.arraycopy(first, 0, queries, 0, ROUNDTRIP_DIMENSION); + System.arraycopy(second, 0, queries, ROUNDTRIP_DIMENSION, ROUNDTRIP_DIMENSION); + return queries; + } + private static long[] roundtripIds() { long[] ids = new long[ROUNDTRIP_VECTOR_COUNT]; for (int i = 0; i < ROUNDTRIP_VECTOR_COUNT; i++) { diff --git a/jni/src/lib.rs b/jni/src/lib.rs index 82f92c7..944bc93 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -369,14 +369,20 @@ fn build_metadata(env: &mut JNIEnv, metadata: VectorIndexMetadata) -> jobject { result.into_raw() } -fn search_params(k: jint, nprobe: jint, ef_search: jint) -> Option { - if k < 0 || nprobe < 0 || ef_search < 0 { +fn search_params( + k: jint, + nprobe: jint, + ef_search: jint, + query_bits: jint, +) -> Option { + if k < 0 || nprobe < 0 || ef_search < 0 || query_bits < 0 { None } else { - Some(VectorSearchParams::with_ef_search( + Some(VectorSearchParams::with_ef_search_and_query_bits( k as usize, nprobe as usize, ef_search as usize, + query_bits as usize, )) } } @@ -693,7 +699,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea Some(reader) => reader, None => return throw_and_return(env, "null native pointer (reader already freed?)"), }; - let params = match search_params(k, nprobe, ef_search) { + let params = match search_params(k, nprobe, ef_search, 0) { Some(params) => params, None => { return throw_and_return( @@ -717,6 +723,46 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea }) } +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchWithQueryBits( + env: JNIEnv, + _class: JClass, + ptr: jlong, + query: JFloatArray, + k: jint, + nprobe: jint, + ef_search: jint, + query_bits: jint, +) -> jobject { + jni_call(env, |env| { + let reader = match deref_reader(ptr) { + Some(reader) => reader, + None => return throw_and_return(env, "null native pointer (reader already freed?)"), + }; + let params = match search_params(k, nprobe, ef_search, query_bits) { + Some(params) => params, + None => { + return throw_and_return( + env, + &format!( + "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", + k, nprobe, ef_search, query_bits + ), + ) + } + }; + let query_buf = match read_float_array(env, &query, "query") { + Ok(buf) => buf, + Err(e) => return throw_and_return(env, &e), + }; + let (ids, dists) = match reader.search(&query_buf, params) { + Ok(result) => result, + Err(e) => return throw_and_return(env, &format!("search: {}", e)), + }; + build_result(env, ids, dists) + }) +} + #[no_mangle] pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchWithRoaringFilter( env: JNIEnv, @@ -733,7 +779,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea Some(reader) => reader, None => return throw_and_return(env, "null native pointer (reader already freed?)"), }; - let params = match search_params(k, nprobe, ef_search) { + let params = match search_params(k, nprobe, ef_search, 0) { Some(params) => params, None => { return throw_and_return( @@ -762,6 +808,52 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea }) } +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchWithRoaringFilterAndQueryBits( + env: JNIEnv, + _class: JClass, + ptr: jlong, + query: JFloatArray, + k: jint, + nprobe: jint, + ef_search: jint, + query_bits: jint, + roaring_filter: JByteArray, +) -> jobject { + jni_call(env, |env| { + let reader = match deref_reader(ptr) { + Some(reader) => reader, + None => return throw_and_return(env, "null native pointer (reader already freed?)"), + }; + let params = match search_params(k, nprobe, ef_search, query_bits) { + Some(params) => params, + None => { + return throw_and_return( + env, + &format!( + "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", + k, nprobe, ef_search, query_bits + ), + ) + } + }; + let query_buf = match read_float_array(env, &query, "query") { + Ok(buf) => buf, + Err(e) => return throw_and_return(env, &e), + }; + let filter_bytes = match read_byte_array(env, roaring_filter) { + Ok(bytes) => bytes, + Err(e) => return throw_and_return(env, &e), + }; + let (ids, dists) = + match reader.search_with_roaring_filter(&query_buf, params, &filter_bytes) { + Ok(result) => result, + Err(e) => return throw_and_return(env, &format!("search_with_filter: {}", e)), + }; + build_result(env, ids, dists) + }) +} + #[no_mangle] pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatch( env: JNIEnv, @@ -781,7 +873,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea if query_count < 0 { return throw_and_return(env, &format!("invalid query count: {}", query_count)); } - let params = match search_params(k, nprobe, ef_search) { + let params = match search_params(k, nprobe, ef_search, 0) { Some(params) => params, None => { return throw_and_return( @@ -806,6 +898,51 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea }) } +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatchWithQueryBits( + env: JNIEnv, + _class: JClass, + ptr: jlong, + queries: JFloatArray, + query_count: jint, + k: jint, + nprobe: jint, + ef_search: jint, + query_bits: jint, +) -> jobject { + jni_call(env, |env| { + let reader = match deref_reader(ptr) { + Some(reader) => reader, + None => return throw_and_return(env, "null native pointer (reader already freed?)"), + }; + if query_count < 0 { + return throw_and_return(env, &format!("invalid query count: {}", query_count)); + } + let params = match search_params(k, nprobe, ef_search, query_bits) { + Some(params) => params, + None => { + return throw_and_return( + env, + &format!( + "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", + k, nprobe, ef_search, query_bits + ), + ) + } + }; + let nq = query_count as usize; + let query_buf = match read_float_array(env, &queries, "queries") { + Ok(buf) => buf, + Err(e) => return throw_and_return(env, &e), + }; + let (ids, dists) = match reader.search_batch(&query_buf, nq, params) { + Ok(result) => result, + Err(e) => return throw_and_return(env, &format!("search_batch: {}", e)), + }; + build_batch_result(env, ids, dists, nq, params.top_k) + }) +} + #[no_mangle] pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatchWithRoaringFilter( env: JNIEnv, @@ -826,7 +963,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea if query_count < 0 { return throw_and_return(env, &format!("invalid query count: {}", query_count)); } - let params = match search_params(k, nprobe, ef_search) { + let params = match search_params(k, nprobe, ef_search, 0) { Some(params) => params, None => { return throw_and_return( @@ -858,6 +995,59 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea }) } +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatchWithRoaringFilterAndQueryBits( + env: JNIEnv, + _class: JClass, + ptr: jlong, + queries: JFloatArray, + query_count: jint, + k: jint, + nprobe: jint, + ef_search: jint, + query_bits: jint, + roaring_filter: JByteArray, +) -> jobject { + jni_call(env, |env| { + let reader = match deref_reader(ptr) { + Some(reader) => reader, + None => return throw_and_return(env, "null native pointer (reader already freed?)"), + }; + if query_count < 0 { + return throw_and_return(env, &format!("invalid query count: {}", query_count)); + } + let params = match search_params(k, nprobe, ef_search, query_bits) { + Some(params) => params, + None => { + return throw_and_return( + env, + &format!( + "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", + k, nprobe, ef_search, query_bits + ), + ) + } + }; + let nq = query_count as usize; + let query_buf = match read_float_array(env, &queries, "queries") { + Ok(buf) => buf, + Err(e) => return throw_and_return(env, &e), + }; + let filter_bytes = match read_byte_array(env, roaring_filter) { + Ok(bytes) => bytes, + Err(e) => return throw_and_return(env, &e), + }; + let (ids, dists) = + match reader.search_batch_with_roaring_filter(&query_buf, nq, params, &filter_bytes) { + Ok(result) => result, + Err(e) => { + return throw_and_return(env, &format!("search_batch_with_filter: {}", e)) + } + }; + build_batch_result(env, ids, dists, nq, params.top_k) + }) +} + #[no_mangle] pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_freeReader( env: JNIEnv, diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py index f66c47c..ea10ed4 100644 --- a/python/paimon_vindex/__init__.py +++ b/python/paimon_vindex/__init__.py @@ -29,6 +29,7 @@ 1: "ivf_pq", 2: "ivf_hnsw_flat", 3: "ivf_hnsw_sq", + 4: "ivf_rq", } METRICS = { @@ -427,7 +428,7 @@ def _filter_args(self, filter_bytes): return None, 0, None return _bytes_buffer(filter_bytes, "filter_bytes") - def search(self, query, top_k, nprobe, ef_search=0, filter_bytes=None): + def search(self, query, top_k, nprobe, ef_search=0, query_bits=0, filter_bytes=None): self._require_open() query = _float32_vector(query, "query") if query.shape[0] != self._metadata.dimension: @@ -439,24 +440,26 @@ def search(self, query, top_k, nprobe, ef_search=0, filter_bytes=None): distances = np.empty(top_k, dtype=np.float32) if filter_bytes is None: - rc = lib.paimon_vindex_reader_search( + rc = lib.paimon_vindex_reader_search_with_query_bits( self._handle, query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), top_k, nprobe, ef_search, + query_bits, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), top_k, ) else: filter_buf, filter_len, _ = self._filter_args(filter_bytes) - rc = lib.paimon_vindex_reader_search_with_roaring_filter( + rc = lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits( self._handle, query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), top_k, nprobe, ef_search, + query_bits, filter_buf, filter_len, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), @@ -467,7 +470,7 @@ def search(self, query, top_k, nprobe, ef_search=0, filter_bytes=None): _check_error("search failed") return ids, distances - def search_batch(self, queries, top_k, nprobe, ef_search=0, filter_bytes=None): + def search_batch(self, queries, top_k, nprobe, ef_search=0, query_bits=0, filter_bytes=None): self._require_open() queries = _float32_matrix(queries, "queries") if queries.shape[1] != self._metadata.dimension: @@ -480,26 +483,28 @@ def search_batch(self, queries, top_k, nprobe, ef_search=0, filter_bytes=None): distances = np.empty((queries.shape[0], top_k), dtype=np.float32) if filter_bytes is None: - rc = lib.paimon_vindex_reader_search_batch( + rc = lib.paimon_vindex_reader_search_batch_with_query_bits( self._handle, queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), queries.shape[0], top_k, nprobe, ef_search, + query_bits, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), result_len, ) else: filter_buf, filter_len, _ = self._filter_args(filter_bytes) - rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter( + rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( self._handle, queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), queries.shape[0], top_k, nprobe, ef_search, + query_bits, filter_buf, filter_len, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py index 38c0ea7..16deebf 100644 --- a/python/paimon_vindex/_ffi.py +++ b/python/paimon_vindex/_ffi.py @@ -197,6 +197,19 @@ class PaimonVindexMetadata(Structure): ] lib.paimon_vindex_reader_search.restype = c_int +lib.paimon_vindex_reader_search_with_query_bits.argtypes = [ + c_void_p, + POINTER(c_float), + c_size_t, + c_size_t, + c_size_t, + c_size_t, + POINTER(c_int64), + POINTER(c_float), + c_size_t, +] +lib.paimon_vindex_reader_search_with_query_bits.restype = c_int + lib.paimon_vindex_reader_search_with_roaring_filter.argtypes = [ c_void_p, POINTER(c_float), @@ -211,6 +224,21 @@ class PaimonVindexMetadata(Structure): ] lib.paimon_vindex_reader_search_with_roaring_filter.restype = c_int +lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits.argtypes = [ + c_void_p, + POINTER(c_float), + c_size_t, + c_size_t, + c_size_t, + c_size_t, + POINTER(c_uint8), + c_size_t, + POINTER(c_int64), + POINTER(c_float), + c_size_t, +] +lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits.restype = c_int + lib.paimon_vindex_reader_search_batch.argtypes = [ c_void_p, POINTER(c_float), @@ -224,6 +252,20 @@ class PaimonVindexMetadata(Structure): ] lib.paimon_vindex_reader_search_batch.restype = c_int +lib.paimon_vindex_reader_search_batch_with_query_bits.argtypes = [ + c_void_p, + POINTER(c_float), + c_size_t, + c_size_t, + c_size_t, + c_size_t, + c_size_t, + POINTER(c_int64), + POINTER(c_float), + c_size_t, +] +lib.paimon_vindex_reader_search_batch_with_query_bits.restype = c_int + lib.paimon_vindex_reader_search_batch_with_roaring_filter.argtypes = [ c_void_p, POINTER(c_float), @@ -238,3 +280,19 @@ class PaimonVindexMetadata(Structure): c_size_t, ] lib.paimon_vindex_reader_search_batch_with_roaring_filter.restype = c_int + +lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits.argtypes = [ + c_void_p, + POINTER(c_float), + c_size_t, + c_size_t, + c_size_t, + c_size_t, + c_size_t, + POINTER(c_uint8), + c_size_t, + POINTER(c_int64), + POINTER(c_float), + c_size_t, +] +lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits.restype = c_int diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py index a6e79bc..9dae5ee 100644 --- a/python/tests/test_vindex.py +++ b/python/tests/test_vindex.py @@ -77,6 +77,15 @@ def test_python_ffi_roundtrips_supported_indexes(): }, 16, ), + ( + { + "index.type": "ivf_rq", + "dimension": "16", + "nlist": "4", + "metric": "l2", + }, + 16, + ), ( { "index.type": "ivf_hnsw_flat", @@ -143,6 +152,37 @@ def test_python_ffi_batch_search(): assert ids[1, 0] == 1 +def test_python_ffi_ivfrq_query_bits(): + index_bytes, data = build_index( + { + "index.type": "ivf_rq", + "dimension": "16", + "nlist": "4", + "metric": "l2", + }, + 16, + n=128, + ) + + with reader_from_bytes(index_bytes) as reader: + for query_bits in (4, 8): + ids, distances = reader.search( + data[7], top_k=5, nprobe=4, query_bits=query_bits + ) + assert ids.shape == (5,) + assert distances.shape == (5,) + assert ids[0] % 4 == 7 % 4 + + ids, distances = reader.search_batch( + np.vstack([data[4], data[7]]), top_k=5, nprobe=4, query_bits=4 + ) + assert ids[0, 0] % 4 == 4 % 4 + assert ids[1, 0] % 4 == 7 % 4 + + with pytest.raises(RuntimeError, match="query_bits"): + reader.search(data[0], top_k=5, nprobe=4, query_bits=7) + + def test_python_ffi_delegates_validation(): options = { "index.type": "ivf_pq", From a1686ba604e04394aacd562868950f4b1792dcfe Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 9 Jul 2026 18:55:16 +0800 Subject: [PATCH 2/4] Fix clippy warnings in IVF RQ code --- core/src/index.rs | 12 +++++------- core/src/ivfrq_io.rs | 42 +++++++++++++++++++++++++++++------------- 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/core/src/index.rs b/core/src/index.rs index 8dac92a..cc66c1b 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -867,13 +867,11 @@ fn validate_config(config: &VectorIndexConfig) -> io::Result<()> { )); } } - VectorIndexConfig::IvfRq { dimension, .. } => { - if !dimension.is_multiple_of(8) { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - format!("dimension {} must be divisible by 8 for IVF_RQ", dimension), - )); - } + VectorIndexConfig::IvfRq { dimension, .. } if !dimension.is_multiple_of(8) => { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("dimension {} must be divisible by 8 for IVF_RQ", dimension), + )); } VectorIndexConfig::IvfHnswFlat { hnsw, .. } | VectorIndexConfig::IvfHnswSq { hnsw, .. } => { validate_hnsw_params(*hnsw)? diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs index 925d434..ea1bbcb 100644 --- a/core/src/ivfrq_io.rs +++ b/core/src/ivfrq_io.rs @@ -49,6 +49,13 @@ const FORMAT_FLAG_ERROR_FACTOR_PRESENT: u32 = 1 << 1; const SUPPORTED_FORMAT_FLAGS: u32 = FORMAT_FLAG_EX_CODES_PRESENT | FORMAT_FLAG_ERROR_FACTOR_PRESENT; const CURRENT_FORMAT_FLAGS: u32 = 0; +struct SortedRQList { + ids: Vec, + id_bytes: Vec, + codes: Vec, + factors: Vec, +} + pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Result<()> { validate_index_shape(index)?; let d = index.d; @@ -64,12 +71,16 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Res }) })?; - let mut sorted_lists: Vec<(Vec, Vec, Vec, Vec)> = - Vec::with_capacity(nlist); + let mut sorted_lists = Vec::with_capacity(nlist); for list_id in 0..nlist { let count = index.ids[list_id].len(); if count == 0 { - sorted_lists.push((Vec::new(), Vec::new(), Vec::new(), Vec::new())); + sorted_lists.push(SortedRQList { + ids: Vec::new(), + id_bytes: Vec::new(), + codes: Vec::new(), + factors: Vec::new(), + }); continue; } @@ -85,7 +96,12 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Res sorted_factors.push(index.factors[list_id][idx]); } let (_, id_bytes) = encode_delta_varint_ids(&sorted_ids); - sorted_lists.push((sorted_ids, id_bytes, sorted_codes, sorted_factors)); + sorted_lists.push(SortedRQList { + ids: sorted_ids, + id_bytes, + codes: sorted_codes, + factors: sorted_factors, + }); } write_u32_le(out, IVF_RQ_MAGIC)?; @@ -127,10 +143,10 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Res for list_id in 0..nlist { list_offsets[list_id] = u64_to_i64(current_offset, "list offset")?; - let count = sorted_lists[list_id].0.len(); + let count = sorted_lists[list_id].ids.len(); list_counts[list_id] = usize_to_i32(count, "list count")?; if count > 0 { - let id_bytes_len = sorted_lists[list_id].1.len(); + let id_bytes_len = sorted_lists[list_id].id_bytes.len(); list_id_bytes_lens[list_id] = usize_to_i32(id_bytes_len, "delta ID section")?; let code_bytes = checked_list_bytes(count, code_size)?; let factor_bytes = checked_list_bytes(count, FACTOR_BYTES)?; @@ -155,15 +171,15 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) -> io::Res write_i32_le(out, list_id_bytes_lens[list_id])?; } - for (sorted_ids, id_bytes, sorted_codes, sorted_factors) in sorted_lists { - if sorted_ids.is_empty() { + for sorted_list in sorted_lists { + if sorted_list.ids.is_empty() { continue; } - write_i64_le(out, sorted_ids[0])?; - write_i32_le(out, id_bytes.len() as i32)?; - out.write_all(&id_bytes)?; - out.write_all(&sorted_codes)?; - write_factors(out, &sorted_factors)?; + write_i64_le(out, sorted_list.ids[0])?; + write_i32_le(out, sorted_list.id_bytes.len() as i32)?; + out.write_all(&sorted_list.id_bytes)?; + out.write_all(&sorted_list.codes)?; + write_factors(out, &sorted_list.factors)?; } Ok(()) From 6ddd64f56504d2a9df4a5c85bf2c88093df451dd Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 9 Jul 2026 19:31:25 +0800 Subject: [PATCH 3/4] Use search params structs across bindings --- README.md | 26 +- c/test_vindex.c | 16 +- cpp/test_vindex.cpp | 10 +- ffi/src/lib.rs | 178 +++-------- include/paimon_vindex.hpp | 107 +++---- .../index/vector/VectorIndexNative.java | 41 +-- .../index/vector/VectorIndexReader.java | 93 ++---- .../index/vector/VectorSearchParams.java | 61 ++++ .../index/vector/VectorIndexJavaApiTest.java | 34 +-- .../VectorIndexNativeHandleSafetyTest.java | 2 +- .../VectorIndexNativePanicBoundaryTest.java | 2 +- .../VectorIndexNativeValidationTest.java | 31 +- jni/src/lib.rs | 288 +++--------------- python/paimon_vindex/__init__.py | 64 ++-- python/paimon_vindex/_ffi.py | 83 +---- python/tests/test_vindex.py | 24 +- 16 files changed, 332 insertions(+), 728 deletions(-) create mode 100644 java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java diff --git a/README.md b/README.md index 28e6853..a949786 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,12 @@ reader.optimize_for_search()?; let params = VectorSearchParams::with_ef_search(10, 16, 80); let (ids, distances) = reader.search(&query, params)?; -let rq_params = VectorSearchParams::with_query_bits(10, 16, 4); +let rq_params = VectorSearchParams { + top_k: 10, + nprobe: 16, + ef_search: 0, + query_bits: 4, +}; let (ids, distances) = reader.search(&query, rq_params)?; ``` @@ -181,8 +186,14 @@ paimon_vindex_reader_optimize_for_search(reader); int64_t ids[10]; float distances[10]; +PaimonVindexSearchParams search_params = { + .top_k = 10, + .nprobe = 16, + .ef_search = 80, + .query_bits = 0, +}; paimon_vindex_reader_search( - reader, query, 10, 16, 80, ids, distances, 10); + reader, query, search_params, ids, distances, 10); paimon_vindex_reader_free(reader); ``` @@ -215,7 +226,7 @@ writer.write_index(output_file); paimon::vindex::Reader reader(input_file); auto metadata = reader.metadata(); reader.optimize_for_search(); -auto result = reader.search(query.data(), 10, 16, 80); +auto result = reader.search(query.data(), paimon::vindex::SearchParams{10, 16, 80}); ``` ### Java/JNI @@ -228,6 +239,7 @@ import java.util.Map; import org.apache.paimon.index.vector.VectorIndexInput; import org.apache.paimon.index.vector.VectorIndexMetadata; import org.apache.paimon.index.vector.VectorIndexReader; +import org.apache.paimon.index.vector.VectorSearchParams; import org.apache.paimon.index.vector.VectorIndexTrainer; import org.apache.paimon.index.vector.VectorIndexTraining; import org.apache.paimon.index.vector.VectorSearchResult; @@ -273,7 +285,7 @@ try (VectorIndexTrainer trainer = VectorIndexTrainer.create(options)) { try (VectorIndexReader reader = new VectorIndexReader(vectorIndexInput)) { VectorIndexMetadata metadata = reader.metadata(); reader.optimizeForSearch(); - VectorSearchResult result = reader.search(query, 10, 16, 80); + VectorSearchResult result = reader.search(query, new VectorSearchParams(10, 16, 80, 0)); } ``` @@ -284,7 +296,7 @@ and validates the options when the trainer is created. ### Python ```python -from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter +from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter class VectorIndexInput: @@ -311,8 +323,8 @@ writer.write(output) reader = VectorIndexReader(VectorIndexInput(index_bytes)) reader.optimize_for_search() -ids, distances = reader.search(query, top_k=10, nprobe=16, ef_search=80) -ids, distances = reader.search(query, top_k=10, nprobe=16, query_bits=4) +ids, distances = reader.search(query, SearchParams(top_k=10, nprobe=16, ef_search=80)) +ids, distances = reader.search(query, SearchParams(top_k=10, nprobe=16, query_bits=4)) ``` The Python package is pure Python and uses `ctypes` to load diff --git a/c/test_vindex.c b/c/test_vindex.c index 6792139..7fa3f90 100644 --- a/c/test_vindex.c +++ b/c/test_vindex.c @@ -262,15 +262,17 @@ static void run_roundtrip( fill_query(query, 0.0f); int64_t result_ids[2] = {0}; float result_distances[2] = {0}; + struct PaimonVindexSearchParams search_params = {2, 4, 16, 0}; if (paimon_vindex_reader_search( - reader, query, 2, 4, 16, result_ids, result_distances, 2) != 0) { + reader, query, search_params, result_ids, result_distances, 2) != 0) { fail_ffi("reader search failed"); } assert_id_in_cluster(result_ids[0], 0); ASSERT_TRUE(isfinite(result_distances[0])); if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { - if (paimon_vindex_reader_search_with_query_bits( - reader, query, 2, 4, 16, 4, result_ids, result_distances, 2) != 0) { + search_params.query_bits = 4; + if (paimon_vindex_reader_search( + reader, query, search_params, result_ids, result_distances, 2) != 0) { fail_ffi("reader search with query bits failed"); } assert_id_in_cluster(result_ids[0], 0); @@ -282,15 +284,17 @@ static void run_roundtrip( fill_query(queries + ROUNDTRIP_DIMENSION, 20.0f); int64_t batch_ids[2] = {0}; float batch_distances[2] = {0}; + struct PaimonVindexSearchParams batch_params = {1, 4, 16, 0}; if (paimon_vindex_reader_search_batch( - reader, queries, 2, 1, 4, 16, batch_ids, batch_distances, 2) != 0) { + reader, queries, 2, batch_params, batch_ids, batch_distances, 2) != 0) { fail_ffi("reader search batch failed"); } assert_id_in_cluster(batch_ids[0], 0); assert_id_in_cluster(batch_ids[1], 1); if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { - if (paimon_vindex_reader_search_batch_with_query_bits( - reader, queries, 2, 1, 4, 16, 8, batch_ids, batch_distances, 2) != 0) { + batch_params.query_bits = 8; + if (paimon_vindex_reader_search_batch( + reader, queries, 2, batch_params, batch_ids, batch_distances, 2) != 0) { fail_ffi("reader search batch with query bits failed"); } assert_id_in_cluster(batch_ids[0], 0); diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp index 5d097f5..be842ae 100644 --- a/cpp/test_vindex.cpp +++ b/cpp/test_vindex.cpp @@ -149,12 +149,13 @@ static void run_roundtrip( reader.optimize_for_search(); auto query = query_for_center(0.0f); - auto result = reader.search(query.data(), 2, 4, 16); + auto result = reader.search(query.data(), paimon::vindex::SearchParams{2, 4, 16}); ASSERT_EQ(result.ids.size(), 2); assert_id_in_cluster(result.ids[0], 0); ASSERT_TRUE(std::isfinite(result.distances[0])); if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { - auto query_bits_result = reader.search(query.data(), 2, 4, 16, 4); + auto query_bits_result = + reader.search(query.data(), paimon::vindex::SearchParams{2, 4, 16, 4}); assert_id_in_cluster(query_bits_result.ids[0], 0); ASSERT_TRUE(std::isfinite(query_bits_result.distances[0])); } @@ -164,12 +165,13 @@ static void run_roundtrip( std::vector queries; queries.insert(queries.end(), query0.begin(), query0.end()); queries.insert(queries.end(), query1.begin(), query1.end()); - auto batch = reader.search_batch(queries.data(), 2, 1, 4, 16); + auto batch = reader.search_batch(queries.data(), 2, paimon::vindex::SearchParams{1, 4, 16}); ASSERT_EQ(batch.ids.size(), 2); assert_id_in_cluster(batch.ids[0], 0); assert_id_in_cluster(batch.ids[1], 1); if (expected_index_type == PAIMON_VINDEX_INDEX_TYPE_IVF_RQ) { - auto query_bits_batch = reader.search_batch(queries.data(), 2, 1, 4, 16, 8); + auto query_bits_batch = + reader.search_batch(queries.data(), 2, paimon::vindex::SearchParams{1, 4, 16, 8}); assert_id_in_cluster(query_bits_batch.ids[0], 0); assert_id_in_cluster(query_bits_batch.ids[1], 1); } diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 60893ed..46ae9c8 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -215,6 +215,15 @@ pub struct PaimonVindexMetadata { pub hnsw_max_level: usize, } +#[repr(C)] +#[derive(Clone, Copy)] +pub struct PaimonVindexSearchParams { + pub top_k: usize, + pub nprobe: usize, + pub ef_search: usize, + pub query_bits: usize, +} + pub struct PaimonVindexTrainerHandle { inner: Option, } @@ -424,6 +433,15 @@ fn copy_search_result( Ok(()) } +fn search_params_from_ffi(params: PaimonVindexSearchParams) -> VectorSearchParams { + VectorSearchParams { + top_k: params.top_k, + nprobe: params.nprobe, + ef_search: params.ef_search, + query_bits: params.query_bits, + } +} + // ======================== Trainer / Writer ======================== #[no_mangle] @@ -665,33 +683,7 @@ pub unsafe extern "C" fn paimon_vindex_reader_optimize_for_search( pub unsafe extern "C" fn paimon_vindex_reader_search( handle: *mut PaimonVindexReaderHandle, query: *const f32, - top_k: usize, - nprobe: usize, - ef_search: usize, - out_ids: *mut i64, - out_distances: *mut f32, - result_len: usize, -) -> c_int { - ffi_status(|| { - let handle = unsafe { reader_mut(handle) }?; - let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; - let params = VectorSearchParams::with_ef_search(top_k, nprobe, ef_search); - let (ids, distances) = handle - .inner - .search(query, params) - .map_err(|e| format!("search: {}", e))?; - copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) - }) -} - -#[no_mangle] -pub unsafe extern "C" fn paimon_vindex_reader_search_with_query_bits( - handle: *mut PaimonVindexReaderHandle, - query: *const f32, - top_k: usize, - nprobe: usize, - ef_search: usize, - query_bits: usize, + params: PaimonVindexSearchParams, out_ids: *mut i64, out_distances: *mut f32, result_len: usize, @@ -699,13 +691,19 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_with_query_bits( ffi_status(|| { let handle = unsafe { reader_mut(handle) }?; let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; - let params = - VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let params = search_params_from_ffi(params); let (ids, distances) = handle .inner .search(query, params) .map_err(|e| format!("search: {}", e))?; - copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) + copy_search_result( + &ids, + &distances, + out_ids, + out_distances, + result_len, + params.top_k, + ) }) } @@ -713,9 +711,7 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_with_query_bits( pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( handle: *mut PaimonVindexReaderHandle, query: *const f32, - top_k: usize, - nprobe: usize, - ef_search: usize, + params: PaimonVindexSearchParams, roaring_filter: *const u8, roaring_filter_len: usize, out_ids: *mut i64, @@ -726,85 +722,28 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter( let handle = unsafe { reader_mut(handle) }?; let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; - let params = VectorSearchParams::with_ef_search(top_k, nprobe, ef_search); + let params = search_params_from_ffi(params); let (ids, distances) = handle .inner .search_with_roaring_filter(query, params, filter) .map_err(|e| format!("search_with_roaring_filter: {}", e))?; - copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) - }) -} - -#[no_mangle] -pub unsafe extern "C" fn paimon_vindex_reader_search_with_roaring_filter_and_query_bits( - handle: *mut PaimonVindexReaderHandle, - query: *const f32, - top_k: usize, - nprobe: usize, - ef_search: usize, - query_bits: usize, - roaring_filter: *const u8, - roaring_filter_len: usize, - out_ids: *mut i64, - out_distances: *mut f32, - result_len: usize, -) -> c_int { - ffi_status(|| { - let handle = unsafe { reader_mut(handle) }?; - let query = unsafe { const_slice(query, handle.inner.dimension(), "query") }?; - let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; - let params = - VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); - let (ids, distances) = handle - .inner - .search_with_roaring_filter(query, params, filter) - .map_err(|e| format!("search_with_roaring_filter: {}", e))?; - copy_search_result(&ids, &distances, out_ids, out_distances, result_len, top_k) - }) -} - -#[no_mangle] -pub unsafe extern "C" fn paimon_vindex_reader_search_batch( - handle: *mut PaimonVindexReaderHandle, - queries: *const f32, - query_count: usize, - top_k: usize, - nprobe: usize, - ef_search: usize, - out_ids: *mut i64, - out_distances: *mut f32, - result_len: usize, -) -> c_int { - ffi_status(|| { - let handle = unsafe { reader_mut(handle) }?; - let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; - let queries = unsafe { const_slice(queries, query_len, "queries") }?; - let expected_len = checked_len(query_count, top_k, "batch result")?; - let params = VectorSearchParams::with_ef_search(top_k, nprobe, ef_search); - let (ids, distances) = handle - .inner - .search_batch(queries, query_count, params) - .map_err(|e| format!("search_batch: {}", e))?; copy_search_result( &ids, &distances, out_ids, out_distances, result_len, - expected_len, + params.top_k, ) }) } #[no_mangle] -pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_query_bits( +pub unsafe extern "C" fn paimon_vindex_reader_search_batch( handle: *mut PaimonVindexReaderHandle, queries: *const f32, query_count: usize, - top_k: usize, - nprobe: usize, - ef_search: usize, - query_bits: usize, + params: PaimonVindexSearchParams, out_ids: *mut i64, out_distances: *mut f32, result_len: usize, @@ -813,9 +752,8 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_query_bits( let handle = unsafe { reader_mut(handle) }?; let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; let queries = unsafe { const_slice(queries, query_len, "queries") }?; - let expected_len = checked_len(query_count, top_k, "batch result")?; - let params = - VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let params = search_params_from_ffi(params); + let expected_len = checked_len(query_count, params.top_k, "batch result")?; let (ids, distances) = handle .inner .search_batch(queries, query_count, params) @@ -836,46 +774,7 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter( handle: *mut PaimonVindexReaderHandle, queries: *const f32, query_count: usize, - top_k: usize, - nprobe: usize, - ef_search: usize, - roaring_filter: *const u8, - roaring_filter_len: usize, - out_ids: *mut i64, - out_distances: *mut f32, - result_len: usize, -) -> c_int { - ffi_status(|| { - let handle = unsafe { reader_mut(handle) }?; - let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; - let queries = unsafe { const_slice(queries, query_len, "queries") }?; - let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; - let expected_len = checked_len(query_count, top_k, "batch result")?; - let params = VectorSearchParams::with_ef_search(top_k, nprobe, ef_search); - let (ids, distances) = handle - .inner - .search_batch_with_roaring_filter(queries, query_count, params, filter) - .map_err(|e| format!("search_batch_with_roaring_filter: {}", e))?; - copy_search_result( - &ids, - &distances, - out_ids, - out_distances, - result_len, - expected_len, - ) - }) -} - -#[no_mangle] -pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( - handle: *mut PaimonVindexReaderHandle, - queries: *const f32, - query_count: usize, - top_k: usize, - nprobe: usize, - ef_search: usize, - query_bits: usize, + params: PaimonVindexSearchParams, roaring_filter: *const u8, roaring_filter_len: usize, out_ids: *mut i64, @@ -887,9 +786,8 @@ pub unsafe extern "C" fn paimon_vindex_reader_search_batch_with_roaring_filter_a let query_len = checked_len(query_count, handle.inner.dimension(), "queries")?; let queries = unsafe { const_slice(queries, query_len, "queries") }?; let filter = unsafe { const_slice(roaring_filter, roaring_filter_len, "roaring_filter") }?; - let expected_len = checked_len(query_count, top_k, "batch result")?; - let params = - VectorSearchParams::with_ef_search_and_query_bits(top_k, nprobe, ef_search, query_bits); + let params = search_params_from_ffi(params); + let expected_len = checked_len(query_count, params.top_k, "batch result")?; let (ids, distances) = handle .inner .search_batch_with_roaring_filter(queries, query_count, params, filter) diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp index 33e105c..8584f09 100644 --- a/include/paimon_vindex.hpp +++ b/include/paimon_vindex.hpp @@ -114,6 +114,25 @@ struct SearchResult { std::vector distances; }; +struct SearchParams { + size_t top_k = 0; + size_t nprobe = 0; + size_t ef_search = 0; + size_t query_bits = 0; + + SearchParams(size_t top_k, size_t nprobe, size_t ef_search = 0, size_t query_bits = 0) + : top_k(top_k), nprobe(nprobe), ef_search(ef_search), query_bits(query_bits) {} + + PaimonVindexSearchParams to_ffi() const { + PaimonVindexSearchParams params; + params.top_k = top_k; + params.nprobe = nprobe; + params.ef_search = ef_search; + params.query_bits = query_bits; + return params; + } +}; + class Training { public: explicit Training(PaimonVindexTrainingHandle* handle = nullptr) : handle_(handle) {} @@ -354,83 +373,53 @@ class Reader { check(paimon_vindex_reader_optimize_for_search(handle_)); } - SearchResult search( - const float* query, - size_t top_k, - size_t nprobe, - size_t ef_search = 0, - size_t query_bits = 0) { + SearchResult search(const float* query, SearchParams params) { SearchResult result; - result.ids.resize(top_k); - result.distances.resize(top_k); - check(paimon_vindex_reader_search_with_query_bits( + result.ids.resize(params.top_k); + result.distances.resize(params.top_k); + check(paimon_vindex_reader_search( handle_, query, - top_k, - nprobe, - ef_search, - query_bits, + params.to_ffi(), result.ids.data(), result.distances.data(), - top_k)); + params.top_k)); return result; } SearchResult search_with_roaring_filter( const float* query, - size_t top_k, - size_t nprobe, - size_t ef_search, - const uint8_t* filter, - size_t filter_len) { - return search_with_roaring_filter(query, top_k, nprobe, ef_search, 0, filter, filter_len); - } - - SearchResult search_with_roaring_filter( - const float* query, - size_t top_k, - size_t nprobe, - size_t ef_search, - size_t query_bits, + SearchParams params, const uint8_t* filter, size_t filter_len) { SearchResult result; - result.ids.resize(top_k); - result.distances.resize(top_k); - check(paimon_vindex_reader_search_with_roaring_filter_and_query_bits( + result.ids.resize(params.top_k); + result.distances.resize(params.top_k); + check(paimon_vindex_reader_search_with_roaring_filter( handle_, query, - top_k, - nprobe, - ef_search, - query_bits, + params.to_ffi(), filter, filter_len, result.ids.data(), result.distances.data(), - top_k)); + params.top_k)); return result; } SearchResult search_batch( const float* queries, size_t query_count, - size_t top_k, - size_t nprobe, - size_t ef_search = 0, - size_t query_bits = 0) { - const size_t result_len = query_count * top_k; + SearchParams params) { + const size_t result_len = query_count * params.top_k; SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch_with_query_bits( + check(paimon_vindex_reader_search_batch( handle_, queries, query_count, - top_k, - nprobe, - ef_search, - query_bits, + params.to_ffi(), result.ids.data(), result.distances.data(), result_len)); @@ -440,36 +429,18 @@ class Reader { SearchResult search_batch_with_roaring_filter( const float* queries, size_t query_count, - size_t top_k, - size_t nprobe, - size_t ef_search, - const uint8_t* filter, - size_t filter_len) { - return search_batch_with_roaring_filter( - queries, query_count, top_k, nprobe, ef_search, 0, filter, filter_len); - } - - SearchResult search_batch_with_roaring_filter( - const float* queries, - size_t query_count, - size_t top_k, - size_t nprobe, - size_t ef_search, - size_t query_bits, + SearchParams params, const uint8_t* filter, size_t filter_len) { - const size_t result_len = query_count * top_k; + const size_t result_len = query_count * params.top_k; SearchResult result; result.ids.resize(result_len); result.distances.resize(result_len); - check(paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( + check(paimon_vindex_reader_search_batch_with_roaring_filter( handle_, queries, query_count, - top_k, - nprobe, - ef_search, - query_bits, + params.to_ffi(), filter, filter_len, result.ids.data(), diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java index d4d4445..49edc6d 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java @@ -49,52 +49,19 @@ private VectorIndexNative() {} static native void optimizeForSearch(long ptr); - static native VectorSearchResult search(long ptr, float[] query, int k, int nprobe, int efSearch); - - static native VectorSearchResult searchWithQueryBits( - long ptr, float[] query, int k, int nprobe, int efSearch, int queryBits); + static native VectorSearchResult search(long ptr, float[] query, VectorSearchParams params); static native VectorSearchResult searchWithRoaringFilter( - long ptr, float[] query, int k, int nprobe, int efSearch, byte[] roaringFilter); - - static native VectorSearchResult searchWithRoaringFilterAndQueryBits( - long ptr, - float[] query, - int k, - int nprobe, - int efSearch, - int queryBits, - byte[] roaringFilter); + long ptr, float[] query, VectorSearchParams params, byte[] roaringFilter); static native VectorSearchBatchResult searchBatch( - long ptr, float[] queries, int queryCount, int k, int nprobe, int efSearch); - - static native VectorSearchBatchResult searchBatchWithQueryBits( - long ptr, - float[] queries, - int queryCount, - int k, - int nprobe, - int efSearch, - int queryBits); + long ptr, float[] queries, int queryCount, VectorSearchParams params); static native VectorSearchBatchResult searchBatchWithRoaringFilter( long ptr, float[] queries, int queryCount, - int k, - int nprobe, - int efSearch, - byte[] roaringFilter); - - static native VectorSearchBatchResult searchBatchWithRoaringFilterAndQueryBits( - long ptr, - float[] queries, - int queryCount, - int k, - int nprobe, - int efSearch, - int queryBits, + VectorSearchParams params, byte[] roaringFilter); static native void freeReader(long ptr); diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java index 847aae5..91de525 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexReader.java @@ -77,53 +77,31 @@ public void optimizeForSearch() { } } - public VectorSearchResult search(float[] query, int topK, int nprobe) { - return search(query, topK, nprobe, 0); - } - - public VectorSearchResult search(float[] query, int topK, int nprobe, int efSearch) { - return search(query, topK, nprobe, efSearch, 0); - } - - public VectorSearchResult search( - float[] query, int topK, int nprobe, int efSearch, int queryBits) { + public VectorSearchResult search(float[] query, VectorSearchParams params) { validateQuery(query); + validateParams(params); synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchWithQueryBits( - requireOpen(), query, topK, nprobe, efSearch, queryBits); + return VectorIndexNative.search(requireOpen(), query, params); } finally { exitNativeHandle(); } } } - public VectorSearchResult search(float[] query, int topK, int nprobe, byte[] roaringFilter) { - return search(query, topK, nprobe, 0, roaringFilter); - } - public VectorSearchResult search( - float[] query, int topK, int nprobe, int efSearch, byte[] roaringFilter) { - return search(query, topK, nprobe, efSearch, 0, roaringFilter); - } - - public VectorSearchResult search( - float[] query, - int topK, - int nprobe, - int efSearch, - int queryBits, - byte[] roaringFilter) { + float[] query, VectorSearchParams params, byte[] roaringFilter) { validateQuery(query); + validateParams(params); if (roaringFilter == null) { throw new NullPointerException("roaringFilter"); } synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchWithRoaringFilterAndQueryBits( - requireOpen(), query, topK, nprobe, efSearch, queryBits, roaringFilter); + return VectorIndexNative.searchWithRoaringFilter( + requireOpen(), query, params, roaringFilter); } finally { exitNativeHandle(); } @@ -131,25 +109,15 @@ public VectorSearchResult search( } public VectorSearchBatchResult searchBatch( - float[] queries, int queryCount, int topK, int nprobe) { - return searchBatch(queries, queryCount, topK, nprobe, 0); - } - - public VectorSearchBatchResult searchBatch( - float[] queries, int queryCount, int topK, int nprobe, int efSearch) { - return searchBatch(queries, queryCount, topK, nprobe, efSearch, 0); - } - - public VectorSearchBatchResult searchBatch( - float[] queries, int queryCount, int topK, int nprobe, int efSearch, int queryBits) { + float[] queries, int queryCount, VectorSearchParams params) { if (queries == null) { throw new NullPointerException("queries"); } + validateParams(params); synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchBatchWithQueryBits( - requireOpen(), queries, queryCount, topK, nprobe, efSearch, queryBits); + return VectorIndexNative.searchBatch(requireOpen(), queries, queryCount, params); } finally { exitNativeHandle(); } @@ -157,46 +125,19 @@ public VectorSearchBatchResult searchBatch( } public VectorSearchBatchResult searchBatch( - float[] queries, int queryCount, int topK, int nprobe, byte[] roaringFilter) { - return searchBatch(queries, queryCount, topK, nprobe, 0, roaringFilter); - } - - public VectorSearchBatchResult searchBatch( - float[] queries, - int queryCount, - int topK, - int nprobe, - int efSearch, - byte[] roaringFilter) { - return searchBatch(queries, queryCount, topK, nprobe, efSearch, 0, roaringFilter); - } - - public VectorSearchBatchResult searchBatch( - float[] queries, - int queryCount, - int topK, - int nprobe, - int efSearch, - int queryBits, - byte[] roaringFilter) { + float[] queries, int queryCount, VectorSearchParams params, byte[] roaringFilter) { if (queries == null) { throw new NullPointerException("queries"); } + validateParams(params); if (roaringFilter == null) { throw new NullPointerException("roaringFilter"); } synchronized (nativeHandleLock) { enterNativeHandle(); try { - return VectorIndexNative.searchBatchWithRoaringFilterAndQueryBits( - requireOpen(), - queries, - queryCount, - topK, - nprobe, - efSearch, - queryBits, - roaringFilter); + return VectorIndexNative.searchBatchWithRoaringFilter( + requireOpen(), queries, queryCount, params, roaringFilter); } finally { exitNativeHandle(); } @@ -225,6 +166,12 @@ private void validateQuery(float[] query) { } } + private void validateParams(VectorSearchParams params) { + if (params == null) { + throw new NullPointerException("params"); + } + } + private long requireOpen() { if (nativePtr == 0L) { throw new IllegalStateException("VectorIndexReader is closed"); diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java b/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java new file mode 100644 index 0000000..b0de1fe --- /dev/null +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorSearchParams.java @@ -0,0 +1,61 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.paimon.index.vector; + +public final class VectorSearchParams { + + private final int topK; + private final int nprobe; + private final int efSearch; + private final int queryBits; + + public VectorSearchParams(int topK, int nprobe) { + this(topK, nprobe, 0, 0); + } + + public VectorSearchParams(int topK, int nprobe, int efSearch, int queryBits) { + this.topK = topK; + this.nprobe = nprobe; + this.efSearch = efSearch; + this.queryBits = queryBits; + } + + public int topK() { + return topK; + } + + public int nprobe() { + return nprobe; + } + + public int efSearch() { + return efSearch; + } + + public int queryBits() { + return queryBits; + } + + public VectorSearchParams withEfSearch(int efSearch) { + return new VectorSearchParams(topK, nprobe, efSearch, queryBits); + } + + public VectorSearchParams withQueryBits(int queryBits) { + return new VectorSearchParams(topK, nprobe, efSearch, queryBits); + } +} diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java index fb3a406..cbe37cc 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java @@ -139,13 +139,13 @@ public void run() { assertThrows(IllegalStateException.class, new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {0.0f}, 1, 1); + reader.search(new float[] {0.0f}, new VectorSearchParams(1, 1)); } }); assertThrows(IllegalStateException.class, new ThrowingRunnable() { @Override public void run() { - reader.searchBatch(new float[] {0.0f}, 1, 1, 1); + reader.searchBatch(new float[] {0.0f}, 1, new VectorSearchParams(1, 1)); } }); } @@ -230,25 +230,25 @@ private static void testReaderAndWriterApiCompile() { reader.dimension(); reader.totalVectors(); reader.optimizeForSearch(); - reader.search(new float[] {0.0f, 1.0f}, 10, 4); - reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32); - reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, 4); - reader.search(new float[] {0.0f, 1.0f}, 10, 4, new byte[] {1, 2}); - reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, new byte[] {1, 2}); - reader.search(new float[] {0.0f, 1.0f}, 10, 4, 32, 4, new byte[] {1, 2}); - reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4); - reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32); - reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32, 4); - reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, new byte[] {1, 2}); + VectorSearchParams params = new VectorSearchParams(10, 4); + VectorSearchParams hnswParams = params.withEfSearch(32); + VectorSearchParams rqParams = hnswParams.withQueryBits(4); + reader.search(new float[] {0.0f, 1.0f}, params); + reader.search(new float[] {0.0f, 1.0f}, hnswParams); + reader.search(new float[] {0.0f, 1.0f}, rqParams); + reader.search(new float[] {0.0f, 1.0f}, params, new byte[] {1, 2}); + reader.search(new float[] {0.0f, 1.0f}, hnswParams, new byte[] {1, 2}); + reader.search(new float[] {0.0f, 1.0f}, rqParams, new byte[] {1, 2}); + reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, params); + reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, hnswParams); + reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, rqParams); + reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, params, new byte[] {1, 2}); reader.searchBatch( - new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32, new byte[] {1, 2}); + new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, hnswParams, new byte[] {1, 2}); reader.searchBatch( new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, - 10, - 4, - 32, - 4, + rqParams, new byte[] {1, 2}); VectorIndexTraining training = diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java index 3609cdd..563205c 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java @@ -53,7 +53,7 @@ private static void testReaderRejectsReentrantCloseDuringNativeCall() { VectorIndexReader reader = new VectorIndexReader(input); input.setReader(reader); try { - VectorSearchResult result = reader.search(new float[] {0.0f}, 1, 1); + VectorSearchResult result = reader.search(new float[] {0.0f}, new VectorSearchParams(1, 1)); assertEquals(1, result.ids().length); assertTrue(input.closeAttempted()); assertTrue(input.closeRejected()); diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java index 963b26f..ec0273f 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java @@ -75,7 +75,7 @@ private static void testObjectEntrypointPanicBecomesRuntimeException() { assertThrows(RuntimeException.class, new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {0.0f}, 2, 1); + reader.search(new float[] {0.0f}, new VectorSearchParams(2, 1)); } }); assertEquals(2L, reader.totalVectors()); diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java index 3ee2164..d8f95c0 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java @@ -102,7 +102,7 @@ private static void testReaderValidationComesFromCore() { new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {0.0f, 1.0f}, 1, 1); + reader.search(new float[] {0.0f, 1.0f}, new VectorSearchParams(1, 1)); } }); assertThrowsMessage( @@ -111,7 +111,7 @@ public void run() { new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {0.0f}, 0, 1); + reader.search(new float[] {0.0f}, new VectorSearchParams(0, 1)); } }); assertThrowsMessage( @@ -120,7 +120,8 @@ public void run() { new ThrowingRunnable() { @Override public void run() { - reader.searchBatch(new float[] {0.0f, 1.0f}, 1, 1, 1); + reader.searchBatch( + new float[] {0.0f, 1.0f}, 1, new VectorSearchParams(1, 1)); } }); } finally { @@ -276,7 +277,7 @@ private static void testReaderRejectsNonFiniteQueries() { new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {Float.NaN}, 1, 1); + reader.search(new float[] {Float.NaN}, new VectorSearchParams(1, 1)); } }, "query contains non-finite value at offset 0: NaN"); @@ -284,7 +285,10 @@ public void run() { new ThrowingRunnable() { @Override public void run() { - reader.searchBatch(new float[] {Float.NEGATIVE_INFINITY}, 1, 1, 1); + reader.searchBatch( + new float[] {Float.NEGATIVE_INFINITY}, + 1, + new VectorSearchParams(1, 1)); } }, "queries contains non-finite value at offset 0: -inf"); @@ -292,7 +296,10 @@ public void run() { new ThrowingRunnable() { @Override public void run() { - reader.search(new float[] {Float.NaN}, 1, 1, new byte[] {(byte) 0xFF}); + reader.search( + new float[] {Float.NaN}, + new VectorSearchParams(1, 1), + new byte[] {(byte) 0xFF}); } }, "query contains non-finite value at offset 0: NaN"); @@ -349,12 +356,13 @@ private static void assertRoundtrip( reader.optimizeForSearch(); - VectorSearchResult single = reader.search(queryForCenter(0.0f), 2, 4, 16); + VectorSearchParams params = new VectorSearchParams(2, 4, 16, 0); + VectorSearchResult single = reader.search(queryForCenter(0.0f), params); assertIdInCluster(single.ids()[0], 0); assertFinite(single.distances()[0], indexType + " single distance"); VectorSearchBatchResult batch = - reader.searchBatch(batchQueries(), 2, 1, 4, 16); + reader.searchBatch(batchQueries(), 2, new VectorSearchParams(1, 4, 16, 0)); assertIdInCluster(batch.ids()[0], 0); assertIdInCluster(batch.ids()[1], 1); assertFinite(batch.distances()[0], indexType + " batch distance 0"); @@ -362,12 +370,13 @@ private static void assertRoundtrip( if ("ivf_rq".equals(indexType)) { VectorSearchResult queryBitsSingle = - reader.search(queryForCenter(0.0f), 2, 4, 16, 4); + reader.search(queryForCenter(0.0f), params.withQueryBits(4)); assertIdInCluster(queryBitsSingle.ids()[0], 0); assertFinite(queryBitsSingle.distances()[0], "ivf_rq queryBits single distance"); VectorSearchBatchResult queryBitsBatch = - reader.searchBatch(batchQueries(), 2, 1, 4, 16, 8); + reader.searchBatch( + batchQueries(), 2, new VectorSearchParams(1, 4, 16, 8)); assertIdInCluster(queryBitsBatch.ids()[0], 0); assertIdInCluster(queryBitsBatch.ids()[1], 1); @@ -377,7 +386,7 @@ private static void assertRoundtrip( new ThrowingRunnable() { @Override public void run() { - reader.search(queryForCenter(0.0f), 2, 4, 16, 7); + reader.search(queryForCenter(0.0f), params.withQueryBits(7)); } }); } diff --git a/jni/src/lib.rs b/jni/src/lib.rs index 944bc93..cfc0ecc 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -369,22 +369,32 @@ fn build_metadata(env: &mut JNIEnv, metadata: VectorIndexMetadata) -> jobject { result.into_raw() } -fn search_params( - k: jint, - nprobe: jint, - ef_search: jint, - query_bits: jint, -) -> Option { - if k < 0 || nprobe < 0 || ef_search < 0 || query_bits < 0 { - None - } else { - Some(VectorSearchParams::with_ef_search_and_query_bits( - k as usize, - nprobe as usize, - ef_search as usize, - query_bits as usize, - )) +fn search_params(env: &mut JNIEnv, params: JObject) -> Result { + if params.is_null() { + return Err("params is null".to_string()); + } + let top_k = call_int_method(env, ¶ms, "topK")?; + let nprobe = call_int_method(env, ¶ms, "nprobe")?; + let ef_search = call_int_method(env, ¶ms, "efSearch")?; + let query_bits = call_int_method(env, ¶ms, "queryBits")?; + if top_k < 0 || nprobe < 0 || ef_search < 0 || query_bits < 0 { + return Err(format!( + "invalid search parameters: topK={}, nprobe={}, efSearch={}, queryBits={}", + top_k, nprobe, ef_search, query_bits + )); } + Ok(VectorSearchParams { + top_k: top_k as usize, + nprobe: nprobe as usize, + ef_search: ef_search as usize, + query_bits: query_bits as usize, + }) +} + +fn call_int_method(env: &mut JNIEnv, object: &JObject, name: &str) -> Result { + env.call_method(object, name, "()I", &[]) + .and_then(|value| value.i()) + .map_err(|e| format!("VectorSearchParams.{}(): {}", name, e)) } // --- Unified Trainer / Writer API --- @@ -690,67 +700,17 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea _class: JClass, ptr: jlong, query: JFloatArray, - k: jint, - nprobe: jint, - ef_search: jint, + params: JObject, ) -> jobject { jni_call(env, |env| { let reader = match deref_reader(ptr) { Some(reader) => reader, None => return throw_and_return(env, "null native pointer (reader already freed?)"), }; - let params = match search_params(k, nprobe, ef_search, 0) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}", - k, nprobe, ef_search - ), - ) - } - }; - let query_buf = match read_float_array(env, &query, "query") { - Ok(buf) => buf, + let params = match search_params(env, params) { + Ok(params) => params, Err(e) => return throw_and_return(env, &e), }; - let (ids, dists) = match reader.search(&query_buf, params) { - Ok(result) => result, - Err(e) => return throw_and_return(env, &format!("search: {}", e)), - }; - build_result(env, ids, dists) - }) -} - -#[no_mangle] -pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchWithQueryBits( - env: JNIEnv, - _class: JClass, - ptr: jlong, - query: JFloatArray, - k: jint, - nprobe: jint, - ef_search: jint, - query_bits: jint, -) -> jobject { - jni_call(env, |env| { - let reader = match deref_reader(ptr) { - Some(reader) => reader, - None => return throw_and_return(env, "null native pointer (reader already freed?)"), - }; - let params = match search_params(k, nprobe, ef_search, query_bits) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", - k, nprobe, ef_search, query_bits - ), - ) - } - }; let query_buf = match read_float_array(env, &query, "query") { Ok(buf) => buf, Err(e) => return throw_and_return(env, &e), @@ -769,9 +729,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea _class: JClass, ptr: jlong, query: JFloatArray, - k: jint, - nprobe: jint, - ef_search: jint, + params: JObject, roaring_filter: JByteArray, ) -> jobject { jni_call(env, |env| { @@ -779,64 +737,10 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea Some(reader) => reader, None => return throw_and_return(env, "null native pointer (reader already freed?)"), }; - let params = match search_params(k, nprobe, ef_search, 0) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}", - k, nprobe, ef_search - ), - ) - } - }; - let query_buf = match read_float_array(env, &query, "query") { - Ok(buf) => buf, - Err(e) => return throw_and_return(env, &e), - }; - let filter_bytes = match read_byte_array(env, roaring_filter) { - Ok(bytes) => bytes, + let params = match search_params(env, params) { + Ok(params) => params, Err(e) => return throw_and_return(env, &e), }; - let (ids, dists) = - match reader.search_with_roaring_filter(&query_buf, params, &filter_bytes) { - Ok(result) => result, - Err(e) => return throw_and_return(env, &format!("search_with_filter: {}", e)), - }; - build_result(env, ids, dists) - }) -} - -#[no_mangle] -pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchWithRoaringFilterAndQueryBits( - env: JNIEnv, - _class: JClass, - ptr: jlong, - query: JFloatArray, - k: jint, - nprobe: jint, - ef_search: jint, - query_bits: jint, - roaring_filter: JByteArray, -) -> jobject { - jni_call(env, |env| { - let reader = match deref_reader(ptr) { - Some(reader) => reader, - None => return throw_and_return(env, "null native pointer (reader already freed?)"), - }; - let params = match search_params(k, nprobe, ef_search, query_bits) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", - k, nprobe, ef_search, query_bits - ), - ) - } - }; let query_buf = match read_float_array(env, &query, "query") { Ok(buf) => buf, Err(e) => return throw_and_return(env, &e), @@ -861,9 +765,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea ptr: jlong, queries: JFloatArray, query_count: jint, - k: jint, - nprobe: jint, - ef_search: jint, + params: JObject, ) -> jobject { jni_call(env, |env| { let reader = match deref_reader(ptr) { @@ -873,63 +775,10 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea if query_count < 0 { return throw_and_return(env, &format!("invalid query count: {}", query_count)); } - let params = match search_params(k, nprobe, ef_search, 0) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}", - k, nprobe, ef_search - ), - ) - } - }; - let nq = query_count as usize; - let query_buf = match read_float_array(env, &queries, "queries") { - Ok(buf) => buf, + let params = match search_params(env, params) { + Ok(params) => params, Err(e) => return throw_and_return(env, &e), }; - let (ids, dists) = match reader.search_batch(&query_buf, nq, params) { - Ok(result) => result, - Err(e) => return throw_and_return(env, &format!("search_batch: {}", e)), - }; - build_batch_result(env, ids, dists, nq, params.top_k) - }) -} - -#[no_mangle] -pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatchWithQueryBits( - env: JNIEnv, - _class: JClass, - ptr: jlong, - queries: JFloatArray, - query_count: jint, - k: jint, - nprobe: jint, - ef_search: jint, - query_bits: jint, -) -> jobject { - jni_call(env, |env| { - let reader = match deref_reader(ptr) { - Some(reader) => reader, - None => return throw_and_return(env, "null native pointer (reader already freed?)"), - }; - if query_count < 0 { - return throw_and_return(env, &format!("invalid query count: {}", query_count)); - } - let params = match search_params(k, nprobe, ef_search, query_bits) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", - k, nprobe, ef_search, query_bits - ), - ) - } - }; let nq = query_count as usize; let query_buf = match read_float_array(env, &queries, "queries") { Ok(buf) => buf, @@ -950,9 +799,7 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea ptr: jlong, queries: JFloatArray, query_count: jint, - k: jint, - nprobe: jint, - ef_search: jint, + params: JObject, roaring_filter: JByteArray, ) -> jobject { jni_call(env, |env| { @@ -963,71 +810,10 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_sea if query_count < 0 { return throw_and_return(env, &format!("invalid query count: {}", query_count)); } - let params = match search_params(k, nprobe, ef_search, 0) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}", - k, nprobe, ef_search - ), - ) - } - }; - let nq = query_count as usize; - let query_buf = match read_float_array(env, &queries, "queries") { - Ok(buf) => buf, + let params = match search_params(env, params) { + Ok(params) => params, Err(e) => return throw_and_return(env, &e), }; - let filter_bytes = match read_byte_array(env, roaring_filter) { - Ok(bytes) => bytes, - Err(e) => return throw_and_return(env, &e), - }; - let (ids, dists) = - match reader.search_batch_with_roaring_filter(&query_buf, nq, params, &filter_bytes) { - Ok(result) => result, - Err(e) => { - return throw_and_return(env, &format!("search_batch_with_filter: {}", e)) - } - }; - build_batch_result(env, ids, dists, nq, params.top_k) - }) -} - -#[no_mangle] -pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_searchBatchWithRoaringFilterAndQueryBits( - env: JNIEnv, - _class: JClass, - ptr: jlong, - queries: JFloatArray, - query_count: jint, - k: jint, - nprobe: jint, - ef_search: jint, - query_bits: jint, - roaring_filter: JByteArray, -) -> jobject { - jni_call(env, |env| { - let reader = match deref_reader(ptr) { - Some(reader) => reader, - None => return throw_and_return(env, "null native pointer (reader already freed?)"), - }; - if query_count < 0 { - return throw_and_return(env, &format!("invalid query count: {}", query_count)); - } - let params = match search_params(k, nprobe, ef_search, query_bits) { - Some(params) => params, - None => { - return throw_and_return( - env, - &format!( - "invalid search parameters: k={}, nprobe={}, efSearch={}, queryBits={}", - k, nprobe, ef_search, query_bits - ), - ) - } - }; let nq = query_count as usize; let query_buf = match read_float_array(env, &queries, "queries") { Ok(buf) => buf, diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py index ea10ed4..4bdd4a3 100644 --- a/python/paimon_vindex/__init__.py +++ b/python/paimon_vindex/__init__.py @@ -52,6 +52,22 @@ class VectorIndexMetadata: hnsw_max_level: Optional[int] = None +@dataclass(frozen=True) +class SearchParams: + top_k: int + nprobe: int + ef_search: int = 0 + query_bits: int = 0 + + def to_ffi(self): + return _ffi.PaimonVindexSearchParams( + self.top_k, + self.nprobe, + self.ef_search, + self.query_bits, + ) + + def _check_error(message="operation failed"): err = lib.paimon_vindex_last_error() if err: @@ -428,7 +444,7 @@ def _filter_args(self, filter_bytes): return None, 0, None return _bytes_buffer(filter_bytes, "filter_bytes") - def search(self, query, top_k, nprobe, ef_search=0, query_bits=0, filter_bytes=None): + def search(self, query, params: SearchParams, filter_bytes=None): self._require_open() query = _float32_vector(query, "query") if query.shape[0] != self._metadata.dimension: @@ -436,41 +452,36 @@ def search(self, query, top_k, nprobe, ef_search=0, query_bits=0, filter_bytes=N f"query length {query.shape[0]} does not match index dimension " f"{self._metadata.dimension}" ) - ids = np.empty(top_k, dtype=np.int64) - distances = np.empty(top_k, dtype=np.float32) + ffi_params = params.to_ffi() + ids = np.empty(params.top_k, dtype=np.int64) + distances = np.empty(params.top_k, dtype=np.float32) if filter_bytes is None: - rc = lib.paimon_vindex_reader_search_with_query_bits( + rc = lib.paimon_vindex_reader_search( self._handle, query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - top_k, - nprobe, - ef_search, - query_bits, + ffi_params, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - top_k, + params.top_k, ) else: filter_buf, filter_len, _ = self._filter_args(filter_bytes) - rc = lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits( + rc = lib.paimon_vindex_reader_search_with_roaring_filter( self._handle, query.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - top_k, - nprobe, - ef_search, - query_bits, + ffi_params, filter_buf, filter_len, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - top_k, + params.top_k, ) if rc != 0: _check_error("search failed") return ids, distances - def search_batch(self, queries, top_k, nprobe, ef_search=0, query_bits=0, filter_bytes=None): + def search_batch(self, queries, params: SearchParams, filter_bytes=None): self._require_open() queries = _float32_matrix(queries, "queries") if queries.shape[1] != self._metadata.dimension: @@ -478,33 +489,28 @@ def search_batch(self, queries, top_k, nprobe, ef_search=0, query_bits=0, filter f"queries length {queries.size} does not match nq * dimension " f"{queries.shape[0] * self._metadata.dimension}" ) - result_len = queries.shape[0] * top_k - ids = np.empty((queries.shape[0], top_k), dtype=np.int64) - distances = np.empty((queries.shape[0], top_k), dtype=np.float32) + ffi_params = params.to_ffi() + result_len = queries.shape[0] * params.top_k + ids = np.empty((queries.shape[0], params.top_k), dtype=np.int64) + distances = np.empty((queries.shape[0], params.top_k), dtype=np.float32) if filter_bytes is None: - rc = lib.paimon_vindex_reader_search_batch_with_query_bits( + rc = lib.paimon_vindex_reader_search_batch( self._handle, queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), queries.shape[0], - top_k, - nprobe, - ef_search, - query_bits, + ffi_params, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), distances.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), result_len, ) else: filter_buf, filter_len, _ = self._filter_args(filter_bytes) - rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits( + rc = lib.paimon_vindex_reader_search_batch_with_roaring_filter( self._handle, queries.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), queries.shape[0], - top_k, - nprobe, - ef_search, - query_bits, + ffi_params, filter_buf, filter_len, ids.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)), diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py index 16deebf..c084a2e 100644 --- a/python/paimon_vindex/_ffi.py +++ b/python/paimon_vindex/_ffi.py @@ -118,6 +118,15 @@ class PaimonVindexMetadata(Structure): ] +class PaimonVindexSearchParams(Structure): + _fields_ = [ + ("top_k", c_size_t), + ("nprobe", c_size_t), + ("ef_search", c_size_t), + ("query_bits", c_size_t), + ] + + lib.paimon_vindex_last_error.argtypes = [] lib.paimon_vindex_last_error.restype = c_char_p @@ -188,34 +197,17 @@ class PaimonVindexMetadata(Structure): lib.paimon_vindex_reader_search.argtypes = [ c_void_p, POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, + PaimonVindexSearchParams, POINTER(c_int64), POINTER(c_float), c_size_t, ] lib.paimon_vindex_reader_search.restype = c_int -lib.paimon_vindex_reader_search_with_query_bits.argtypes = [ - c_void_p, - POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, - c_size_t, - POINTER(c_int64), - POINTER(c_float), - c_size_t, -] -lib.paimon_vindex_reader_search_with_query_bits.restype = c_int - lib.paimon_vindex_reader_search_with_roaring_filter.argtypes = [ c_void_p, POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, + PaimonVindexSearchParams, POINTER(c_uint8), c_size_t, POINTER(c_int64), @@ -224,55 +216,22 @@ class PaimonVindexMetadata(Structure): ] lib.paimon_vindex_reader_search_with_roaring_filter.restype = c_int -lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits.argtypes = [ - c_void_p, - POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, - c_size_t, - POINTER(c_uint8), - c_size_t, - POINTER(c_int64), - POINTER(c_float), - c_size_t, -] -lib.paimon_vindex_reader_search_with_roaring_filter_and_query_bits.restype = c_int - lib.paimon_vindex_reader_search_batch.argtypes = [ c_void_p, POINTER(c_float), c_size_t, - c_size_t, - c_size_t, - c_size_t, + PaimonVindexSearchParams, POINTER(c_int64), POINTER(c_float), c_size_t, ] lib.paimon_vindex_reader_search_batch.restype = c_int -lib.paimon_vindex_reader_search_batch_with_query_bits.argtypes = [ - c_void_p, - POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, - c_size_t, - c_size_t, - POINTER(c_int64), - POINTER(c_float), - c_size_t, -] -lib.paimon_vindex_reader_search_batch_with_query_bits.restype = c_int - lib.paimon_vindex_reader_search_batch_with_roaring_filter.argtypes = [ c_void_p, POINTER(c_float), c_size_t, - c_size_t, - c_size_t, - c_size_t, + PaimonVindexSearchParams, POINTER(c_uint8), c_size_t, POINTER(c_int64), @@ -280,19 +239,3 @@ class PaimonVindexMetadata(Structure): c_size_t, ] lib.paimon_vindex_reader_search_batch_with_roaring_filter.restype = c_int - -lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits.argtypes = [ - c_void_p, - POINTER(c_float), - c_size_t, - c_size_t, - c_size_t, - c_size_t, - c_size_t, - POINTER(c_uint8), - c_size_t, - POINTER(c_int64), - POINTER(c_float), - c_size_t, -] -lib.paimon_vindex_reader_search_batch_with_roaring_filter_and_query_bits.restype = c_int diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py index 9dae5ee..e4c0f94 100644 --- a/python/tests/test_vindex.py +++ b/python/tests/test_vindex.py @@ -20,7 +20,7 @@ import numpy as np import pytest -from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter +from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter class VectorIndexInput: @@ -116,11 +116,10 @@ def test_python_ffi_roundtrips_supported_indexes(): assert reader.dimension == d assert metadata.total_vectors == 512 - ids, distances = reader.search(data[0], top_k=5, nprobe=4, ef_search=32) + params = SearchParams(top_k=5, nprobe=4, ef_search=32) + ids, distances = reader.search(data[0], params) reader.optimize_for_search() - optimized_ids, optimized_distances = reader.search( - data[0], top_k=5, nprobe=4, ef_search=32 - ) + optimized_ids, optimized_distances = reader.search(data[0], params) assert ids.shape == (5,) assert distances.shape == (5,) assert ids[0] == 0 @@ -143,8 +142,7 @@ def test_python_ffi_batch_search(): with reader_from_bytes(index_bytes) as reader: ids, distances = reader.search_batch( np.vstack([data[0], data[1]]), - top_k=2, - nprobe=2, + SearchParams(top_k=2, nprobe=2), ) assert ids.shape == (2, 2) assert distances.shape == (2, 2) @@ -167,20 +165,20 @@ def test_python_ffi_ivfrq_query_bits(): with reader_from_bytes(index_bytes) as reader: for query_bits in (4, 8): ids, distances = reader.search( - data[7], top_k=5, nprobe=4, query_bits=query_bits + data[7], SearchParams(top_k=5, nprobe=4, query_bits=query_bits) ) assert ids.shape == (5,) assert distances.shape == (5,) assert ids[0] % 4 == 7 % 4 ids, distances = reader.search_batch( - np.vstack([data[4], data[7]]), top_k=5, nprobe=4, query_bits=4 + np.vstack([data[4], data[7]]), SearchParams(top_k=5, nprobe=4, query_bits=4) ) assert ids[0, 0] % 4 == 4 % 4 assert ids[1, 0] % 4 == 7 % 4 with pytest.raises(RuntimeError, match="query_bits"): - reader.search(data[0], top_k=5, nprobe=4, query_bits=7) + reader.search(data[0], SearchParams(top_k=5, nprobe=4, query_bits=7)) def test_python_ffi_delegates_validation(): @@ -205,8 +203,8 @@ def test_python_ffi_delegates_validation(): index_bytes, data = build_index(options, 16) with reader_from_bytes(index_bytes) as reader: with pytest.raises(RuntimeError, match="query length 15"): - reader.search(np.zeros(15, dtype=np.float32), top_k=5, nprobe=2) + reader.search(np.zeros(15, dtype=np.float32), SearchParams(top_k=5, nprobe=2)) with pytest.raises(RuntimeError, match="k must be greater than 0"): - reader.search(data[0], top_k=0, nprobe=2) + reader.search(data[0], SearchParams(top_k=0, nprobe=2)) with pytest.raises(RuntimeError, match="queries length 15"): - reader.search_batch(np.zeros((1, 15), dtype=np.float32), top_k=5, nprobe=2) + reader.search_batch(np.zeros((1, 15), dtype=np.float32), SearchParams(top_k=5, nprobe=2)) From 6317ac22090dd4814fee6452fed765462d7b3c07 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 9 Jul 2026 19:36:29 +0800 Subject: [PATCH 4/4] Update wheel smoke tests for search params --- .github/workflows/ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4dbb1bd..97e4025 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -254,7 +254,7 @@ jobs: python - <<'PY' import io import numpy as np - from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter + from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter class Input: def __init__(self, data): @@ -275,7 +275,7 @@ jobs: reader = VectorIndexReader(Input(output.getvalue())) reader.optimize_for_search() - result_ids, distances = reader.search(data[0], top_k=2, nprobe=2) + result_ids, distances = reader.search(data[0], SearchParams(top_k=2, nprobe=2)) assert result_ids[0] == 1, result_ids assert np.isfinite(distances[0]), distances reader.close() @@ -292,7 +292,7 @@ jobs: python -c @' import io import numpy as np - from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter + from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter class Input: def __init__(self, data): @@ -313,7 +313,7 @@ jobs: reader = VectorIndexReader(Input(output.getvalue())) reader.optimize_for_search() - result_ids, distances = reader.search(data[0], top_k=2, nprobe=2) + result_ids, distances = reader.search(data[0], SearchParams(top_k=2, nprobe=2)) assert result_ids[0] == 1, result_ids assert np.isfinite(distances[0]), distances reader.close()