lance-format / lance-format/lance

Refactor IVF index build/optimize into a unified pipeline

Open
#6,686 1 comment 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

A-index chore
Dominant language
Rust
Stars
7.1k
Forks
852
Avg merge
3d 18h
Merged PRs (30d)
272

Description

Motivation

The IVF index builder (IvfIndexBuilder in lance/src/index/vector/builder.rs) handles
both fresh index creation and incremental optimization, but the two code paths are tangled
together inside build_partitions. The split, join, and no-adjustment cases each have
distinct routing logic for how data flows through the pipeline, making it hard to follow
and extend.

The core operations are the same in every case:

  1. Decide on centroids (train from scratch or update existing)
  2. Decide on a quantizer (train from scratch or reuse existing)
  3. Identify which vectors need (re-)quantization and shuffle them
  4. For each partition, combine existing quantized data with newly shuffled data and build the sub-index
  5. Write the final index files

But currently these steps are interleaved with builder state management, and the
split/join logic is embedded deep inside build_partitions rather than being expressed
as a centroid update that feeds into the standard pipeline.

Proposed design

A top-level function with clear stages:

async fn build_ivf_index(
    dataset: &Dataset,
    column: &str,
    existing_index: Option<&dyn VectorIndex>,
    unindexed: Option<impl RecordBatchStream>,
    options: &BuildOptions,
) -> Result<IndexOutput> {
    // 1. Centroids
    let (ivf_model, reassign_ids) = if let Some(existing) = existing_index {
        update_ivf(existing.ivf_model(), dataset, &options).await?
    } else {
        (train_ivf(dataset, &options).await?, None)
    };

    // 2. Quantizer
    let quantizer = if let Some(existing) = existing_index {
        existing.quantizer()
    } else {
        train_quantizer(dataset, &ivf_model, &options).await?
    };

    // 3. Shuffle vectors that need (re-)quantization
    let vectors_to_process = new_row_ids(dataset, existing_index)
        .union(reassign_ids.unwrap_or_default());
    let shuffle_reader = shuffle_and_quantize(
        dataset, vectors_to_process, &ivf_model, &quantizer,
    ).await?;

    // 4-5. Build per-partition sub-indices and write final files
    write_partitions(
        existing_index,    // source of unchanged PQ codes
        &shuffle_reader,   // source of newly quantized data
        &reassign_ids,     // which partitions had data removed
        &ivf_model,
        &quantizer,
        &options,
    ).await
}
Key design points

update_ivf returns reassign_ids: This is the unifying abstraction. Split, join,
and no-op all produce a potentially updated IVF model and a set of row IDs whose partition
assignments changed. For splits, reassign_ids contains all row IDs from affected
partitions (split targets + neighbors). For joins, it contains the deleted partition's
row IDs. For no-op, it's empty.

shuffle_and_quantize takes a row ID set: It doesn't know why vectors need
processing, just which ones. It streams raw vectors from the dataset, assigns each to
its nearest centroid, computes PQ residuals, encodes to PQ codes, and writes to temp
files via the existing shuffler infrastructure. This is the same pipeline used today
for both fresh builds and the streaming split we just added.

write_partitions takes two data sources: For each partition, it reads:

  • Unchanged PQ codes from the existing index (skipping partitions that were
    split/joined, since those vectors were re-shuffled)
  • Newly quantized data from the shuffle reader

It combines them, builds the sub-index (FlatIndex or HNSW), and writes to the final
index files. The sub-index type is generic and handled entirely within this function.

write_partitions needs to know which partitions had data removed: When a partition
is split or joined, its existing PQ codes in the old index are stale. reassign_ids
(or a derived set of affected partition IDs) tells write_partitions to skip reading
from the existing index for those partitions. Their data comes entirely from the shuffle
reader.

How each case maps to this model
Case update_ivf reassign_ids shuffle_and_quantize write_partitions
Fresh build train_ivf None All vectors Shuffle reader only
Optimize (no split/join) Pass through existing Empty New unindexed only Existing + shuffle
Optimize with split Train k=2 per oversized partition, update centroids Affected partition row IDs New + reassigned Existing (unaffected) + shuffle
Optimize with join Remove centroid, update model Deleted partition row IDs New + reassigned Existing (unaffected) + shuffle
Delta merge (no rebalance) Pass through existing Empty New unindexed only Multiple existing + shuffle
HNSW

HNSW is a per-partition sub-index and fits entirely within write_partitions. After
combining existing and new PQ codes for a partition, the sub-index is built via
S::index_vectors(&storage, params) — this is O(1) for FlatIndex and O(n log n) for
HNSW. The upstream pipeline (centroid training, shuffling, quantization) is identical
regardless of sub-index type.

IVF_HNSW variants have much larger target partition sizes (1M vs 8K for IVF_PQ), so splits
are rarer and per-partition data is larger. This doesn't affect the pipeline structure, but
write_partitions holds more data per partition in memory for these variants.

What this replaces

  • IvfIndexBuilder struct and its build() method
  • build_partitions() with its inline split/join/merge routing
  • split_partitions_streaming() (becomes part of update_ivf)
  • join_partition() / join_partition_impl() (becomes part of update_ivf)
  • take_partition_batches() (becomes internal to write_partitions)
  • merge_partitions() (becomes internal to write_partitions)
  • The PartitionAdjustment enum and AssignResult / SplitResult structs
  • The dual code paths in ivf.rs optimize_vector_indices that construct different
    IvfIndexBuilder variants

Non-goals

  • Changing the on-disk format or quantizer types
  • Changing how HNSW graphs are built
  • Retraining quantizers during optimize (currently never done; users create a new index)
  • Changing the shuffle file format or shuffler implementations

Related work

  • Streaming partition splits (already landed) — introduced reshuffle_partitions and
    compute_split_centroids which are precursors to update_ivf and
    shuffle_and_quantize

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with IvfIndexBuilder and build_partitions in lance/src/index/vector/builder.rs, then trace the dual optimize_vector_indices paths in ivf.rs and the existing split/join helpers named in the issue. Compare fresh builds, optimization, splits, joins, and delta merges against the proposed staged pipeline. Done means these cases share the unified flow without changing the on-disk format, quantizer types, HNSW construction, or shuffle format.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
databases, machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.