feat: add public vector segment builder API

Xuanwo · Xuanwo · commit 691cecb9af7c · 2026-03-20T23:35:44.000+08:00
diff --git a/rust/lance-index/src/lib.rs b/rust/lance-index/src/lib.rs
@@ -33,7 +33,7 @@ pub mod types;
 pub mod vector;
 
 pub use crate::traits::*;
-pub use crate::types::IndexSegment;
+pub use crate::types::{IndexSegment, IndexSegmentPlan};
 
 pub const INDEX_FILE_NAME: &str = "index.idx";
 /// The name of the auxiliary index file.
diff --git a/rust/lance-index/src/traits.rs b/rust/lance-index/src/traits.rs
@@ -126,6 +126,9 @@ pub trait IndexDescription: Send + Sync {
 #[async_trait]
 pub trait DatasetIndexExt {
     type IndexBuilder<'a>
+    where
+        Self: 'a;
+    type IndexSegmentBuilder<'a>
     where
         Self: 'a;
 
@@ -145,6 +148,21 @@ pub trait DatasetIndexExt {
         params: &'a dyn IndexParams,
     ) -> Self::IndexBuilder<'a>;
 
+    /// Create a builder for building index segments from partial index outputs.
+    ///
+    /// The staging UUID identifies a directory containing previously-built shard
+    /// outputs. The caller supplies the partial index metadata returned by
+    /// `execute_uncommitted()` so the builder can plan segment grouping without
+    /// rediscovering shard coverage.
+    ///
+    /// This is the canonical entry point for distributed vector segment build.
+    /// After building the physical segments, publish them as a
+    /// logical index with [`Self::commit_existing_index_segments`].
+    fn create_index_segment_builder<'a>(
+        &'a self,
+        staging_index_uuid: String,
+    ) -> Self::IndexSegmentBuilder<'a>;
+
     /// Create indices on columns.
     ///
     /// Upon finish, a new dataset version is generated.
@@ -275,6 +293,11 @@ pub trait DatasetIndexExt {
     async fn index_statistics(&self, index_name: &str) -> Result<String>;
 
     /// Commit one or more existing physical index segments as a logical index.
+    ///
+    /// This publishes already-built physical segments. It does not build
+    /// or merge index data; callers should first build segments with
+    /// [`Self::create_index_segment_builder`] or another index-specific build
+    /// path and then pass the resulting segments here.
     async fn commit_existing_index_segments(
         &mut self,
         index_name: &str,
diff --git a/rust/lance-index/src/types.rs b/rust/lance-index/src/types.rs
@@ -3,6 +3,8 @@
 
 use std::sync::Arc;
 
+use crate::IndexType;
+use lance_table::format::IndexMetadata;
 use roaring::RoaringBitmap;
 use uuid::Uuid;
 
@@ -73,3 +75,58 @@ impl IndexSegment {
         )
     }
 }
+
+/// A plan for building one physical segment from one or more vector index
+/// partial indices.
+#[derive(Debug, Clone, PartialEq)]
+pub struct IndexSegmentPlan {
+    staging_index_uuid: Uuid,
+    segment: IndexSegment,
+    partial_indices: Vec<IndexMetadata>,
+    estimated_bytes: u64,
+    requested_index_type: Option<IndexType>,
+}
+
+impl IndexSegmentPlan {
+    /// Create a plan for one built segment.
+    pub fn new(
+        staging_index_uuid: Uuid,
+        segment: IndexSegment,
+        partial_indices: Vec<IndexMetadata>,
+        estimated_bytes: u64,
+        requested_index_type: Option<IndexType>,
+    ) -> Self {
+        Self {
+            staging_index_uuid,
+            segment,
+            partial_indices,
+            estimated_bytes,
+            requested_index_type,
+        }
+    }
+
+    /// Return the staging index UUID that owns the partial shard outputs.
+    pub fn staging_index_uuid(&self) -> Uuid {
+        self.staging_index_uuid
+    }
+
+    /// Return the segment metadata that should be committed after this plan is built.
+    pub fn segment(&self) -> &IndexSegment {
+        &self.segment
+    }
+
+    /// Return the uncommitted partial index metadata that should be combined into the segment.
+    pub fn partial_indices(&self) -> &[IndexMetadata] {
+        &self.partial_indices
+    }
+
+    /// Return the estimated number of bytes covered by this plan.
+    pub fn estimated_bytes(&self) -> u64 {
+        self.estimated_bytes
+    }
+
+    /// Return the requested logical index type, if one was supplied to the planner.
+    pub fn requested_index_type(&self) -> Option<IndexType> {
+        self.requested_index_type
+    }
+}
diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs
@@ -582,6 +582,7 @@ impl IndexDescription for IndexDescriptionImpl {
 #[async_trait]
 impl DatasetIndexExt for Dataset {
     type IndexBuilder<'a> = CreateIndexBuilder<'a>;
+    type IndexSegmentBuilder<'a> = create::IndexSegmentBuilder<'a>;
 
     /// Create a builder for creating an index on columns.
     ///
@@ -627,6 +628,13 @@ impl DatasetIndexExt for Dataset {
         CreateIndexBuilder::new(self, columns, index_type, params)
     }
 
+    fn create_index_segment_builder<'a>(
+        &'a self,
+        staging_index_uuid: String,
+    ) -> create::IndexSegmentBuilder<'a> {
+        create::IndexSegmentBuilder::new(self, staging_index_uuid)
+    }
+
     #[instrument(skip_all)]
     async fn create_index(
         &mut self,
diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs
@@ -17,10 +17,10 @@ use crate::{
         vector_index_details,
     },
 };
-use futures::future::BoxFuture;
+use futures::future::{BoxFuture, try_join_all};
 use lance_core::datatypes::format_field_path;
 use lance_index::progress::{IndexBuildProgress, NoopIndexBuildProgress};
-use lance_index::{IndexParams, IndexType, scalar::CreatedIndex};
+use lance_index::{IndexParams, IndexSegment, IndexSegmentPlan, IndexType, scalar::CreatedIndex};
 use lance_index::{
     metrics::NoOpMetricsCollector,
     scalar::{LANCE_SCALAR_INDEX, ScalarIndexParams, inverted::tokenizer::InvertedIndexParams},
@@ -196,6 +196,7 @@ impl<'a> CreateIndexBuilder<'a> {
                 .map_err(|e| Error::index(format!("Invalid UUID string provided: {}", e)))?,
             None => Uuid::new_v4(),
         };
+        let mut output_index_uuid = index_id;
         let created_index = match (self.index_type, self.params.index_name()) {
             (
                 IndexType::Bitmap
@@ -323,7 +324,7 @@ impl<'a> CreateIndexBuilder<'a> {
                     if let Some(fragments) = &self.fragments {
                         // For distributed indexing, build only on specified fragments
                         // This creates temporary index metadata without committing
-                        Box::pin(build_distributed_vector_index(
+                        let shard_uuid = Box::pin(build_distributed_vector_index(
                             self.dataset,
                             column,
                             &index_name,
@@ -334,6 +335,7 @@ impl<'a> CreateIndexBuilder<'a> {
                             self.progress.clone(),
                         ))
                         .await?;
+                        output_index_uuid = shard_uuid;
                     } else {
                         // Standard full dataset indexing
                         Box::pin(build_vector_index(
@@ -359,7 +361,14 @@ impl<'a> CreateIndexBuilder<'a> {
                     .await?;
                 }
                 // Capture file sizes after vector index creation
-                let index_dir = self.dataset.indices_dir().child(index_id.to_string());
+                let index_dir = if self.fragments.is_some() {
+                    self.dataset
+                        .indices_dir()
+                        .child(index_id.to_string())
+                        .child(format!("partial_{}", output_index_uuid))
+                } else {
+                    self.dataset.indices_dir().child(index_id.to_string())
+                };
                 let files =
                     list_index_files_with_sizes(&self.dataset.object_store, &index_dir).await?;
                 CreatedIndex {
@@ -420,7 +429,7 @@ impl<'a> CreateIndexBuilder<'a> {
         };
 
         Ok(IndexMetadata {
-            uuid: index_id,
+            uuid: output_index_uuid,
             name: index_name,
             fields: vec![field.id],
             dataset_version: self.dataset.manifest.version,
@@ -494,6 +503,91 @@ impl<'a> IntoFuture for CreateIndexBuilder<'a> {
     }
 }
 
+/// Build physical index segments from previously-written partial index outputs.
+///
+/// Use [`DatasetIndexExt::create_index_segment_builder`] to open a staging root
+/// and then either:
+///
+/// - call [`Self::plan`] and orchestrate individual segment builds externally, or
+/// - call [`Self::build_all`] to build all segments on the current node.
+///
+/// This builder only builds physical segments. Publishing those segments as
+/// a logical index still requires [`DatasetIndexExt::commit_existing_index_segments`].
+/// Together these two APIs form the canonical distributed vector segment build workflow.
+#[derive(Clone)]
+pub struct IndexSegmentBuilder<'a> {
+    dataset: &'a Dataset,
+    staging_index_uuid: String,
+    partial_indices: Vec<IndexMetadata>,
+    target_segment_bytes: Option<u64>,
+}
+
+impl<'a> IndexSegmentBuilder<'a> {
+    pub(crate) fn new(dataset: &'a Dataset, staging_index_uuid: String) -> Self {
+        Self {
+            dataset,
+            staging_index_uuid,
+            partial_indices: Vec::new(),
+            target_segment_bytes: None,
+        }
+    }
+
+    /// Provide the partial index metadata returned by `execute_uncommitted()`
+    /// for this staging root.
+    pub fn with_partial_indices(mut self, partial_indices: Vec<IndexMetadata>) -> Self {
+        self.partial_indices = partial_indices;
+        self
+    }
+
+    /// Set the target size, in bytes, for merged built segments.
+    ///
+    /// When set, shard outputs will be grouped into larger built segments up to
+    /// approximately this size. When unset, each shard output becomes one built
+    /// segment.
+    pub fn with_target_segment_bytes(mut self, bytes: u64) -> Self {
+        self.target_segment_bytes = Some(bytes);
+        self
+    }
+
+    /// Plan how partial indices should be grouped into built segments.
+    pub async fn plan(&self) -> Result<Vec<IndexSegmentPlan>> {
+        if self.partial_indices.is_empty() {
+            return Err(Error::invalid_input(
+                "IndexSegmentBuilder requires at least one partial index; \
+                 call with_partial_indices(...) with execute_uncommitted() outputs"
+                    .to_string(),
+            ));
+        }
+
+        crate::index::vector::ivf::plan_staging_segments(
+            &self
+                .dataset
+                .indices_dir()
+                .child(self.staging_index_uuid.as_str()),
+            &self.partial_indices,
+            None,
+            self.target_segment_bytes,
+        )
+        .await
+    }
+
+    /// Build one segment from a previously-generated plan.
+    pub async fn build(&self, plan: &IndexSegmentPlan) -> Result<IndexSegment> {
+        crate::index::vector::ivf::build_staging_segment(
+            self.dataset.object_store(),
+            &self.dataset.indices_dir(),
+            plan,
+        )
+        .await
+    }
+
+    /// Plan and build all segments from this staging root.
+    pub async fn build_all(&self) -> Result<Vec<IndexSegment>> {
+        let plans = self.plan().await?;
+        try_join_all(plans.iter().map(|plan| self.build(plan))).await
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
