feat: distributed vector search via index segment selection (#24)

Exposes Lance's segment-model APIs through the C ABI so a distributed
query engine (Velox, Presto worker, etc.) can fan a single k-NN query
out across workers, each scanning a slice of the logical index's
physical segments. Tracks
[lance#6309](lance-format/lance#6309).

## Distributed query pattern

```
Coordinator                                  Worker(s)
─────────────────                            ───────────────
open dataset
list segments  ──────── slice ──────────►   open same dataset
                                             scanner.nearest(q, k)
                                             scanner.index_segments(my_slice)
                                             return partial top-k stream
heap-merge partial top-k  ◄─────────────────  (Velox top-k operator handles this)
```

## Summary

**`lance_dataset_index_segment_count(ds, name)`** — number of physical
segments in a logical vector index. Returns 0 + `LANCE_ERR_NOT_FOUND`
for an unknown name.

**`lance_dataset_index_segments(ds, name, out_uuids)`** — fills a
caller-allocated buffer (`count * 16` bytes) with each segment's 16-byte
UUID (RFC 4122).

**`lance_scanner_set_index_segments(scanner, segment_uuids, len)`** —
restricts the next `lance_scanner_nearest()` query to a subset of
segments. `len=0` (any pointer) clears the restriction.

**C++ wrappers**:
- `Dataset::index_segment_count(name)` → `uint64_t`
- `Dataset::index_segments(name)` → `std::vector<std::array<uint8_t,
16>>`
- `Scanner::index_segments(uuids)` (typed vector overload + raw
`uint8_t*` + len overload) — fluent

## Lance dep bump

To get `Scanner::with_index_segments()` (merged in lance #6376) we bump
from crates.io `lance = \"3.0.1\"` to a `git+rev` pin at lance commit
`d630106d` (release tag `v5.0.0-beta.5`). beta-5 keeps arrow on 57.0.0 —
no transitive arrow churn. The `DatasetIndexExt` trait moved from
`lance_index` to `lance::index`; one import path adjusted in
`src/index.rs`.

When lance publishes 5.0.0 stable, the git+rev can be replaced with the
version pin.

## Test plan

- [x] `cargo fmt` clean
- [x] `cargo clippy --all-targets -- -D warnings` clean
- [x] `cargo test` — **75 passed** (70 from main + 5 new)
- [x] `cargo test --test compile_and_run_test -- --ignored` — 2 passed
(C + C++ smoke)

New tests:
- `test_index_segment_count_and_list` — build IVF index, count = 1, list
returns a non-zero UUID.
- `test_index_segment_count_unknown_index` — unknown name → `NotFound`.
- `test_scanner_set_index_segments_with_listed_uuids` — end-to-end k=5
nearest restricted to listed segment UUID, returns 5 results.
- `test_scanner_set_index_segments_unknown_uuid` — bogus UUID is
accepted at setter time, surfaces as an error at scan materialize time
with a message containing "segment".
- `test_scanner_set_index_segments_null_safety` — NULL scanner / NULL
pointer with len>0 / NULL with len=0 (clears).

## Follow-ups (not in this PR)

- Per-segment metadata: today we only expose UUID. A future pass could
add fragment_bitmap / dataset_version / num_indexed_rows so coordinators
can balance work by segment size.
- Distributed build: `commit_existing_index_segments()` and
`merge_existing_index_segments()` exist upstream — they'd let workers
each train one segment and the coordinator commit them atomically.
- Once lance publishes 5.0.0 stable, replace the git+rev pin with a
version pin.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: jja725

Cargo.lock

@@ -18,32 +18,33 @@ rust-version = "1.91.0"
 crate-type = ["cdylib", "staticlib", "rlib"]
 
 [dependencies]
-lance = { version = "4.0.1", features = ["substrait"] }
-lance-core = "4.0.1"
-lance-file = "4.0.1"
-lance-index = "4.0.1"
-lance-io = "4.0.1"
-lance-linalg = "4.0.1"
-arrow = { version = "57.0.0", features = ["prettyprint", "ffi"] }
-arrow-array = "57.0.0"
-arrow-schema = "57.0.0"
+lance = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6", features = ["substrait"] }
+lance-core = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-file = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-index = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-io = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-linalg = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+arrow = { version = "58.0.0", features = ["prettyprint", "ffi"] }
+arrow-array = "58.0.0"
+arrow-schema = "58.0.0"
 half = "2"
 tokio = { version = "1", features = ["rt-multi-thread", "sync"] }
 futures = "0.3"
 log = "0.4"
 pin-project = "1.0"
 snafu = "0.9"
+uuid = { version = "1", features = ["v4"] }
 
 [dev-dependencies]
-lance = { version = "4.0.1", features = ["substrait"] }
-lance-datafusion = { version = "4.0.1", features = ["substrait"] }
-lance-datagen = "4.0.1"
-lance-file = "4.0.1"
-lance-table = "4.0.1"
-datafusion = { version = "52.1.0", default-features = false }
+lance = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6", features = ["substrait"] }
+lance-datafusion = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6", features = ["substrait"] }
+lance-datagen = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-file = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+lance-table = { git = "https://github.com/lance-format/lance.git", rev = "e0e977a6" }
+datafusion = { version = "53.1.0", default-features = false }
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
-arrow-array = "57.0.0"
-arrow-schema = "57.0.0"
+arrow-array = "58.0.0"
+arrow-schema = "58.0.0"
 tempfile = "3"
 
 [profile.release]

Cargo.toml

@@ -707,6 +707,41 @@ uint64_t lance_dataset_index_count(const LanceDataset* dataset);
  */
 const char* lance_dataset_index_list_json(const LanceDataset* dataset);
 
+/* ─── Distributed vector search: index segment enumeration ─── */
+
+/**
+ * Count the segments that make up a logical vector index.
+ *
+ * A logical index is a set of physical segments (one per distributed-build
+ * worker, or one per fragment range). Each segment has a stable UUID. Returns
+ * 0 if the index does not exist (also sets `LANCE_ERR_NOT_FOUND`) or on error.
+ */
+uint64_t lance_dataset_index_segment_count(
+    const LanceDataset* dataset,
+    const char* index_name
+);
+
+/**
+ * Fill `out_uuids` with the UUIDs of the segments that make up a logical index.
+ * Each UUID is written as 16 raw bytes (RFC 4122 layout).
+ *
+ * @param out_uuids  Caller-allocated buffer for the UUIDs (byte length >= capacity * 16).
+ * @param capacity   Number of UUIDs the buffer can hold.
+ * @param out_count  Optional (may be NULL). On success, receives the number of
+ *                   UUIDs actually written.
+ *
+ * Returns 0 on success, -1 on error.  If the index has more segments than
+ * `capacity`, returns LANCE_ERR_INVALID_ARGUMENT without writing anything;
+ * the caller can retry with a larger buffer.
+ */
+int32_t lance_dataset_index_segments(
+    const LanceDataset* dataset,
+    const char* index_name,
+    uint8_t* out_uuids,
+    size_t capacity,
+    uint64_t* out_count
+);
+
 /* ─── Vector search (Phase 2) ─── */
 
 /**
@@ -736,6 +771,26 @@ int32_t lance_scanner_set_metric(LanceScanner* scanner, LanceMetricType metric);
 int32_t lance_scanner_set_use_index(LanceScanner* scanner, bool enable);
 int32_t lance_scanner_set_prefilter(LanceScanner* scanner, bool enable);
 
+/**
+ * Restrict the next k-NN query to a specific subset of vector index segments.
+ *
+ * Used by distributed query engines (e.g. Velox) to fan a single k-NN query
+ * out across workers, each handling a slice of segments. The coordinator gets
+ * the segment list via `lance_dataset_index_segments()`.
+ *
+ * @param segment_uuids Pointer to `len` 16-byte UUIDs concatenated end-to-end
+ *                      (total byte length = `len * 16`). Each UUID identifies
+ *                      one physical segment of a logical index.
+ * @param len           Number of UUIDs. Pass 0 (and segment_uuids may be NULL)
+ *                      to clear any previously-set segment restriction.
+ * @return 0 on success, -1 on error.
+ */
+int32_t lance_scanner_set_index_segments(
+    LanceScanner* scanner,
+    const uint8_t* segment_uuids,
+    size_t len
+);
+
 /* ─── Full-text search (Phase 2) ─── */
 
 /**

include/lance/lance.h

@@ -17,6 +17,7 @@
 
 #include "lance/lance.h"
 
+#include <array>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -533,6 +534,31 @@ class Dataset {
         return out;
     }
 
+    /// Number of segments that make up a logical vector index.
+    /// Throws lance::Error with code NotFound if the index does not exist.
+    uint64_t index_segment_count(const std::string& index_name) const {
+        uint64_t n = lance_dataset_index_segment_count(handle_.get(), index_name.c_str());
+        if (n == 0 && lance_last_error_code() != LANCE_OK) check_error();
+        return n;
+    }
+
+    /// UUIDs of the physical segments that make up a logical vector index.
+    /// Each UUID is a 16-byte array (RFC 4122 layout). Used by distributed
+    /// query engines to fan k-NN out across workers — see
+    /// `Scanner::index_segments`.
+    std::vector<std::array<uint8_t, 16>> index_segments(const std::string& index_name) const {
+        uint64_t count = index_segment_count(index_name);
+        std::vector<std::array<uint8_t, 16>> out(count);
+        if (count == 0) return out;
+        uint64_t written = 0;
+        if (lance_dataset_index_segments(handle_.get(), index_name.c_str(),
+                                          reinterpret_cast<uint8_t*>(out.data()),
+                                          static_cast<size_t>(count), &written) != 0)
+            check_error();
+        out.resize(static_cast<size_t>(written));
+        return out;
+    }
+
     /// Access the underlying C handle (does not transfer ownership).
     const LanceDataset* c_handle() const { return handle_.get(); }
 
@@ -601,6 +627,21 @@ class Scanner {
         return substrait_filter(bytes.data(), bytes.size());
     }
 
+    /// Restrict the next k-NN query to a subset of vector index segments.
+    /// Pass `len` 16-byte UUIDs concatenated as a single byte buffer
+    /// (total bytes = `len * 16`). Pass len=0 (and any pointer) to clear.
+    Scanner& index_segments(const uint8_t* uuids, size_t len) {
+        if (lance_scanner_set_index_segments(handle_.get(), uuids, len) != 0)
+            check_error();
+        return *this;
+    }
+
+    /// Restrict the next k-NN query to a subset of vector index segments
+    /// (typed vector overload).
+    Scanner& index_segments(const std::vector<std::array<uint8_t, 16>>& uuids) {
+        return index_segments(reinterpret_cast<const uint8_t*>(uuids.data()), uuids.size());
+    }
+
     /// Materialize the scan as an ArrowArrayStream (blocking).
     void to_arrow_stream(ArrowArrayStream* out) {
         if (lance_scanner_to_arrow_stream(handle_.get(), out) != 0)

include/lance/lance.hpp

@@ -8,9 +8,10 @@
 
 use std::ffi::{CString, c_char};
 
+use lance::index::DatasetIndexExt;
 use lance_core::Result;
+use lance_index::IndexType;
 use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams};
-use lance_index::{DatasetIndexExt, IndexType};
 
 use crate::dataset::LanceDataset;
 use crate::error::{LanceErrorCode, ffi_try, set_last_error};
@@ -153,6 +154,146 @@ pub unsafe extern "C" fn lance_dataset_index_count(dataset: *const LanceDataset)
     }
 }
 
+/// Count the segments that make up a logical index.
+///
+/// A logical index is a set of physical segments (one per distributed-build worker
+/// or one per fragment range). Each segment has a stable UUID. Returns 0 if the
+/// index does not exist (also sets `LANCE_ERR_NOT_FOUND`).
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn lance_dataset_index_segment_count(
+    dataset: *const LanceDataset,
+    index_name: *const c_char,
+) -> u64 {
+    if dataset.is_null() || index_name.is_null() {
+        set_last_error(
+            LanceErrorCode::InvalidArgument,
+            "dataset and index_name must not be NULL",
+        );
+        return 0;
+    }
+    let ds = unsafe { &*dataset };
+    let name = match unsafe { helpers::parse_c_string(index_name) } {
+        Ok(Some(s)) => s,
+        Ok(None) => {
+            set_last_error(
+                LanceErrorCode::InvalidArgument,
+                "index_name must not be empty",
+            );
+            return 0;
+        }
+        Err(err) => {
+            crate::error::set_lance_error(&err);
+            return 0;
+        }
+    };
+    let snap = ds.snapshot();
+    match block_on(snap.load_indices()) {
+        Ok(indices) => {
+            let count = indices
+                .iter()
+                .filter(|i| !lance_index::is_system_index(i) && i.name == name)
+                .count();
+            if count == 0 {
+                set_last_error(
+                    LanceErrorCode::NotFound,
+                    format!("index '{}' not found", name),
+                );
+                return 0;
+            }
+            crate::error::clear_last_error();
+            count as u64
+        }
+        Err(err) => {
+            crate::error::set_lance_error(&err);
+            0
+        }
+    }
+}
+
+/// Fill `out_uuids` with the UUIDs of the segments that make up a logical index.
+///
+/// Each UUID is written as 16 raw bytes (RFC 4122 layout).
+///
+/// - `capacity`: number of UUIDs the caller allocated space for in `out_uuids`
+///   (byte length must be at least `capacity * 16`).
+/// - `out_count`: if non-NULL, receives the number of UUIDs actually written.
+///
+/// Returns 0 on success, -1 on error. If the index has more segments than
+/// `capacity`, returns `LANCE_ERR_INVALID_ARGUMENT` without writing anything.
+/// Callers can retry with a larger buffer.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn lance_dataset_index_segments(
+    dataset: *const LanceDataset,
+    index_name: *const c_char,
+    out_uuids: *mut u8,
+    capacity: usize,
+    out_count: *mut u64,
+) -> i32 {
+    ffi_try!(
+        unsafe {
+            dataset_index_segments_inner(dataset, index_name, out_uuids, capacity, out_count)
+        },
+        neg
+    )
+}
+
+unsafe fn dataset_index_segments_inner(
+    dataset: *const LanceDataset,
+    index_name: *const c_char,
+    out_uuids: *mut u8,
+    capacity: usize,
+    out_count: *mut u64,
+) -> Result<i32> {
+    if dataset.is_null() || index_name.is_null() || out_uuids.is_null() {
+        return Err(lance_core::Error::InvalidInput {
+            source: "dataset, index_name, and out_uuids must not be NULL".into(),
+            location: snafu::location!(),
+        });
+    }
+    let ds = unsafe { &*dataset };
+    let name = unsafe { helpers::parse_c_string(index_name)? }.ok_or_else(|| {
+        lance_core::Error::InvalidInput {
+            source: "index_name must not be empty".into(),
+            location: snafu::location!(),
+        }
+    })?;
+    let snap = ds.snapshot();
+    let indices = block_on(snap.load_indices())?;
+    let segments: Vec<_> = indices
+        .iter()
+        .filter(|i| !lance_index::is_system_index(i) && i.name == name)
+        .collect();
+    if segments.is_empty() {
+        return Err(lance_core::Error::IndexNotFound {
+            identity: format!("name='{}'", name),
+            location: snafu::location!(),
+        });
+    }
+    if segments.len() > capacity {
+        return Err(lance_core::Error::InvalidInput {
+            source: format!(
+                "out_uuids capacity ({}) too small for {} segments",
+                capacity,
+                segments.len()
+            )
+            .into(),
+            location: snafu::location!(),
+        });
+    }
+    // SAFETY: caller guarantees out_uuids has at least `capacity * 16` bytes,
+    // and we verified `segments.len() <= capacity` above.
+    for (i, seg) in segments.iter().enumerate() {
+        let bytes = seg.uuid.as_bytes();
+        unsafe {
+            std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_uuids.add(i * 16), 16);
+        }
+    }
+    if !out_count.is_null() {
+        unsafe { *out_count = segments.len() as u64 };
+    }
+    Ok(0)
+}
+
 /// Drop an index by name.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn lance_dataset_drop_index(