Unified API or Separate APIs for Distributed Index Building？

[wayneli-vt]

This is a follow-up to issue #5164.
As the scale of unstructured data continues to grow, the need for a robust distributed index-building capability is becoming increasingly critical. Because different index types, such as the Range-based BTree Index, have different distributed build strategies, the current build implementation is no longer sufficient for all index types.

Current Implementation

Current implementation for distributed indexing is an extension of the single-node logic. It adapts the CreateIndexBuilder by adding a fragments parameter. This approach allowed us to reuse the existing single-node build process with minimal changes, achieving distributed capabilities quickly.
The build process is divided into the following three stages:

1. build_distributed_index: An instance of CreateIndexBuilder is created, and the presence of fragments determines if it's a distributed build.
2. merge_metadata: After the distributed build completes, the Dataset.merge_index_metadata() method is called to consolidate the metadata. (This is invoked at the Python/Java layer by directly calling into Index, and there isn't a unified entry point for this at the Rust level).
3. commit_index: The final index information is committed to the dataset's manifest.

Limitations

This design is built on two key assumptions:

1. All computation for index creation occurs within Lance itself.
2. The build process can be partitioned at the fragment level.

As we expand support for more advanced indexing strategies, we're encountering new requirements that go beyond this initial scope. A prime example is the Range-Based Distributed B-Tree Index, where

1. The expensive computation (global sorting) is performed by external big data engines (like Spark, Flink, etc.).
2. The data is partitioned by value range, not by fragment.

The existing CreateIndexBuilder was not originally designed to handle these scenarios, suggesting a need for a more adaptable approach.

Proposal

As mentioned in (#5202 (comment)), similar to compaction API, I'd like to propose creating a dedicated, unified API for distributed index building at the Rust level, perhaps a DistributedIndexBuilder, which would encapsulate the entire lifecycle:

1. build_distributed_index: Generalize beyond fragment-level splitting to support arbitrary data partitions and accept pre-processed data from external compute engines.
2. merge_metadata: Move this logic from a standalone Dataset interface into the DistributedIndexBuilder, abstracting the Index details away from the Python/Java bindings.
3. commit_index: Encapsulate the final commit step within the builder as well.

Would this unified builder be flexible enough to accommodate the distributed build requirements for various index types, including vector indexes?
On a related note, I've noticed that vector indexes currently use IvfIndexBuilder internally. Could this new, unified API also integrate the logic from IvfIndexBuilder, creating a consistent interface for both single-node and distributed builds across all index families?
Here is a rough sketch of the proposed structure:

// The common configuration base for all index builders.
pub struct IndexBuilderBase<'a> {
    dataset: &'a mut Dataset,
    columns: Vec<String>,
    index_type: IndexType,
    params: &'a dyn IndexParams,
    name: Option<String>,
    replace: bool,
    train: bool,
    index_uuid: Option<String>,
}
// Basic build impls (e.g., .name(), .replace()) would be here.

// A builder for standard, monolithic index creation.
pub struct CreateIndexBuilder<'a> {
    base: IndexBuilderBase<'a>,
}
// Impl for `execute (build_index + commit_index)` would be here.
...

impl<'a> Deref for CreateIndexBuilder<'a> {
    type Target = IndexBuilderBase<'a>;
    fn deref(&self) -> &Self::Target { &self.base }
}
impl<'a> DerefMut for CreateIndexBuilder<'a> {
    fn deref_mut(&mut self) -> &mut Self::Target { &mut self.base }
}

// A builder for distributed, multi-stage index creation.
// Or a more specific DistributedBTreeIndexBuilder
pub struct DistributedIndexBuilder<'a> {
    base: IndexBuilderBase<'a>,
    // For fragment-based builds
    fragments: Option<Vec<u32>>,
    // For different-kind-based builds such as range-based
    // (partition_id, preprocessed_data)
    external_data: Option<(u32, Box<dyn RecordBatchReader + Send + 'static>)>,
}
// Impls for `build_distributed_index`, `merge_metadata`, and `commit_index` would be here.
...

I'm still debating whether a generic DistributedIndexBuilder is the right path, or if a more specific extension, like creating a new BTreeIndexBuilder, would be a better initial step.
@westonpace @wjones127 @BubbleCal @chenghao-guo I would greatly appreciate your thoughts and any alternative suggestions on this!

[steFaiz]

Maybe we should start a discussion about this? I think designing the API for a distributed index is quite complex—for example, distributed construction, updates, and remapping. We need to abstract out common traits and allow different indexes to have their own implementations.

[chenghao-guo]

I think the proposal to generalize beyond fragment-level splitting is crucial for long-term viability. As you've noted with the B-Tree example, different index types have fundamentally different distributed build strategies. I am not sure if build_distributed_index can cover all index types, B-tree and FTS shall be fine per my understanding and can benefit from this.

While for IVF index, I am trying to draft a PR to split the process and make it possible for ray to call that, like Training centroids globally -> Assignment and Per-Worker Writing -> Merging the partial index. If I need the merging the partial index to be distributed by ray/spark, could this situation benefit from your design?

[wayneli-vt]

I think the proposal to generalize beyond fragment-level splitting is crucial for long-term viability. As you've noted with the B-Tree example, different index types have fundamentally different distributed build strategies. I am not sure if build_distributed_index can cover all index types, B-tree and FTS shall be fine per my understanding and can benefit from this.

While for IVF index, I am trying to draft a PR to split the process and make it possible for ray to call that, like Training centroids globally -> Assignment and Per-Worker Writing -> Merging the partial index. If I need the merging the partial index to be distributed by ray/spark, could this situation benefit from your design?

