[Rust] Standalone vector index build and ANN search with Python API by xuchen-plus · Pull Request #797 · lakesoul-io/LakeSoul

xuchen-plus · 2026-07-15T08:09:15Z

Summary

This PR implements single-machine vector index building and approximate nearest neighbor (ANN) search for LakeSoul, based on IVF+RaBitQ. Includes full Python Table API for end-to-end workflow.

Key Features

1. New crate: `lakesoul-vector`

IVF+RaBitQ core (pure computation, no IO dependency)
Two-pass streaming index build: reservoir sampling → K-Means clustering → RaBitQ quantization
Delta segment design for incremental index updates (only write new vectors, never modify existing files)
Object store persistence (S3/local) via ManifestStore with CAS concurrency safety
AVX2 FastScan SIMD batch distance computation (32 vectors/batch)

2. Per-shard index building (`lakesoul-io`)

VectorShardIndexBuilder: reads parquet files, builds IVF+RaBitQ index per (partition, bucket) shard
Auto-derives index prefix from file paths — no manual index_prefix parameter needed
Fresh mode (full build) and Loaded mode (incremental update)
Index stored at {table_path}/_vector_index/{col}/{partition}/{bucket}/

3. Vector search in read path (`lakesoul-io`)

LakeSoulReader::start() auto-injects vector search results as PK filter
Each per-bucket reader searches only its own bucket's index, returns top-K candidates
Merge + exact-distance re-rank at Python layer via rerank_by_distance()

4. Python Table API

# Build index
table.build_vector_index(column="vec", dim=200, nlist=16)

# Search with re-rank
ds = table.scan().options(reader_options={
    "vector_search_column": "vec",
    "vector_search_query": "0.1,0.2,...",
    "vector_search_top_k": "3",
}).to_arrow_dataset()

# Re-rank by exact distance
from lakesoul.vector_index import rerank_by_distance
result = rerank_by_distance(result_table, query_vec, "vec", top_k=3)

5. E2E tests (glove-200d, multi-bucket verified)

Local mode: Writer → build_shard_vector_index → sync_reader → recall (hash_bucket_num=4, Recall@3=1.00)
Catalog mode: create_table → write_arrow → build_vector_index → scan → recall
Incremental write + index update verified

Dependency Graph

lakesoul-vector (pure rabitq, no IO dep)
    ↑
lakesoul-io (builder + search + reader)
    ↑
lakesoul-python (PyO3 + Python API)

Testing

cargo check
uv run --directory python python tests/vector/test_e2e_glove.py
uv run --directory python python tests/vector/test_e2e_glove.py --use-catalog  # requires PG

Signed-off-by: chenxu <chenxu@dmetasoul.com>

mag1c1an1 · 2026-07-15T08:13:32Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6a483b6b55

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

- Extract first 700 train vectors + 5 test vectors from glove-200d into python/tests/vector/data/ (~554KB total) - Use relative paths in both Python tests and Rust tests - Fix VectorShardIndexBuilder::new() calls (remove removed index_prefix param) - Fix type annotation for Vec::new() in Rust test - CI can now run tests without external data dependency Signed-off-by: chenxu <chenxu@dmetasoul.com>

- Fix all clippy warnings in lakesoul-io (unused imports, identity map, collapsible ifs) - Move vector_search.rs into lakesoul-io/src/vector/search.rs - Use LakeSoulReader in VectorShardIndexBuilder::read_all_batches for proper merge-on-read handling - Return empty result (not full scan) when vector search yields no hits - Fix partition_desc construction to follow table's partition_by order - Reduce catalog test nlist to 8 for small per-bucket data (~125 vecs) - Clean up index files on disk between catalog test runs - Use new vector as query in incremental test for reliable verification Signed-off-by: chenxu <chenxu@dmetasoul.com>

Signed-off-by: chenxu <chenxu@dmetasoul.com>

- Remove identity map, use to_vec() instead of iter().copied().collect() - Allow clippy::useless_conversion and clippy::too_many_arguments - Replace needless_range_loop with iterator-based zip/enumerate Signed-off-by: chenxu <chenxu@dmetasoul.com>

- Auto-detect parquet schema from first file for proper column projection when using LakeSoulReader in read_all_batches - Use LocalFileSystem::new() (root at /) in tests so absolute index_prefix paths work correctly (matching Python behavior) - Mark test_read_parquet_schema and test_build_and_list_files as #[ignore] (require prepare_data.py) Signed-off-by: chenxu <chenxu@dmetasoul.com>

- Remove empty lines after doc comments (cluster.rs, ivf/mod.rs) - Add #![allow(dead_code)] to vendored rabitq modules with unused API surface (fastscan, fastscan_kernel, kmeans) - Wrap madvise call in unsafe{} block for Rust 2024 compliance (memory.rs) - Fix Cargo.toml: remove non-existent spill_bench, add mor_bench Signed-off-by: chenxu <chenxu@dmetasoul.com>

Signed-off-by: chenxu <chenxu@dmetasoul.com>

mag1c1an1 · 2026-07-20T06:14:57Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8a2ee0fe12

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+            let u = u.trim_end_matches('/');
+            std::path::Path::new(u)
+                .parent()
+                .and_then(|p| p.to_str())
+                .map(|s| s.to_string())


Preserve URL authorities when deriving the scan prefix

