feat: add Python vector segment builder API

Author: Xuanwo

python/python/lance/__init__.py

@@ -29,6 +29,7 @@
 from .lance import (
     DatasetBasePath,
     FFILanceTableProvider,
+    IndexSegmentBuilder,
     ScanStatistics,
     bytes_read_counter,
     iops_counter,
@@ -64,6 +65,7 @@
     "FragmentMetadata",
     "Index",
     "IndexFile",
+    "IndexSegmentBuilder",
     "LanceDataset",
     "LanceFragment",
     "LanceOperation",

python/python/lance/dataset.py

@@ -47,7 +47,7 @@
 from .dependencies import numpy as np
 from .dependencies import pandas as pd
 from .fragment import DataFile, FragmentMetadata, LanceFragment
-from .indices import IndexConfig, SupportedDistributedIndices
+from .indices import IndexConfig, IndexSegment, SupportedDistributedIndices
 from .lance import (
     CleanupStats,
     Compaction,
@@ -3391,6 +3391,26 @@ def merge_index_metadata(
         self._ds.merge_index_metadata(index_uuid, t, batch_readhead)
         return None
 
+    def create_index_segment_builder(self, staging_index_uuid: str):
+        """
+        Create a builder for turning partial index outputs into committed segments.
+
+        The caller should pass the shared index UUID used during
+        :meth:`create_index` with ``fragment_ids=...`` and ``index_uuid=...``.
+        Then provide the returned partial index metadata through
+        :meth:`IndexSegmentBuilder.with_partial_indices`.
+        """
+        return self._ds.create_index_segment_builder(staging_index_uuid)
+
+    def commit_existing_index_segments(
+        self, index_name: str, column: str, segments: List[IndexSegment]
+    ) -> LanceDataset:
+        """
+        Commit built index segments as one logical index.
+        """
+        self._ds.commit_existing_index_segments(index_name, column, segments)
+        return self
+
     def session(self) -> Session:
         """
         Return the dataset session, which holds the dataset's state.

python/python/lance/indices/__init__.py

@@ -5,9 +5,18 @@
 
 from lance.indices.builder import IndexConfig, IndicesBuilder
 from lance.indices.ivf import IvfModel
+from lance.lance.indices import IndexSegment, IndexSegmentPlan
 from lance.indices.pq import PqModel
 
-__all__ = ["IndicesBuilder", "IndexConfig", "PqModel", "IvfModel", "IndexFileVersion"]
+__all__ = [
+    "IndicesBuilder",
+    "IndexConfig",
+    "PqModel",
+    "IvfModel",
+    "IndexFileVersion",
+    "IndexSegment",
+    "IndexSegmentPlan",
+]
 
 
 class IndexFileVersion(str, Enum):

python/python/lance/lance/__init__.pyi

@@ -61,6 +61,8 @@ from .fragment import (
     RowIdMeta as RowIdMeta,
 )
 from .indices import IndexDescription as IndexDescription
+from .indices import IndexSegment as IndexSegment
+from .indices import IndexSegmentPlan as IndexSegmentPlan
 from .lance import PySearchFilter
 from .optimize import (
     Compaction as Compaction,
@@ -185,6 +187,15 @@ class LanceColumnStatistics:
 class _Session:
     def size_bytes(self) -> int: ...
 
+class IndexSegmentBuilder:
+    @property
+    def staging_index_uuid(self) -> str: ...
+    def with_partial_indices(self, partial_indices: List[Index]) -> Self: ...
+    def with_target_segment_bytes(self, bytes: int) -> Self: ...
+    def plan(self) -> List[IndexSegmentPlan]: ...
+    def build(self, plan: IndexSegmentPlan) -> IndexSegment: ...
+    def build_all(self) -> List[IndexSegment]: ...
+
 class LanceBlobFile:
     def close(self): ...
     def is_closed(self) -> bool: ...
@@ -360,6 +371,12 @@ class _Dataset:
     def merge_index_metadata(
         self, index_uuid: str, index_type: str, batch_readhead: Optional[int] = None
     ): ...
+    def create_index_segment_builder(
+        self, staging_index_uuid: str
+    ) -> IndexSegmentBuilder: ...
+    def commit_existing_index_segments(
+        self, index_name: str, column: str, segments: List[IndexSegment]
+    ) -> None: ...
     def count_fragments(self) -> int: ...
     def num_small_files(self, max_rows_per_group: int) -> int: ...
     def get_fragments(self) -> List[_Fragment]: ...

python/python/lance/lance/indices/__init__.pyi

@@ -16,11 +16,27 @@ from datetime import datetime
 from typing import Optional
 
 import pyarrow as pa
+from ...dataset import Index
 
 class IndexConfig:
     index_type: str
     config: str
 
+class IndexSegment:
+    uuid: str
+    fragment_ids: set[int]
+    index_version: int
+
+    def __repr__(self) -> str: ...
+
+class IndexSegmentPlan:
+    staging_index_uuid: str
+    segment: IndexSegment
+    partial_indices: list[Index]
+    estimated_bytes: int
+
+    def __repr__(self) -> str: ...
+
 def train_ivf_model(
     dataset,
     column: str,

python/src/dataset.rs

@@ -87,7 +87,7 @@ use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
 use crate::error::PythonErrorExt;
 use crate::file::object_store_from_uri_or_path;
 use crate::fragment::FileFragment;
-use crate::indices::{PyIndexConfig, PyIndexDescription};
+use crate::indices::{PyIndexConfig, PyIndexDescription, PyIndexSegment, PyIndexSegmentPlan};
 use crate::namespace::extract_namespace_arc;
 use crate::rt;
 use crate::scanner::ScanStatistics;
@@ -323,6 +323,88 @@ impl MergeInsertBuilder {
     }
 }
 
+#[pyclass(name = "IndexSegmentBuilder", module = "lance", subclass)]
+#[derive(Clone)]
+pub struct PyIndexSegmentBuilder {
+    dataset: Arc<LanceDataset>,
+    staging_index_uuid: String,
+    partial_indices: Vec<IndexMetadata>,
+    target_segment_bytes: Option<u64>,
+}
+
+#[pymethods]
+impl PyIndexSegmentBuilder {
+    #[getter]
+    fn staging_index_uuid(&self) -> String {
+        self.staging_index_uuid.clone()
+    }
+
+    fn with_partial_indices<'a>(
+        mut slf: PyRefMut<'a, Self>,
+        partial_indices: &Bound<'_, PyAny>,
+    ) -> PyResult<PyRefMut<'a, Self>> {
+        let mut indices = Vec::new();
+        for item in partial_indices.try_iter()? {
+            indices.push(item?.extract::<PyLance<IndexMetadata>>()?.0);
+        }
+        slf.partial_indices = indices;
+        Ok(slf)
+    }
+
+    fn with_target_segment_bytes<'a>(
+        mut slf: PyRefMut<'a, Self>,
+        bytes: u64,
+    ) -> PyResult<PyRefMut<'a, Self>> {
+        slf.target_segment_bytes = Some(bytes);
+        Ok(slf)
+    }
+
+    fn plan(&self, py: Python<'_>) -> PyResult<Vec<Py<PyIndexSegmentPlan>>> {
+        let mut builder = self
+            .dataset
+            .create_index_segment_builder(self.staging_index_uuid.clone())
+            .with_partial_indices(self.partial_indices.clone());
+        if let Some(target_segment_bytes) = self.target_segment_bytes {
+            builder = builder.with_target_segment_bytes(target_segment_bytes);
+        }
+        let plans = rt().block_on(Some(py), builder.plan())?.infer_error()?;
+        plans.into_iter()
+            .map(|plan| Py::new(py, PyIndexSegmentPlan::from_inner(plan)))
+            .collect()
+    }
+
+    fn build(
+        &self,
+        py: Python<'_>,
+        plan: &Bound<'_, PyAny>,
+    ) -> PyResult<Py<PyIndexSegment>> {
+        let plan = plan.extract::<PyRef<'_, PyIndexSegmentPlan>>()?;
+        let builder = self
+            .dataset
+            .create_index_segment_builder(self.staging_index_uuid.clone())
+            .with_partial_indices(self.partial_indices.clone());
+        let segment = rt()
+            .block_on(Some(py), builder.build(&plan.inner))?
+            .infer_error()?;
+        Py::new(py, PyIndexSegment::from_inner(segment))
+    }
+
+    fn build_all(&self, py: Python<'_>) -> PyResult<Vec<Py<PyIndexSegment>>> {
+        let mut builder = self
+            .dataset
+            .create_index_segment_builder(self.staging_index_uuid.clone())
+            .with_partial_indices(self.partial_indices.clone());
+        if let Some(target_segment_bytes) = self.target_segment_bytes {
+            builder = builder.with_target_segment_bytes(target_segment_bytes);
+        }
+        let segments = rt().block_on(Some(py), builder.build_all())?.infer_error()?;
+        segments
+            .into_iter()
+            .map(|segment| Py::new(py, PyIndexSegment::from_inner(segment)))
+            .collect()
+    }
+}
+
 impl MergeInsertBuilder {
     fn build_stats<'a>(stats: &MergeStats, py: Python<'a>) -> PyResult<Bound<'a, PyDict>> {
         let dict = PyDict::new(py);
@@ -2019,6 +2101,35 @@ impl Dataset {
         Ok(PyLance(index_metadata))
     }
 
+    fn create_index_segment_builder(&self, staging_index_uuid: String) -> PyResult<PyIndexSegmentBuilder> {
+        Ok(PyIndexSegmentBuilder {
+            dataset: self.ds.clone(),
+            staging_index_uuid,
+            partial_indices: Vec::new(),
+            target_segment_bytes: None,
+        })
+    }
+
+    fn commit_existing_index_segments(
+        &mut self,
+        index_name: &str,
+        column: &str,
+        segments: Vec<PyRef<'_, PyIndexSegment>>,
+    ) -> PyResult<()> {
+        let mut new_self = self.ds.as_ref().clone();
+        let segments = segments
+            .into_iter()
+            .map(|segment| segment.inner.clone())
+            .collect();
+        rt().block_on(
+            None,
+            new_self.commit_existing_index_segments(index_name, column, segments),
+        )?
+        .infer_error()?;
+        self.ds = Arc::new(new_self);
+        Ok(())
+    }
+
     fn drop_index(&mut self, name: &str) -> PyResult<()> {
         let mut new_self = self.ds.as_ref().clone();
         rt().block_on(None, new_self.drop_index(name))?

python/src/indices.rs

@@ -32,12 +32,14 @@ use pyo3::{
 use lance::index::DatasetIndexInternalExt;
 
 use crate::fragment::FileFragment;
-use crate::utils::PyJson;
+use crate::utils::{PyJson, PyLance};
 use crate::{
     dataset::Dataset, error::PythonErrorExt, file::object_store_from_uri_or_path_no_options, rt,
 };
 use lance::index::vector::ivf::write_ivf_pq_file_from_existing_index;
-use lance_index::{DatasetIndexExt, IndexDescription, IndexSegment, IndexType};
+use lance_index::{
+    DatasetIndexExt, IndexDescription, IndexSegment, IndexSegmentPlan, IndexType,
+};
 use uuid::Uuid;
 
 #[pyclass(name = "IndexConfig", module = "lance.indices", get_all)]
@@ -58,6 +60,93 @@ impl PyIndexConfig {
     }
 }
 
+#[pyclass(name = "IndexSegment", module = "lance.indices")]
+#[derive(Debug, Clone)]
+pub struct PyIndexSegment {
+    pub(crate) inner: IndexSegment,
+}
+
+impl PyIndexSegment {
+    pub(crate) fn from_inner(inner: IndexSegment) -> Self {
+        Self { inner }
+    }
+}
+
+#[pymethods]
+impl PyIndexSegment {
+    #[getter]
+    fn uuid(&self) -> String {
+        self.inner.uuid().to_string()
+    }
+
+    #[getter]
+    fn fragment_ids(&self) -> HashSet<u32> {
+        self.inner.fragment_bitmap().iter().collect()
+    }
+
+    #[getter]
+    fn index_version(&self) -> i32 {
+        self.inner.index_version()
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "IndexSegment(uuid={}, fragment_ids={:?}, index_version={})",
+            self.uuid(),
+            self.fragment_ids(),
+            self.index_version()
+        )
+    }
+}
+
+#[pyclass(name = "IndexSegmentPlan", module = "lance.indices")]
+#[derive(Debug, Clone)]
+pub struct PyIndexSegmentPlan {
+    pub(crate) inner: IndexSegmentPlan,
+}
+
+impl PyIndexSegmentPlan {
+    pub(crate) fn from_inner(inner: IndexSegmentPlan) -> Self {
+        Self { inner }
+    }
+}
+
+#[pymethods]
+impl PyIndexSegmentPlan {
+    #[getter]
+    fn staging_index_uuid(&self) -> String {
+        self.inner.staging_index_uuid().to_string()
+    }
+
+    #[getter]
+    fn segment(&self) -> PyIndexSegment {
+        PyIndexSegment::from_inner(self.inner.segment().clone())
+    }
+
+    #[getter]
+    fn partial_indices(&self) -> Vec<PyLance<lance_table::format::IndexMetadata>> {
+        self.inner
+            .partial_indices()
+            .iter()
+            .cloned()
+            .map(PyLance)
+            .collect()
+    }
+
+    #[getter]
+    fn estimated_bytes(&self) -> u64 {
+        self.inner.estimated_bytes()
+    }
+    fn __repr__(&self) -> String {
+        format!(
+            "IndexSegmentPlan(staging_index_uuid={}, partial_indices={}, estimated_bytes={})",
+            self.staging_index_uuid(),
+            self.inner.partial_indices().len(),
+            self.estimated_bytes()
+        )
+    }
+}
+
 #[pyclass(name = "IvfModel", module = "lance.indices")]
 #[derive(Debug, Clone)]
 pub struct PyIvfModel {
@@ -619,6 +708,8 @@ pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     indices.add_wrapped(wrap_pyfunction!(load_shuffled_vectors))?;
     indices.add_class::<PyIvfModel>()?;
     indices.add_class::<PyIndexConfig>()?;
+    indices.add_class::<PyIndexSegment>()?;
+    indices.add_class::<PyIndexSegmentPlan>()?;
     indices.add_class::<PyIndexDescription>()?;
     indices.add_class::<PyIndexSegmentDescription>()?;
     indices.add_wrapped(wrap_pyfunction!(get_ivf_model))?;