Thanks for the great work on the draft PR! It really helps to make the distributed index build concrete.

To make sure I'm following your implementation correctly, let me summarize the workflow as I see:

1. Global Training: A single-node process samples the entire dataset to compute the initial centroids.
2. Distributed Build: The process is distributed across workers, where each worker operates on a set of fragments.
3. Metadata Merge: A final, single-node process merges the metadata from all the sub-indices into the final index file.

If this understanding is correct, we could add a preparation step at first before build_distributed_index. The centroids generated in Global Training phase could be passed in as preprocessed_data to the main build process.

Alternatively, we could create a dedicated IVFDistributedIndexBuilder for it. This would grant us the maximum flexibility to accommodate the different build processes required by various distributed indexes.

Both looks good to me! And thanks again for driving this forward!

[westonpace]

Yes, this API could be extended to IVF/PQ building. It would be possible to make an even more complex IVF build API if one wanted but this might be a good step.

Adapting the Proposed API

Build / Global Training

In the IVF / PQ build this is where we would calculate the centroids (this is done with a small sample of the dataset) and for each centroid we would calculate a set of PQ centroids.

Distributed Build

In the IVF / PQ build this is where we would apply the transformations to the vector column for every row (find the nearest centroid, calculate the residual, apply PQ compression, ...)

Merge

In the IVF / PQ build this is where we would sort the transformed vectors by their assigned partition (this is a fast sort since we know the range of partitions and can easily do a radix sort).

A More Complicated API

It is possible to make the API for building IVF / PQ even more complicated if one wanted to 😆. The initial calculation of the centroids uses the k-means algorithm to calculate the centroid values. This operation itself is often parallelized in many different iterations. When creating a very large index (e.g. 20K partitions) with the default sample rate (256) we have ~5M vectors in our initial sampling. We can parallelize across this data when calculating distances from the centroids for each sample.

So you could conceivably have a very complicated parallel build API that allowed for this kmeans step to be parallelized in addition to the distributed build step.

flowchart LR
    A[Calc Initial Centroids] -->B[Calc Distance]
    A --> C[Calc Distance]
    A --> D[Calc Distance]
    B --> E[Update Centroids]
    C --> E
    D --> E
    E --> F[Calc Distance]
    E --> G[Calc Distance]
    E --> H[Calc Distance]
    F --> I[...]
    G --> I
    H --> I
    I --> J[Transform Vector]
    I --> K[Transform Vector]
    I --> L[Transform Vector]
    J --> M[Shuffle / Merge]
    K --> M
    L --> M

Loading

Should someone do this? Maybe not 😅.

I guess my point is that someone will always be able to come up with a more complex API to optimize a specific index build process and that is ok. However, the Build / Distribute / Merge API is simple, applies to many indexes, and easy to use.

Follow-up Question

How would the build / distribute / merge API apply to the range-based btree? Would the build step be calculating the bounds of the partition ranges and the distribute step be the shuffling of the data?

[wayneli-vt]

Thank you so much for your thoughtful and detailed reply!

Range-based BTree Index

How would the build / distribute / merge API apply to the range-based btree? Would the build step be calculating the bounds of the partition ranges and the distribute step be the shuffling of the data?

Yes, your understanding is right. The process for the Range-based BTree Index maps perfectly to this API:

1. Global Training: Calculate the range boundaries for partitioning.
2. Distributed Build: Build a sub-index within each partition.
3. Merge: Merge the metadata from all sub-indexes.

Currently, the Global Training step is implemented in Spark, which involves a full table scan to sample the data and then calculates the ranges. I agree that we can optimize this by moving sampling into Lance and using the take operation to precisely fetch only the sampled data.

Vector Indexes (IVF/PQ)

I completely agree that this API is flexible enough to be extended to IVF/PQ index building.
Regarding my earlier point about using external_data for centroids, @chenghao-guo , we might have a more elegant solution. It makes more sense for the centroids to be treated as part of the index metadata. Specifically, we could add centroids as a property to the IndexBuilder, which could then be broadcast to each distributed task.

The external_data parameter could be reserved for ingesting data that has been pre-processed (e.g., shuffled) by an external compute engine. This would keep the API's intent clearer.

Extending the API for Other Operations

I'm excited that this Global Training / Distribute Build / Merge pattern can be extended to other distributed index operations like update and remap.
To use the Range-based BTree Index as an example, a distributed update could work as follows:

1. Global Training: Retrieve the range boundaries from the existing index.
2. Distribute Build: The data from UnIndexedFragments can be shuffled and sorted based on those ranges, and then merge them with existing ordered sub-indexes.
3. Merge: Update new index metadata.

This would pave the way for a fully distributed index management lifecycle. (Considering the cost of a full rebuild, perhaps this is a lower-priority feature to consider for the future).

Next Steps & A Quick Question

Thank you once again, @westonpace and @chenghao-guo , for your invaluable guidance!
As a first step, I'll implement the new DistributedBTreeIndexBuilder, and this pattern can extend to other index types later on. To make sure I'm on the right track, should this new distributed API be integrated into existing #5202 , or would a new, separate PR after #5202 is merged be better?
Looking forward to your thoughts!

[chenghao-guo]

Hi @wayneli-vt ,Thank you very much for your valuable suggestions.
Regarding your first point, I agree that it would be best to propose a new, separate Pull Request after #5202 has been merged. This approach will help ensure that we avoid introducing any potential breaking changes, allowing for its smooth integration.
Concerning the vector index, I have also encountered some other recall related challenges while implementing #5117. However, I am actively working to adhere to your proposed design. Once these tasks are completed, you will be well-positioned to refactor all interfaces into a unified API, which I believe will result in a much more elegant and less conflicting design.