When a Python scan receives a file such as s3://bucket/table/part.parquet, treating the URL as a filesystem path derives s3:/bucket/table, and .with_prefix then makes register_object_store process an S3 URL without a host. This affects ordinary Python S3 scans, not only vector searches, and can fail before the correctly formed file URL is registered; parse URLs by scheme or avoid setting this prefix for non-local paths.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+        for batch in &all_batches {
+            total += batch.ids.len();
+            builder.insert_batch(batch.clone())?;
+        }


Insert only newly committed vectors during incremental builds

On every subsequent build_vector_index call, _group_files_by_shard supplies all files in the latest snapshot, and this loaded path inserts every row again without checking existing IDs. Rebuilding an unchanged table therefore appends a full duplicate copy each time; updates and compactions also retain obsolete vectors, causing unbounded index growth and duplicate or stale ANN candidates that can displace valid top-K results.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+    // Build object_store_options for LakeSoulReader (converted from store_config)
+    let object_store_options = store_config.clone();


Translate S3 credentials for the LakeSoul reader

For S3 index builds where credentials or a custom endpoint are supplied through LakeSoulCatalog.object_store_options rather than environment variables, _default_object_store_config converts them to keys such as access_key_id and endpoint, but cloning that map into LakeSoulIOConfig does not convert them back. The reader's S3 implementation only reads fs.s3a.access.key, fs.s3a.secret.key, fs.s3a.endpoint, and related keys, so the index store can authenticate while the reader that loads the Parquet vectors cannot.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+        let pk_column = io_config
+            .primary_keys
+            .first()
+            .cloned()
+            .unwrap_or_else(|| "id".to_string());


Retain the actual primary key for compacted scans

For a CompactionCommit, NativeMetadataClient.get_scan_plan_partitions deliberately passes an empty primary-key list, so this fallback selects a column literally named id. Any compacted table whose primary key has another name then builds the ANN ID filter against a nonexistent or unrelated column, making vector search fail or return incorrect rows even though the index was built with the real metadata primary key.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+        Err(e) => {
+            warn!("Failed to load vector index at '{}': {:?}", prefix, e);
+            return Ok(None);
+        }


Propagate index load failures instead of returning no matches

If loading an existing index fails because of an object-store outage, denied permission, corrupt segment, or incompatible manifest, this branch converts the failure into Ok(None). The caller interprets that as an empty ANN result and injects a constant-false filter, silently reporting that the table has no matching rows; only a verified missing-index condition should be optional, while operational and corruption errors must propagate.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-20T06:20:26Z

+            Err(e) => {
+                tracing::warn!("Invalid vector search query: {}", e);
+                return Ok(filters);


Reject malformed vector queries instead of scanning all rows

When vector_search_query contains an invalid float, returning the original filters silently disables vector search. A caller expecting top-K ANN results can therefore receive the entire table and trigger a much more expensive scan; propagate the parse error so invalid search input cannot be mistaken for a successful query.

Useful? React with 👍 / 👎.

dmetasoul01 added 6 commits July 15, 2026 14:07

[Rust] Impl vector index and search

4d1bedc

Signed-off-by: chenxu <chenxu@dmetasoul.com>

adjust crate structure and add index filter in io reader

e697b7b

Signed-off-by: chenxu <chenxu@dmetasoul.com>

add python e2e test

524ada9

Signed-off-by: chenxu <chenxu@dmetasoul.com>

add skill

19b2e94

Signed-off-by: chenxu <chenxu@dmetasoul.com>

use python table api for read/write

a483b77

Signed-off-by: chenxu <chenxu@dmetasoul.com>

update readme in rust/lakesoul-vector

6a483b6

Signed-off-by: chenxu <chenxu@dmetasoul.com>

xuchen-plus changed the title ~~[Rust] Single-machine vector index build and ANN search with Python API~~ [Rust] Standalone vector index build and ANN search with Python API Jul 15, 2026

chatgpt-codex-connector Bot reviewed Jul 15, 2026

View reviewed changes

xuchen-plus mentioned this pull request Jul 15, 2026

[Umbrella] LakeSoul Vector Index Roadmap #798

Open

dmetasoul01 added 5 commits July 15, 2026 20:42

fix tests

7ccd07d

Signed-off-by: chenxu <chenxu@dmetasoul.com>

xuchen-plus force-pushed the impl_vector_search branch from e5eb115 to e55b697 Compare July 16, 2026 10:21

dmetasoul01 added 3 commits July 17, 2026 14:45

fix build

6ff435f

Signed-off-by: chenxu <chenxu@dmetasoul.com>

fix build

8a2ee0f

Signed-off-by: chenxu <chenxu@dmetasoul.com>

xuchen-plus merged commit 6ee67e6 into lakesoul-io:main Jul 20, 2026
22 of 23 checks passed

xuchen-plus deleted the impl_vector_search branch July 20, 2026 06:15

chatgpt-codex-connector Bot reviewed Jul 20, 2026

View reviewed changes

		// Build object_store_options for LakeSoulReader (converted from store_config)
		let object_store_options = store_config.clone();

Uh oh!

Conversation

xuchen-plus commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key Features

1. New crate: lakesoul-vector

2. Per-shard index building (lakesoul-io)

3. Vector search in read path (lakesoul-io)

4. Python Table API

5. E2E tests (glove-200d, multi-bucket verified)

Dependency Graph

Testing

Uh oh!

mag1c1an1 commented Jul 15, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mag1c1an1 commented Jul 20, 2026

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

xuchen-plus commented Jul 15, 2026 •

edited

Loading

1. New crate: `lakesoul-vector`

2. Per-shard index building (`lakesoul-io`)

3. Vector search in read path (`lakesoul-io`)