This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new 81f5f8d Expand release test coverage (#41)
81f5f8d is described below
commit 81f5f8d1cb0053eae55eef337cd0689564775a09
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Jun 12 11:36:00 2026 +0800
Expand release test coverage (#41)
---
.github/workflows/ci.yml | 106 ++++++++
README.md | 17 ++
c/test_vindex.c | 231 ++++++++++++++--
core/Cargo.toml | 4 +
core/benches/ivfhnswsq_filter_bench.rs | 293 +++++++++++++++++++++
core/tests/storage_format_fixtures.rs | 31 +++
cpp/test_vindex.cpp | 140 ++++++++--
.../vector/VectorIndexNativeValidationTest.java | 167 +++++++++++-
8 files changed, 931 insertions(+), 58 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 07bb7b4..d1e5275 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -210,3 +210,109 @@ jobs:
pytest -v
env:
PAIMON_VINDEX_LIB_PATH: ${{ github.workspace }}/target/release
+
+ python-wheel:
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [ubuntu-latest, macos-latest, windows-latest]
+ runs-on: ${{ matrix.os }}
+ steps:
+ - uses: actions/checkout@v4
+
+ - name: Rust Cache
+ uses: actions/cache@v4
+ with:
+ path: |
+ ~/.cargo/registry
+ ~/.cargo/git
+ target
+ key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+ restore-keys: |
+ ${{ runner.os }}-cargo-
+
+ - name: Set up Python
+ uses: actions/setup-python@v5
+ with:
+ python-version: '3.12'
+
+ - name: Build FFI library
+ run: cargo build --release -p paimon-vindex-ffi
+
+ - name: Build wheel
+ working-directory: python
+ run: |
+ python -m pip install --upgrade pip build
+ python -m build --wheel
+
+ - name: Test wheel
+ if: runner.os != 'Windows'
+ working-directory: python
+ run: |
+ python -m pip install dist/*.whl
+ cd ..
+ python - <<'PY'
+ import io
+ import numpy as np
+ from paimon_vindex import VectorIndexReader, VectorIndexWriter
+
+ class Input:
+ def __init__(self, data):
+ self.data = data
+
+ def pread_many(self, ranges):
+ return [self.data[pos : pos + length] for pos, length in
ranges]
+
+ data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1,
10.0]], dtype=np.float32)
+ ids = np.array([1, 2, 3, 4], dtype=np.int64)
+ output = io.BytesIO()
+ writer = VectorIndexWriter({"index.type": "ivf_flat", "dimension":
"2", "nlist": "2", "metric": "l2"})
+ writer.train(data)
+ writer.add_vectors(ids, data)
+ writer.write(output)
+ writer.close()
+
+ reader = VectorIndexReader(Input(output.getvalue()))
+ reader.optimize_for_search()
+ result_ids, distances = reader.search(data[0], top_k=2, nprobe=2)
+ assert result_ids[0] == 1, result_ids
+ assert np.isfinite(distances[0]), distances
+ reader.close()
+ PY
+
+ - name: Test wheel
+ if: runner.os == 'Windows'
+ working-directory: python
+ shell: pwsh
+ run: |
+ $wheel = Get-ChildItem dist/*.whl | Select-Object -First 1
+ python -m pip install $wheel.FullName
+ Set-Location ..
+ python -c @'
+ import io
+ import numpy as np
+ from paimon_vindex import VectorIndexReader, VectorIndexWriter
+
+ class Input:
+ def __init__(self, data):
+ self.data = data
+
+ def pread_many(self, ranges):
+ return [self.data[pos : pos + length] for pos, length in
ranges]
+
+ data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1,
10.0]], dtype=np.float32)
+ ids = np.array([1, 2, 3, 4], dtype=np.int64)
+ output = io.BytesIO()
+ writer = VectorIndexWriter({"index.type": "ivf_flat", "dimension":
"2", "nlist": "2", "metric": "l2"})
+ writer.train(data)
+ writer.add_vectors(ids, data)
+ writer.write(output)
+ writer.close()
+
+ reader = VectorIndexReader(Input(output.getvalue()))
+ reader.optimize_for_search()
+ result_ids, distances = reader.search(data[0], top_k=2, nprobe=2)
+ assert result_ids[0] == 1, result_ids
+ assert np.isfinite(distances[0]), distances
+ reader.close()
+ '@
diff --git a/README.md b/README.md
index 3ea1782..82bd6f1 100644
--- a/README.md
+++ b/README.md
@@ -320,6 +320,23 @@ cargo bench -p paimon-vindex-core --bench ann_bench --
--nocapture
Benchmark rows report `disk_scope=index_bytes`, which is the serialized vector
index file.
+`IVF_HNSW_SQ` filtered-search fallback can be profiled separately with a
+filter-heavy benchmark. It compares filtered batch search before and after
+`optimize_for_search` and verifies that both paths return the same results:
+
+```bash
+cargo bench -p paimon-vindex-core --bench ivfhnswsq_filter_bench -- --nocapture
+```
+
+Useful knobs include:
+
+```bash
+FILTER_BENCH_N=50000 FILTER_BENCH_NQ=500 FILTER_BENCH_D=128 \
+FILTER_BENCH_NLIST=64 FILTER_BENCH_NPROBE=32 FILTER_BENCH_EF_SEARCH=80 \
+FILTER_BENCH_FILTER_STRIDES=1,4,16,64 \
+cargo bench -p paimon-vindex-core --bench ivfhnswsq_filter_bench -- --nocapture
+```
+
## Development
Common Rust commands:
diff --git a/c/test_vindex.c b/c/test_vindex.c
index 916889e..b08660e 100644
--- a/c/test_vindex.c
+++ b/c/test_vindex.c
@@ -49,12 +49,33 @@ struct MemBuffer {
size_t pos;
};
+enum {
+ ROUNDTRIP_DIMENSION = 2,
+ ROUNDTRIP_NLIST = 4,
+ ROUNDTRIP_PER_LIST = 128,
+ ROUNDTRIP_VECTOR_COUNT = ROUNDTRIP_NLIST * ROUNDTRIP_PER_LIST,
+};
+
static void fail_ffi(const char *message) {
const char *err = paimon_vindex_last_error();
fprintf(stderr, "%s: %s\n", message, err == NULL ? "(no error)" : err);
abort();
}
+static void assert_last_error_contains(const char *needle) {
+ const char *err = paimon_vindex_last_error();
+ if (err == NULL || strstr(err, needle) == NULL) {
+ fprintf(
+ stderr,
+ "FAIL %s:%d: last error should contain '%s', got '%s'\n",
+ __FILE__,
+ __LINE__,
+ needle,
+ err == NULL ? "(null)" : err);
+ abort();
+ }
+}
+
static int mem_write(void *ctx, const uint8_t *data, uintptr_t len) {
struct MemBuffer *buf = (struct MemBuffer *)ctx;
if (len > SIZE_MAX - buf->len) {
@@ -106,10 +127,57 @@ static int mem_read_at(void *ctx, uint64_t offset,
uint8_t *dst, uintptr_t len)
return 0;
}
-static void test_basic_roundtrip(void) {
- const char *keys[] = {"index.type", "dimension", "nlist", "metric"};
- const char *values[] = {"ivf_flat", "2", "2", "l2"};
- PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(keys, values,
4);
+static int failing_write(void *ctx, const uint8_t *data, uintptr_t len) {
+ (void)ctx;
+ (void)data;
+ (void)len;
+ return -1;
+}
+
+static int failing_flush(void *ctx) {
+ (void)ctx;
+ return -1;
+}
+
+static int failing_read_at(void *ctx, uint64_t offset, uint8_t *dst, uintptr_t
len) {
+ (void)ctx;
+ (void)offset;
+ (void)dst;
+ (void)len;
+ return -1;
+}
+
+static int64_t cluster_base_id(size_t cluster) {
+ return (int64_t)((cluster + 1) * 100000);
+}
+
+static void fill_roundtrip_data(float *data, int64_t *ids) {
+ for (size_t i = 0; i < ROUNDTRIP_VECTOR_COUNT; i++) {
+ 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;
+ ids[i] = cluster_base_id(cluster) + (int64_t)local;
+ }
+}
+
+static void assert_id_in_cluster(int64_t id, size_t cluster) {
+ int64_t base = cluster_base_id(cluster);
+ ASSERT_TRUE(id >= base);
+ ASSERT_TRUE(id < base + ROUNDTRIP_PER_LIST);
+}
+
+static void run_roundtrip(
+ const char *name,
+ const char *const *keys,
+ const char *const *values,
+ uintptr_t num_options,
+ uint32_t expected_index_type,
+ uintptr_t expected_pq_m,
+ uintptr_t expected_hnsw_m) {
+ PaimonVindexWriterHandle *writer =
+ paimon_vindex_writer_open(keys, values, num_options);
if (writer == NULL) {
fail_ffi("writer open failed");
}
@@ -120,17 +188,16 @@ static void test_basic_roundtrip(void) {
}
ASSERT_EQ_I64(dimension, 2);
- const float data[] = {
- 0.0f, 0.0f,
- 1.0f, 0.0f,
- 10.0f, 10.0f,
- 11.0f, 10.0f,
- };
- const int64_t ids[] = {100, 101, 200, 201};
- if (paimon_vindex_writer_train(writer, data, 4) != 0) {
+ float *data = (float *)malloc(sizeof(float) * ROUNDTRIP_VECTOR_COUNT *
ROUNDTRIP_DIMENSION);
+ int64_t *ids = (int64_t *)malloc(sizeof(int64_t) * ROUNDTRIP_VECTOR_COUNT);
+ ASSERT_TRUE(data != NULL);
+ ASSERT_TRUE(ids != NULL);
+ fill_roundtrip_data(data, ids);
+
+ if (paimon_vindex_writer_train(writer, data, ROUNDTRIP_VECTOR_COUNT) != 0)
{
fail_ffi("writer train failed");
}
- if (paimon_vindex_writer_add_vectors(writer, ids, data, 4) != 0) {
+ if (paimon_vindex_writer_add_vectors(writer, ids, data,
ROUNDTRIP_VECTOR_COUNT) != 0) {
fail_ffi("writer add failed");
}
@@ -160,11 +227,13 @@ static void test_basic_roundtrip(void) {
if (paimon_vindex_reader_metadata(reader, &metadata) != 0) {
fail_ffi("reader metadata failed");
}
- ASSERT_EQ_I64(metadata.index_type, PAIMON_VINDEX_INDEX_TYPE_IVF_FLAT);
+ 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.nlist, 2);
- ASSERT_EQ_I64(metadata.total_vectors, 4);
+ ASSERT_EQ_I64(metadata.nlist, 4);
+ ASSERT_EQ_I64(metadata.total_vectors, ROUNDTRIP_VECTOR_COUNT);
+ ASSERT_EQ_I64(metadata.pq_m, expected_pq_m);
+ ASSERT_EQ_I64(metadata.hnsw_m, expected_hnsw_m);
if (paimon_vindex_reader_optimize_for_search(reader) != 0) {
fail_ffi("reader optimize_for_search failed");
@@ -174,28 +243,142 @@ static void test_basic_roundtrip(void) {
int64_t result_ids[2] = {0};
float result_distances[2] = {0};
if (paimon_vindex_reader_search(
- reader, query, 2, 2, 0, result_ids, result_distances, 2) != 0) {
+ reader, query, 2, 4, 16, result_ids, result_distances, 2) != 0) {
fail_ffi("reader search failed");
}
- ASSERT_EQ_I64(result_ids[0], 100);
+ assert_id_in_cluster(result_ids[0], 0);
ASSERT_TRUE(isfinite(result_distances[0]));
- const float queries[] = {0.0f, 0.0f, 10.0f, 10.0f};
+ const float queries[] = {0.0f, 0.0f, 20.0f, 20.0f};
int64_t batch_ids[2] = {0};
float batch_distances[2] = {0};
if (paimon_vindex_reader_search_batch(
- reader, queries, 2, 1, 2, 0, batch_ids, batch_distances, 2) != 0) {
+ reader, queries, 2, 1, 4, 16, batch_ids, batch_distances, 2) != 0)
{
fail_ffi("reader search batch failed");
}
- ASSERT_EQ_I64(batch_ids[0], 100);
- ASSERT_EQ_I64(batch_ids[1], 200);
+ assert_id_in_cluster(batch_ids[0], 0);
+ assert_id_in_cluster(batch_ids[1], 1);
paimon_vindex_reader_free(reader);
free(buf.data);
- printf("PASS test_basic_roundtrip\n");
+ free(data);
+ free(ids);
+ printf("PASS %s\n", name);
+}
+
+static PaimonVindexWriterHandle *new_trained_flat_writer(void) {
+ const char *keys[] = {"index.type", "dimension", "nlist", "metric"};
+ const char *values[] = {"ivf_flat", "1", "1", "l2"};
+ PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(keys, values,
4);
+ if (writer == NULL) {
+ fail_ffi("writer open failed");
+ }
+
+ const float data[] = {0.0f, 1.0f};
+ const int64_t ids[] = {1, 2};
+ if (paimon_vindex_writer_train(writer, data, 2) != 0) {
+ fail_ffi("writer train failed");
+ }
+ if (paimon_vindex_writer_add_vectors(writer, ids, data, 2) != 0) {
+ fail_ffi("writer add failed");
+ }
+ return writer;
+}
+
+static void test_output_write_callback_error_propagates(void) {
+ PaimonVindexWriterHandle *writer = new_trained_flat_writer();
+ struct PaimonVindexOutputFile output = {
+ .ctx = NULL,
+ .write_fn = failing_write,
+ .flush_fn = mem_flush,
+ .get_pos_fn = NULL,
+ };
+
+ ASSERT_TRUE(paimon_vindex_writer_write_index(writer, output) != 0);
+ assert_last_error_contains("write callback failed");
+ paimon_vindex_writer_free(writer);
+ printf("PASS output_write_callback_error_propagates\n");
+}
+
+static void test_output_flush_callback_error_propagates(void) {
+ PaimonVindexWriterHandle *writer = new_trained_flat_writer();
+ struct MemBuffer buf = {0};
+ struct PaimonVindexOutputFile output = {
+ .ctx = &buf,
+ .write_fn = mem_write,
+ .flush_fn = failing_flush,
+ .get_pos_fn = mem_pos,
+ };
+
+ ASSERT_TRUE(paimon_vindex_writer_write_index(writer, output) != 0);
+ assert_last_error_contains("flush callback failed");
+ paimon_vindex_writer_free(writer);
+ free(buf.data);
+ printf("PASS output_flush_callback_error_propagates\n");
+}
+
+static void test_input_read_callback_error_propagates(void) {
+ struct PaimonVindexInputFile input = {
+ .ctx = NULL,
+ .read_at_fn = failing_read_at,
+ };
+
+ PaimonVindexReaderHandle *reader = paimon_vindex_reader_open(input);
+ ASSERT_TRUE(reader == NULL);
+ assert_last_error_contains("read_at callback failed");
+ printf("PASS input_read_callback_error_propagates\n");
+}
+
+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"};
+ run_roundtrip(
+ "ivf_flat_roundtrip",
+ flat_keys,
+ flat_values,
+ 4,
+ PAIMON_VINDEX_INDEX_TYPE_IVF_FLAT,
+ 0,
+ 0);
+
+ const char *pq_keys[] = {"index.type", "dimension", "nlist", "metric",
"pq.m"};
+ const char *pq_values[] = {"ivf_pq", "2", "4", "l2", "1"};
+ run_roundtrip(
+ "ivf_pq_roundtrip",
+ pq_keys,
+ pq_values,
+ 5,
+ PAIMON_VINDEX_INDEX_TYPE_IVF_PQ,
+ 1,
+ 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"};
+ run_roundtrip(
+ "ivf_hnsw_flat_roundtrip",
+ hnsw_flat_keys,
+ hnsw_flat_values,
+ 5,
+ PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_FLAT,
+ 0,
+ 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"};
+ run_roundtrip(
+ "ivf_hnsw_sq_roundtrip",
+ hnsw_sq_keys,
+ hnsw_sq_values,
+ 5,
+ PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_SQ,
+ 0,
+ 4);
}
int main(void) {
- test_basic_roundtrip();
+ test_supported_index_roundtrips();
+ test_output_write_callback_error_propagates();
+ test_output_flush_callback_error_propagates();
+ test_input_read_callback_error_propagates();
return 0;
}
diff --git a/core/Cargo.toml b/core/Cargo.toml
index f331ef7..82f67c0 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -42,3 +42,7 @@ harness = false
[[bench]]
name = "ann_bench"
harness = false
+
+[[bench]]
+name = "ivfhnswsq_filter_bench"
+harness = false
diff --git a/core/benches/ivfhnswsq_filter_bench.rs
b/core/benches/ivfhnswsq_filter_bench.rs
new file mode 100644
index 0000000..4738280
--- /dev/null
+++ b/core/benches/ivfhnswsq_filter_bench.rs
@@ -0,0 +1,293 @@
+use paimon_vindex_core::distance::MetricType;
+use paimon_vindex_core::hnsw::HnswBuildParams;
+use paimon_vindex_core::index::{
+ VectorIndexConfig, VectorIndexReader, VectorIndexWriter,
VectorSearchParams,
+};
+use paimon_vindex_core::io::PosWriter;
+use roaring::RoaringTreemap;
+use std::env;
+use std::io::{self, Cursor};
+use std::time::{Duration, Instant};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let cfg = Config::from_env()?;
+ cfg.validate()?;
+
+ let (data, queries) = generate_clustered_data(cfg.n, cfg.nq, cfg.d,
cfg.clusters, cfg.seed);
+ let ids: Vec<i64> = (0..cfg.n as i64).collect();
+
+ let start = Instant::now();
+ let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfHnswSq {
+ dimension: cfg.d,
+ nlist: cfg.nlist,
+ metric: MetricType::L2,
+ hnsw: cfg.hnsw_params(),
+ })?;
+ writer.train(&data, cfg.n)?;
+ writer.add_vectors(&ids, &data, cfg.n)?;
+ let mut index_bytes = Vec::new();
+ writer.write(&mut PosWriter::new(&mut index_bytes))?;
+ let build = start.elapsed();
+
+ println!(
+
"n={},nq={},d={},k={},nlist={},nprobe={},ef_search={},index_bytes={},build_ms={}",
+ cfg.n,
+ cfg.nq,
+ cfg.d,
+ cfg.k,
+ cfg.nlist,
+ cfg.nprobe,
+ cfg.ef_search,
+ index_bytes.len(),
+ build.as_millis()
+ );
+
println!("filter_stride,allowed,optimized,warmup_ms,search_ms,us_per_query,qps");
+
+ for stride in &cfg.filter_strides {
+ let filter_bytes = filter_bytes(cfg.n, *stride)?;
+ let allowed = cfg.n.div_ceil(*stride);
+ let baseline = run_case(&cfg, &index_bytes, &queries, &filter_bytes,
false)?;
+ let optimized = run_case(&cfg, &index_bytes, &queries, &filter_bytes,
true)?;
+ assert_same_results(*stride, &baseline.result, &optimized.result);
+ print_row(
+ *stride,
+ allowed,
+ false,
+ baseline.warmup,
+ baseline.search,
+ cfg.nq,
+ );
+ print_row(
+ *stride,
+ allowed,
+ true,
+ optimized.warmup,
+ optimized.search,
+ cfg.nq,
+ );
+ }
+
+ Ok(())
+}
+
+struct Config {
+ n: usize,
+ nq: usize,
+ d: usize,
+ k: usize,
+ nlist: usize,
+ nprobe: usize,
+ ef_search: usize,
+ hnsw_m: usize,
+ hnsw_ef_construction: usize,
+ hnsw_max_level: usize,
+ clusters: usize,
+ seed: u64,
+ filter_strides: Vec<usize>,
+}
+
+impl Config {
+ fn from_env() -> Result<Self, Box<dyn std::error::Error>> {
+ Ok(Self {
+ n: read_env("FILTER_BENCH_N", 50_000)?,
+ nq: read_env("FILTER_BENCH_NQ", 500)?,
+ d: read_env("FILTER_BENCH_D", 128)?,
+ k: read_env("FILTER_BENCH_K", 10)?,
+ nlist: read_env("FILTER_BENCH_NLIST", 64)?,
+ nprobe: read_env("FILTER_BENCH_NPROBE", 32)?,
+ ef_search: read_env("FILTER_BENCH_EF_SEARCH", 80)?,
+ hnsw_m: read_env("FILTER_BENCH_HNSW_M", 20)?,
+ hnsw_ef_construction:
read_env("FILTER_BENCH_HNSW_EF_CONSTRUCTION", 150)?,
+ hnsw_max_level: read_env("FILTER_BENCH_HNSW_MAX_LEVEL", 7)?,
+ clusters: read_env("FILTER_BENCH_CLUSTERS", 32)?,
+ seed: read_env("FILTER_BENCH_SEED", 42)?,
+ filter_strides: read_strides("FILTER_BENCH_FILTER_STRIDES", &[1,
4, 16, 64])?,
+ })
+ }
+
+ fn validate(&self) -> Result<(), Box<dyn std::error::Error>> {
+ if self.n == 0 || self.nq == 0 || self.d == 0 || self.k == 0 {
+ return Err("FILTER_BENCH_N, NQ, D, and K must be greater than
0".into());
+ }
+ if self.nlist == 0 || self.nprobe == 0 || self.nprobe > self.nlist ||
self.nlist > self.n {
+ return Err("FILTER_BENCH_NLIST/NPROBE must satisfy 0 < nprobe <=
nlist <= n".into());
+ }
+ if self.clusters == 0 {
+ return Err("FILTER_BENCH_CLUSTERS must be greater than 0".into());
+ }
+ if self.filter_strides.is_empty() || self.filter_strides.contains(&0) {
+ return Err("FILTER_BENCH_FILTER_STRIDES must contain positive
integers".into());
+ }
+ Ok(())
+ }
+
+ fn hnsw_params(&self) -> HnswBuildParams {
+ HnswBuildParams {
+ m: self.hnsw_m,
+ ef_construction: self.hnsw_ef_construction,
+ max_level: self.hnsw_max_level,
+ }
+ .sanitized()
+ }
+}
+
+struct CaseResult {
+ warmup: Duration,
+ search: Duration,
+ result: (Vec<i64>, Vec<f32>),
+}
+
+fn run_case(
+ cfg: &Config,
+ index_bytes: &[u8],
+ queries: &[f32],
+ filter_bytes: &[u8],
+ optimized: bool,
+) -> io::Result<CaseResult> {
+ let mut reader =
VectorIndexReader::open(Cursor::new(index_bytes.to_vec()))?;
+ let warmup_start = Instant::now();
+ if optimized {
+ reader.optimize_for_search()?;
+ }
+ let warmup = warmup_start.elapsed();
+
+ let params = VectorSearchParams::with_ef_search(cfg.k, cfg.nprobe,
cfg.ef_search);
+ let _ = reader.search_batch_with_roaring_filter(queries, cfg.nq, params,
filter_bytes)?;
+
+ let start = Instant::now();
+ let result = reader.search_batch_with_roaring_filter(queries, cfg.nq,
params, filter_bytes)?;
+ let search = start.elapsed();
+ Ok(CaseResult {
+ warmup,
+ search,
+ result,
+ })
+}
+
+fn filter_bytes(n: usize, stride: usize) -> io::Result<Vec<u8>> {
+ let mut filter = RoaringTreemap::new();
+ for id in (0..n as u64).step_by(stride) {
+ filter.insert(id);
+ }
+ let mut bytes = Vec::new();
+ filter.serialize_into(&mut bytes)?;
+ Ok(bytes)
+}
+
+fn assert_same_results(
+ stride: usize,
+ expected: &(Vec<i64>, Vec<f32>),
+ actual: &(Vec<i64>, Vec<f32>),
+) {
+ assert_eq!(
+ actual.0, expected.0,
+ "ids should match for stride {}",
+ stride
+ );
+ assert_eq!(
+ actual.1.len(),
+ expected.1.len(),
+ "distance count should match for stride {}",
+ stride
+ );
+ for (actual, expected) in actual.1.iter().zip(expected.1.iter()) {
+ assert!(
+ (actual - expected).abs() < 1e-4,
+ "distance {} should match {} for stride {}",
+ actual,
+ expected,
+ stride
+ );
+ }
+}
+
+fn print_row(
+ filter_stride: usize,
+ allowed: usize,
+ optimized: bool,
+ warmup: Duration,
+ search: Duration,
+ nq: usize,
+) {
+ println!(
+ "{},{},{},{},{},{:.2},{:.2}",
+ filter_stride,
+ allowed,
+ optimized,
+ warmup.as_millis(),
+ search.as_millis(),
+ search.as_secs_f64() * 1_000_000.0 / nq as f64,
+ nq as f64 / search.as_secs_f64(),
+ );
+}
+
+fn generate_clustered_data(
+ n: usize,
+ nq: usize,
+ d: usize,
+ clusters: usize,
+ seed: u64,
+) -> (Vec<f32>, Vec<f32>) {
+ let mut rng = Lcg::new(seed);
+ let mut centers = vec![0.0f32; clusters * d];
+ for value in &mut centers {
+ *value = rng.next_f32() * 30.0;
+ }
+
+ let mut data = vec![0.0f32; n * d];
+ for i in 0..n {
+ let cluster = i % clusters;
+ for j in 0..d {
+ data[i * d + j] = centers[cluster * d + j] + rng.next_f32();
+ }
+ }
+
+ let mut queries = vec![0.0f32; nq * d];
+ for qi in 0..nq {
+ let source = (qi * 9973) % n;
+ queries[qi * d..(qi + 1) * d].copy_from_slice(&data[source *
d..(source + 1) * d]);
+ }
+ (data, queries)
+}
+
+fn read_env<T>(name: &str, default: T) -> Result<T, Box<dyn std::error::Error>>
+where
+ T: std::str::FromStr,
+ T::Err: std::error::Error + 'static,
+{
+ match env::var(name) {
+ Ok(value) => Ok(value.parse()?),
+ Err(env::VarError::NotPresent) => Ok(default),
+ Err(err) => Err(Box::new(err)),
+ }
+}
+
+fn read_strides(name: &str, default: &[usize]) -> Result<Vec<usize>, Box<dyn
std::error::Error>> {
+ match env::var(name) {
+ Ok(value) => value
+ .split(',')
+ .map(|part| {
+ part.trim()
+ .parse::<usize>()
+ .map_err(|e| -> Box<dyn std::error::Error> { Box::new(e) })
+ })
+ .collect(),
+ Err(env::VarError::NotPresent) => Ok(default.to_vec()),
+ Err(err) => Err(Box::new(err)),
+ }
+}
+
+struct Lcg {
+ state: u64,
+}
+
+impl Lcg {
+ fn new(seed: u64) -> Self {
+ Self { state: seed }
+ }
+
+ fn next_f32(&mut self) -> f32 {
+ self.state =
self.state.wrapping_mul(6364136223846793005).wrapping_add(1);
+ ((self.state >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0
+ }
+}
diff --git a/core/tests/storage_format_fixtures.rs
b/core/tests/storage_format_fixtures.rs
index bf5e453..fea6da3 100644
--- a/core/tests/storage_format_fixtures.rs
+++ b/core/tests/storage_format_fixtures.rs
@@ -84,6 +84,37 @@ fn
storage_format_v1_golden_fixtures_match_current_writers_and_readers() {
}
}
+#[test]
+fn storage_format_v1_golden_fixtures_support_search_warmup() {
+ for case in fixture_cases() {
+ let fixture = hex_to_bytes(case.fixture_hex);
+
+ let mut baseline =
VectorIndexReader::open(Cursor::new(fixture.clone())).unwrap();
+ let expected = baseline.search(&case.query, case.params).unwrap();
+
+ let mut optimized =
VectorIndexReader::open(Cursor::new(fixture)).unwrap();
+ optimized.optimize_for_search().unwrap();
+ let actual = optimized.search(&case.query, case.params).unwrap();
+
+ assert_eq!(actual.0, expected.0, "{} optimized ids", case.name);
+ assert_eq!(
+ actual.1.len(),
+ expected.1.len(),
+ "{} optimized distance count",
+ case.name
+ );
+ for (actual, expected) in actual.1.iter().zip(expected.1.iter()) {
+ assert!(
+ (actual - expected).abs() < 1e-4,
+ "{} optimized distance {} should match {}",
+ case.name,
+ actual,
+ expected
+ );
+ }
+ }
+}
+
#[test]
#[ignore]
fn print_storage_format_v1_fixture_hex() {
diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp
index c1eeb04..96bd2dd 100644
--- a/cpp/test_vindex.cpp
+++ b/cpp/test_vindex.cpp
@@ -45,6 +45,11 @@ struct MemBuffer {
size_t pos = 0;
};
+constexpr size_t kRoundtripDimension = 2;
+constexpr size_t kRoundtripNlist = 4;
+constexpr size_t kRoundtripPerList = 128;
+constexpr size_t kRoundtripVectorCount = kRoundtripNlist * kRoundtripPerList;
+
static paimon::vindex::OutputFile make_output(MemBuffer& buf) {
paimon::vindex::OutputFile out;
out.write_fn = [&buf](const uint8_t* data, size_t len) -> int {
@@ -67,25 +72,51 @@ static paimon::vindex::InputFile make_input(const
MemBuffer& buf) {
return in;
}
-static void test_basic_roundtrip() {
- std::vector<std::pair<std::string, std::string>> options = {
- {"index.type", "ivf_flat"},
- {"dimension", "2"},
- {"nlist", "2"},
- {"metric", "l2"},
- };
+static int64_t cluster_base_id(size_t cluster) {
+ return static_cast<int64_t>((cluster + 1) * 100000);
+}
+
+static std::vector<float> roundtrip_data() {
+ std::vector<float> data(kRoundtripVectorCount * kRoundtripDimension);
+ for (size_t i = 0; i < kRoundtripVectorCount; i++) {
+ size_t cluster = i / kRoundtripPerList;
+ size_t local = i % kRoundtripPerList;
+ float center = static_cast<float>(cluster) * 20.0f;
+ data[i * kRoundtripDimension] = center + static_cast<float>(local %
16) * 0.001f;
+ data[i * kRoundtripDimension + 1] = center + static_cast<float>(local
/ 16) * 0.001f;
+ }
+ return data;
+}
+
+static std::vector<int64_t> roundtrip_ids() {
+ std::vector<int64_t> ids(kRoundtripVectorCount);
+ for (size_t i = 0; i < kRoundtripVectorCount; i++) {
+ size_t cluster = i / kRoundtripPerList;
+ size_t local = i % kRoundtripPerList;
+ ids[i] = cluster_base_id(cluster) + static_cast<int64_t>(local);
+ }
+ return ids;
+}
+
+static void assert_id_in_cluster(int64_t id, size_t cluster) {
+ int64_t base = cluster_base_id(cluster);
+ ASSERT_TRUE(id >= base);
+ ASSERT_TRUE(id < base + static_cast<int64_t>(kRoundtripPerList));
+}
+
+static void run_roundtrip(
+ const char* name,
+ const std::vector<std::pair<std::string, std::string>>& options,
+ uint32_t expected_index_type,
+ size_t expected_pq_m,
+ size_t expected_hnsw_m) {
paimon::vindex::Writer writer(options);
ASSERT_EQ(writer.dimension(), 2);
- std::vector<float> data = {
- 0.0f, 0.0f,
- 1.0f, 0.0f,
- 10.0f, 10.0f,
- 11.0f, 10.0f,
- };
- std::vector<int64_t> ids = {100, 101, 200, 201};
- writer.train(data.data(), 4);
- writer.add_vectors(ids.data(), data.data(), 4);
+ std::vector<float> data = roundtrip_data();
+ std::vector<int64_t> ids = roundtrip_ids();
+ writer.train(data.data(), kRoundtripVectorCount);
+ writer.add_vectors(ids.data(), data.data(), kRoundtripVectorCount);
MemBuffer buf;
writer.write_index(make_output(buf));
@@ -93,29 +124,84 @@ static void test_basic_roundtrip() {
paimon::vindex::Reader reader(make_input(buf));
auto metadata = reader.metadata();
- ASSERT_EQ(metadata.index_type, PAIMON_VINDEX_INDEX_TYPE_IVF_FLAT);
+ ASSERT_EQ(metadata.index_type, expected_index_type);
ASSERT_EQ(metadata.dimension, 2);
- ASSERT_EQ(metadata.nlist, 2);
+ ASSERT_EQ(metadata.nlist, 4);
ASSERT_EQ(metadata.metric, PAIMON_VINDEX_METRIC_L2);
- ASSERT_EQ(metadata.total_vectors, 4);
+ ASSERT_EQ(metadata.total_vectors, kRoundtripVectorCount);
+ ASSERT_EQ(metadata.pq_m, expected_pq_m);
+ ASSERT_EQ(metadata.hnsw_m, expected_hnsw_m);
reader.optimize_for_search();
const float query[] = {0.0f, 0.0f};
- auto result = reader.search(query, 2, 2);
+ auto result = reader.search(query, 2, 4, 16);
ASSERT_EQ(result.ids.size(), 2);
- ASSERT_EQ(result.ids[0], 100);
+ assert_id_in_cluster(result.ids[0], 0);
ASSERT_TRUE(std::isfinite(result.distances[0]));
- const float queries[] = {0.0f, 0.0f, 10.0f, 10.0f};
- auto batch = reader.search_batch(queries, 2, 1, 2);
+ const float queries[] = {0.0f, 0.0f, 20.0f, 20.0f};
+ auto batch = reader.search_batch(queries, 2, 1, 4, 16);
ASSERT_EQ(batch.ids.size(), 2);
- ASSERT_EQ(batch.ids[0], 100);
- ASSERT_EQ(batch.ids[1], 200);
- printf("PASS test_basic_roundtrip\n");
+ assert_id_in_cluster(batch.ids[0], 0);
+ assert_id_in_cluster(batch.ids[1], 1);
+ printf("PASS %s\n", name);
+}
+
+static void test_supported_index_roundtrips() {
+ run_roundtrip(
+ "ivf_flat_roundtrip",
+ {
+ {"index.type", "ivf_flat"},
+ {"dimension", "2"},
+ {"nlist", "4"},
+ {"metric", "l2"},
+ },
+ PAIMON_VINDEX_INDEX_TYPE_IVF_FLAT,
+ 0,
+ 0);
+
+ run_roundtrip(
+ "ivf_pq_roundtrip",
+ {
+ {"index.type", "ivf_pq"},
+ {"dimension", "2"},
+ {"nlist", "4"},
+ {"metric", "l2"},
+ {"pq.m", "1"},
+ },
+ PAIMON_VINDEX_INDEX_TYPE_IVF_PQ,
+ 1,
+ 0);
+
+ run_roundtrip(
+ "ivf_hnsw_flat_roundtrip",
+ {
+ {"index.type", "ivf_hnsw_flat"},
+ {"dimension", "2"},
+ {"nlist", "4"},
+ {"metric", "l2"},
+ {"hnsw.m", "4"},
+ },
+ PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_FLAT,
+ 0,
+ 4);
+
+ run_roundtrip(
+ "ivf_hnsw_sq_roundtrip",
+ {
+ {"index.type", "ivf_hnsw_sq"},
+ {"dimension", "2"},
+ {"nlist", "4"},
+ {"metric", "l2"},
+ {"hnsw.m", "4"},
+ },
+ PAIMON_VINDEX_INDEX_TYPE_IVF_HNSW_SQ,
+ 0,
+ 4);
}
int main() {
- test_basic_roundtrip();
+ test_supported_index_roundtrips();
return 0;
}
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 6e40171..320f2ec 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,6 +23,11 @@ import java.util.Map;
public class VectorIndexNativeValidationTest {
+ private static final int ROUNDTRIP_DIMENSION = 2;
+ 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;
+
public static void main(String[] args) {
if (args.length != 1) {
throw new IllegalArgumentException("native library path is
required");
@@ -34,6 +39,7 @@ public class VectorIndexNativeValidationTest {
testWriterRejectsNonFiniteValues();
testReaderValidationComesFromCore();
testReaderRejectsNonFiniteQueries();
+ testSupportedIndexRoundtrips();
}
private static void testWriterValidationComesFromCore() {
@@ -72,7 +78,8 @@ public class VectorIndexNativeValidationTest {
}
private static void testReaderValidationComesFromCore() {
- VectorIndexReader reader = new VectorIndexReader(new
ByteArraySeekableInputStream(buildIndexBytes()));
+ VectorIndexReader reader =
+ new VectorIndexReader(new
ByteArraySeekableInputStream(buildIndexBytes()));
try {
assertThrowsMessage(
RuntimeException.class,
@@ -143,7 +150,8 @@ public class VectorIndexNativeValidationTest {
}
private static void testReaderRejectsNonFiniteQueries() {
- VectorIndexReader reader = new VectorIndexReader(new
ByteArraySeekableInputStream(buildIndexBytes()));
+ VectorIndexReader reader =
+ new VectorIndexReader(new
ByteArraySeekableInputStream(buildIndexBytes()));
try {
assertInvalidInput(
new ThrowingRunnable() {
@@ -174,12 +182,66 @@ public class VectorIndexNativeValidationTest {
}
}
+ 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_hnsw_flat",
+ ivfHnswOptions("ivf_hnsw_flat", ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST),
+ 0,
+ 4);
+ runRoundtrip(
+ "ivf_hnsw_sq",
+ ivfHnswOptions("ivf_hnsw_sq", ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST),
+ 0,
+ 4);
+ }
+
+ private static void runRoundtrip(
+ String indexType, Map<String, String> options, int expectedPqM,
int expectedHnswM) {
+ byte[] indexBytes =
+ buildIndexBytes(
+ options, roundtripData(), roundtripIds(),
ROUNDTRIP_VECTOR_COUNT);
+ VectorIndexReader reader =
+ new VectorIndexReader(new
ByteArraySeekableInputStream(indexBytes));
+ try {
+ VectorIndexMetadata metadata = reader.metadata();
+ assertEquals(indexType, metadata.indexType());
+ assertEquals(ROUNDTRIP_DIMENSION, metadata.dimension());
+ assertEquals(ROUNDTRIP_NLIST, metadata.nlist());
+ assertEquals("l2", metadata.metric());
+ assertEquals((long) ROUNDTRIP_VECTOR_COUNT,
metadata.totalVectors());
+ assertEquals(expectedPqM, metadata.pqM());
+ assertEquals(expectedHnswM, metadata.hnswM());
+
+ reader.optimizeForSearch();
+
+ VectorSearchResult single = reader.search(new float[] {0.0f,
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);
+ 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");
+ } finally {
+ reader.close();
+ }
+ }
+
private static byte[] buildIndexBytes() {
- VectorIndexWriter writer = new VectorIndexWriter(ivfFlatOptions());
+ return buildIndexBytes(ivfFlatOptions(), new float[] {0.0f, 1.0f}, new
long[] {1L, 2L}, 2);
+ }
+
+ private static byte[] buildIndexBytes(
+ Map<String, String> options, float[] data, long[] ids, int
vectorCount) {
+ VectorIndexWriter writer = new VectorIndexWriter(options);
ByteArrayPositionOutputStream output = new
ByteArrayPositionOutputStream();
try {
- writer.train(new float[] {0.0f, 1.0f}, 2);
- writer.addVectors(new long[] {1L, 2L}, new float[] {0.0f, 1.0f},
2);
+ writer.train(data, vectorCount);
+ writer.addVectors(ids, data, vectorCount);
writer.writeIndex(output);
return output.toByteArray();
} finally {
@@ -188,14 +250,105 @@ public class VectorIndexNativeValidationTest {
}
private static Map<String, String> ivfFlatOptions() {
+ return ivfFlatOptions(1, 1);
+ }
+
+ private static Map<String, String> ivfFlatOptions(int dimension, int
nlist) {
Map<String, String> options = new HashMap<String, String>();
options.put("index.type", "ivf_flat");
- options.put("dimension", "1");
- options.put("nlist", "1");
+ options.put("dimension", Integer.toString(dimension));
+ options.put("nlist", Integer.toString(nlist));
options.put("metric", "l2");
return options;
}
+ private static Map<String, String> ivfPqOptions(int dimension, int nlist,
int m) {
+ Map<String, String> options = ivfFlatOptions(dimension, nlist);
+ options.put("index.type", "ivf_pq");
+ options.put("pq.m", Integer.toString(m));
+ options.put("use-opq", "false");
+ return options;
+ }
+
+ private static Map<String, String> ivfHnswOptions(String indexType, int
dimension, int nlist) {
+ Map<String, String> options = ivfFlatOptions(dimension, nlist);
+ options.put("index.type", indexType);
+ options.put("hnsw.m", "4");
+ options.put("hnsw.ef-construction", "16");
+ options.put("hnsw.max-level", "4");
+ return options;
+ }
+
+ private static float[] roundtripData() {
+ float[] data = new float[ROUNDTRIP_VECTOR_COUNT * ROUNDTRIP_DIMENSION];
+ for (int i = 0; i < ROUNDTRIP_VECTOR_COUNT; i++) {
+ 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;
+ }
+ return data;
+ }
+
+ private static long[] roundtripIds() {
+ long[] ids = new long[ROUNDTRIP_VECTOR_COUNT];
+ for (int i = 0; i < ROUNDTRIP_VECTOR_COUNT; i++) {
+ int cluster = i / ROUNDTRIP_PER_LIST;
+ int local = i % ROUNDTRIP_PER_LIST;
+ ids[i] = clusterBaseId(cluster) + local;
+ }
+ return ids;
+ }
+
+ private static long clusterBaseId(int cluster) {
+ return (cluster + 1L) * 100000L;
+ }
+
+ private static void assertEquals(int expected, int actual) {
+ if (expected != actual) {
+ throw new AssertionError("expected " + expected + " but got " +
actual);
+ }
+ }
+
+ private static void assertEquals(long expected, long actual) {
+ if (expected != actual) {
+ throw new AssertionError("expected " + expected + " but got " +
actual);
+ }
+ }
+
+ private static void assertEquals(Object expected, Object actual) {
+ if (!expected.equals(actual)) {
+ throw new AssertionError("expected " + expected + " but got " +
actual);
+ }
+ }
+
+ private static void assertArrayEquals(long[] expected, long[] actual) {
+ if (expected.length != actual.length) {
+ throw new AssertionError(
+ "expected length " + expected.length + " but got " +
actual.length);
+ }
+ for (int i = 0; i < expected.length; i++) {
+ if (expected[i] != actual[i]) {
+ throw new AssertionError(
+ "expected[" + i + "] " + expected[i] + " but got " +
actual[i]);
+ }
+ }
+ }
+
+ private static void assertIdInCluster(long id, int cluster) {
+ long base = clusterBaseId(cluster);
+ if (id < base || id >= base + ROUNDTRIP_PER_LIST) {
+ throw new AssertionError("id " + id + " should be in cluster " +
cluster);
+ }
+ }
+
+ private static void assertFinite(float value, String label) {
+ if (!Float.isFinite(value)) {
+ throw new AssertionError(label + " should be finite but was " +
value);
+ }
+ }
+
private static void assertThrowsMessage(
Class<? extends Throwable> expected, String expectedMessage,
ThrowingRunnable runnable) {
try {
